jellyfin-ha-src/tests/Jellyfin.Server.Tests/Migrations/PostgreSqlTestServer.cs

using System;
using System.Threading;
using System.Threading.Tasks;
using DotNet.Testcontainers.Builders;
using Npgsql;
using Testcontainers.PostgreSql;

namespace Jellyfin.Server.Tests.Migrations;

/// <summary>
/// Hands out a PostgreSQL server for the tests that need one. A server named by
/// <c>JELLYFIN_TEST_POSTGRES</c> is used as is, so CI can attach a service container instead of running
/// a docker daemon of its own; without it a container is started through testcontainers.
/// </summary>
public sealed class PostgreSqlTestServer : IAsyncDisposable
{
    /// <summary>
    /// The connection string of an already running server. Must be able to create databases.
    /// </summary>
    public const string ConnectionStringVariable = "JELLYFIN_TEST_POSTGRES";

    private readonly PostgreSqlContainer? _container;

    private PostgreSqlTestServer(PostgreSqlContainer? container, string connectionString)
    {
        _container = container;
        ConnectionString = connectionString;
    }

    /// <summary>
    /// Gets the connection string of the server holding the test databases.
    /// </summary>
    public string ConnectionString { get; }

    /// <summary>
    /// Starts or attaches to a PostgreSQL server and waits until it accepts connections.
    /// </summary>
    /// <returns>The running server.</returns>
    public static async Task<PostgreSqlTestServer> StartAsync()
    {
        var provided = Environment.GetEnvironmentVariable(ConnectionStringVariable);
        if (!string.IsNullOrWhiteSpace(provided))
        {
            var attached = new PostgreSqlTestServer(null, provided);
            await attached.WaitUntilReadyAsync().ConfigureAwait(false);
            return attached;
        }

        var container = new PostgreSqlBuilder("postgres:16-alpine")
            .WithWaitStrategy(Wait.ForUnixContainer().UntilCommandIsCompleted("pg_isready"))
            .Build();
        await container.StartAsync().ConfigureAwait(false);

        var started = new PostgreSqlTestServer(container, container.GetConnectionString());
        await started.WaitUntilReadyAsync().ConfigureAwait(false);
        return started;
    }

    /// <summary>
    /// Creates an empty database and returns a connection string pointing at it.
    /// </summary>
    /// <param name="name">The database name.</param>
    /// <param name="cancellationToken">The cancellation token.</param>
    /// <returns>The connection string of the new database.</returns>
    public async Task<string> CreateDatabaseAsync(string name, CancellationToken cancellationToken)
    {
        await using var adminDataSource = new NpgsqlDataSourceBuilder(ConnectionString).Build();

        // A server handed in through the environment outlives the run, so a second run finds the
        // databases the first one left behind.
        await using var drop = adminDataSource.CreateCommand($"DROP DATABASE IF EXISTS {name}");
        await drop.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);

        await using var command = adminDataSource.CreateCommand($"CREATE DATABASE {name}");
        await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);

        return new NpgsqlConnectionStringBuilder(ConnectionString) { Database = name }.ConnectionString;
    }

    /// <inheritdoc/>
    public async ValueTask DisposeAsync()
    {
        if (_container is not null)
        {
            await _container.DisposeAsync().ConfigureAwait(false);
        }
    }

    /// <summary>
    /// Waits until a real connection is accepted. <c>pg_isready</c> also answers for the short-lived server
    /// the entrypoint runs while it initializes the data directory.
    /// </summary>
    private async Task WaitUntilReadyAsync()
    {
        await using var dataSource = new NpgsqlDataSourceBuilder(ConnectionString).Build();
        for (var attempt = 1; ; attempt++)
        {
            try
            {
                await using var command = dataSource.CreateCommand("SELECT 1");
                await command.ExecuteNonQueryAsync(CancellationToken.None).ConfigureAwait(false);
                return;
            }
            catch (NpgsqlException) when (attempt < 60)
            {
                await Task.Delay(TimeSpan.FromSeconds(1)).ConfigureAwait(false);
            }
        }
    }
}