tiku-backend.net/Tiku.IntegrationTests/Infrastructure/PostgresTestDatabase.cs

using Microsoft.EntityFrameworkCore;
using Npgsql;
using Tiku.Infrastructure.Persistence;

namespace Tiku.IntegrationTests.Infrastructure;

internal sealed class PostgresTestDatabase : IDisposable
{
    private const string DatabasePrefix = "tiku_it_";

    private static readonly Lazy<PostgresTestDatabaseTemplate> Template = new(
        PostgresTestDatabaseTemplate.Create,
        LazyThreadSafetyMode.ExecutionAndPublication);

    private readonly string adminConnectionString;
    private bool disposed;

    private PostgresTestDatabase(string adminConnectionString, string databaseName, string connectionString)
    {
        this.adminConnectionString = adminConnectionString;
        DatabaseName = databaseName;
        ConnectionString = connectionString;
    }

    public string DatabaseName { get; }

    public string ConnectionString { get; }

    public void Dispose()
    {
        if (disposed) return;

        disposed = true;
        var builder = new NpgsqlConnectionStringBuilder(adminConnectionString);
        using var adminConnection = new NpgsqlConnection(builder.ConnectionString);
        adminConnection.Open();

        using (var terminateConnections = adminConnection.CreateCommand())
        {
            terminateConnections.CommandText =
                "SELECT pg_terminate_backend(pid) FROM pg_stat_activity WHERE datname = $1 AND pid <> pg_backend_pid()";
            terminateConnections.Parameters.AddWithValue(DatabaseName);
            terminateConnections.ExecuteNonQuery();
        }

        using var dropDatabase = adminConnection.CreateCommand();
        dropDatabase.CommandText = $"DROP DATABASE IF EXISTS {QuoteIdentifier(DatabaseName)}";
        dropDatabase.ExecuteNonQuery();
    }

    public static PostgresTestDatabase Create()
    {
        var template = Template.Value;
        var adminConnectionString = template.AdminConnectionString;
        var databaseName = $"{DatabasePrefix}{Guid.NewGuid():N}";
        var adminBuilder = new NpgsqlConnectionStringBuilder(adminConnectionString);

        using (var adminConnection = new NpgsqlConnection(adminBuilder.ConnectionString))
        {
            adminConnection.Open();
            using var createDatabase = adminConnection.CreateCommand();
            createDatabase.CommandText =
                $"CREATE DATABASE {QuoteIdentifier(databaseName)} TEMPLATE {QuoteIdentifier(template.DatabaseName)}";
            createDatabase.ExecuteNonQuery();
        }

        var databaseBuilder = new NpgsqlConnectionStringBuilder(adminBuilder.ConnectionString)
        {
            Database = databaseName,
            Pooling = true,
            MaxPoolSize = 10
        };
        var database = new PostgresTestDatabase(
            adminBuilder.ConnectionString,
            databaseName,
            databaseBuilder.ConnectionString);

        return database;
    }

    internal static string QuoteIdentifier(string identifier)
    {
        if (!identifier.StartsWith(DatabasePrefix, StringComparison.Ordinal) ||
            identifier.Any(character => !char.IsAsciiLetterOrDigit(character) && character != '_'))
            throw new InvalidOperationException("Refusing to use an unsafe integration-test database name.");

        return $"\"{identifier}\"";
    }
}

internal sealed class PostgresTestDatabaseTemplate : IDisposable
{
    private const string TemplatePrefix = "tiku_it_template_";
    private bool disposed;

    private PostgresTestDatabaseTemplate(string adminConnectionString, string databaseName)
    {
        AdminConnectionString = adminConnectionString;
        DatabaseName = databaseName;
    }

    public string AdminConnectionString { get; }

    public string DatabaseName { get; }

    public void Dispose()
    {
        if (disposed) return;

        disposed = true;
        using var adminConnection = new NpgsqlConnection(AdminConnectionString);
        adminConnection.Open();
        using var dropDatabase = adminConnection.CreateCommand();
        dropDatabase.CommandText =
            $"DROP DATABASE IF EXISTS {PostgresTestDatabase.QuoteIdentifier(DatabaseName)}";
        dropDatabase.ExecuteNonQuery();
    }

    public static PostgresTestDatabaseTemplate Create()
    {
        var adminConnectionString = Environment.GetEnvironmentVariable("TIKU_TEST_POSTGRES_ADMIN") ??
                                    DockerPostgresTestDefaults.AdminConnectionString;
        var adminBuilder = new NpgsqlConnectionStringBuilder(adminConnectionString);
        if (!string.Equals(adminBuilder.Database, "postgres", StringComparison.OrdinalIgnoreCase))
            throw new InvalidOperationException(
                "TIKU_TEST_POSTGRES_ADMIN must target the postgres maintenance database.");

        var databaseName = $"{TemplatePrefix}{Environment.ProcessId}_{Guid.NewGuid():N}";
        var template = new PostgresTestDatabaseTemplate(adminBuilder.ConnectionString, databaseName);

        try
        {
            using (var adminConnection = new NpgsqlConnection(template.AdminConnectionString))
            {
                adminConnection.Open();
                using var createDatabase = adminConnection.CreateCommand();
                createDatabase.CommandText =
                    $"CREATE DATABASE {PostgresTestDatabase.QuoteIdentifier(databaseName)}";
                createDatabase.ExecuteNonQuery();
            }

            var databaseBuilder = new NpgsqlConnectionStringBuilder(template.AdminConnectionString)
            {
                Database = databaseName,
                Pooling = false
            };
            var options = new DbContextOptionsBuilder<TikuDbContext>()
                .UseNpgsql(databaseBuilder.ConnectionString, npgsql =>
                    npgsql.MigrationsAssembly(typeof(TikuDbContext).Assembly.FullName))
                .Options;
            using (var dbContext = new TikuDbContext(options))
            {
                dbContext.Database.Migrate();
            }

            using (var adminConnection = new NpgsqlConnection(template.AdminConnectionString))
            {
                adminConnection.Open();
                using var preventConnections = adminConnection.CreateCommand();
                preventConnections.CommandText =
                    $"ALTER DATABASE {PostgresTestDatabase.QuoteIdentifier(databaseName)} ALLOW_CONNECTIONS false";
                preventConnections.ExecuteNonQuery();
            }

            AppDomain.CurrentDomain.ProcessExit += (_, _) => template.Dispose();
            return template;
        }
        catch
        {
            template.Dispose();
            throw;
        }
    }
}