git.stella-ops.org/src/Findings/StellaOps.Findings.Ledger/Infrastructure/Postgres/LedgerDataSource.cs

using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using Npgsql;
using StellaOps.Findings.Ledger.Options;

namespace StellaOps.Findings.Ledger.Infrastructure.Postgres;

public sealed class LedgerDataSource : IAsyncDisposable
{
    private readonly NpgsqlDataSource _dataSource;
    private readonly LedgerServiceOptions.DatabaseOptions _options;
    private readonly ILogger<LedgerDataSource> _logger;

    public LedgerDataSource(
        IOptions<LedgerServiceOptions> options,
        ILogger<LedgerDataSource> logger)
    {
        ArgumentNullException.ThrowIfNull(options);
        _options = options.Value.Database;
        _logger = logger ?? throw new ArgumentNullException(nameof(logger));

        var builder = new NpgsqlDataSourceBuilder(_options.ConnectionString);
        _dataSource = builder.Build();
    }

    public int CommandTimeoutSeconds => _options.CommandTimeoutSeconds;

    public async ValueTask DisposeAsync()
    {
        await _dataSource.DisposeAsync().ConfigureAwait(false);
    }

    public Task<NpgsqlConnection> OpenConnectionAsync(string tenantId, CancellationToken cancellationToken)
        => OpenConnectionInternalAsync(tenantId, cancellationToken);

    private async Task<NpgsqlConnection> OpenConnectionInternalAsync(string tenantId, CancellationToken cancellationToken)
    {
        var connection = await _dataSource.OpenConnectionAsync(cancellationToken).ConfigureAwait(false);

        try
        {
            await ConfigureSessionAsync(connection, tenantId, cancellationToken).ConfigureAwait(false);
        }
        catch
        {
            await connection.DisposeAsync().ConfigureAwait(false);
            throw;
        }

        return connection;
    }

    private async Task ConfigureSessionAsync(NpgsqlConnection connection, string tenantId, CancellationToken cancellationToken)
    {
        try
        {
            await using (var command = new NpgsqlCommand("SET TIME ZONE 'UTC';", connection))
            {
                command.CommandTimeout = _options.CommandTimeoutSeconds;
                await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
            }

            if (!string.IsNullOrWhiteSpace(tenantId))
            {
                await using var tenantCommand = new NpgsqlCommand("SELECT set_config('app.current_tenant', @tenant, false);", connection);
                tenantCommand.CommandTimeout = _options.CommandTimeoutSeconds;
                tenantCommand.Parameters.AddWithValue("tenant", tenantId);
                await tenantCommand.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
            }
        }
        catch (Exception ex)
        {
            if (_logger.IsEnabled(LogLevel.Error))
            {
                _logger.LogError(ex, "Failed to configure PostgreSQL session for tenant {TenantId}.", tenantId);
            }

            throw;
        }
    }
}