yuanshen/Yuanshen.Gameserver/Database/PostgresAccountRepository.cs

using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using Npgsql;
using Yuanshen.Gameserver.Options;

namespace Yuanshen.Gameserver.Database;

internal sealed class PostgresAccountRepository : IAccountRepository, IAsyncDisposable
{
    private const string CreateTableSql =
        """
        CREATE TABLE IF NOT EXISTS accounts (
            uid         integer GENERATED ALWAYS AS IDENTITY PRIMARY KEY,
            account_key text UNIQUE NOT NULL,
            created_at  timestamptz NOT NULL DEFAULT now()
        )
        """;

    private const string GetOrCreateSql =
        """
        INSERT INTO accounts (account_key) VALUES (@key)
        ON CONFLICT (account_key) DO UPDATE SET account_key = EXCLUDED.account_key
        RETURNING uid
        """;

    private readonly NpgsqlDataSource _dataSource;
    private readonly ILogger _logger;
    private readonly SemaphoreSlim _schemaLock = new(1, 1);
    private volatile bool _schemaReady;

    public PostgresAccountRepository(IOptions<DatabaseOptions> options, ILogger<PostgresAccountRepository> logger)
    {
        _logger = logger;
        _dataSource = new NpgsqlDataSourceBuilder(options.Value.ConnectionString).Build();
    }

    public async ValueTask<uint> GetOrCreateUidAsync(string accountKey)
    {
        if (string.IsNullOrWhiteSpace(accountKey))
            accountKey = "default";

        try
        {
            await EnsureSchemaAsync();

            await using NpgsqlCommand command = _dataSource.CreateCommand(GetOrCreateSql);
            command.Parameters.AddWithValue("key", accountKey);

            object? result = await command.ExecuteScalarAsync();
            uint uid = Convert.ToUInt32(result);
            _logger.LogDebug("Account '{key}' resolved to uid {uid}", accountKey, uid);
            return uid;
        }
        catch (Exception exception)
        {
            _logger.LogError("Failed to resolve account '{key}': {message}", accountKey, exception.Message);
            return 1;
        }
    }

    private async ValueTask EnsureSchemaAsync()
    {
        if (_schemaReady)
            return;

        await _schemaLock.WaitAsync();
        try
        {
            if (_schemaReady)
                return;

            await using NpgsqlCommand command = _dataSource.CreateCommand(CreateTableSql);
            await command.ExecuteNonQueryAsync();
            _schemaReady = true;
        }
        finally
        {
            _schemaLock.Release();
        }
    }

    public ValueTask DisposeAsync()
    {
        _schemaLock.Dispose();
        return _dataSource.DisposeAsync();
    }
}