scadaproj/ZB.MOM.WW.Auth/src/ZB.MOM.WW.Auth.ApiKeys/Sqlite/SqliteAuthStoreMigrator.cs

using System.Globalization;
using Microsoft.Data.Sqlite;
using ZB.MOM.WW.Auth.Abstractions.ApiKeys;

namespace ZB.MOM.WW.Auth.ApiKeys.Sqlite;

/// <summary>Thrown when the auth store cannot be migrated to the supported schema.</summary>
public sealed class AuthStoreMigrationException(string message) : InvalidOperationException(message);

/// <summary>
/// Creates the API-key store schema and records the applied version. Idempotent: it
/// is safe to run repeatedly. Refuses to run against a database whose on-disk version
/// is newer than this build supports.
/// </summary>
public sealed class SqliteAuthStoreMigrator(AuthSqliteConnectionFactory connectionFactory) : IApiKeyStoreMigrator
{
    /// <summary>Applies the schema migration to the auth store.</summary>
    /// <param name="cancellationToken">Cancellation token.</param>
    /// <exception cref="AuthStoreMigrationException">
    /// The on-disk schema version is newer than <see cref="SqliteAuthSchema.CurrentVersion"/>.
    /// </exception>
    public async Task MigrateAsync(CancellationToken cancellationToken)
    {
        await using SqliteConnection connection =
            await connectionFactory.OpenConnectionAsync(cancellationToken).ConfigureAwait(false);

        await using SqliteTransaction transaction =
            (SqliteTransaction)await connection.BeginTransactionAsync(System.Data.IsolationLevel.Serializable, cancellationToken).ConfigureAwait(false);

        int existingVersion =
            await ReadExistingSchemaVersionAsync(connection, transaction, cancellationToken).ConfigureAwait(false);

        if (existingVersion > SqliteAuthSchema.CurrentVersion)
        {
            throw new AuthStoreMigrationException(
                $"Auth database schema version {existingVersion} is newer than supported version {SqliteAuthSchema.CurrentVersion}.");
        }

        await ApplySchemaAsync(connection, transaction, cancellationToken).ConfigureAwait(false);
        await WriteSchemaVersionAsync(connection, transaction, cancellationToken).ConfigureAwait(false);

        await transaction.CommitAsync(cancellationToken).ConfigureAwait(false);
    }

    private static async Task<int> ReadExistingSchemaVersionAsync(
        SqliteConnection connection,
        SqliteTransaction transaction,
        CancellationToken cancellationToken)
    {
        await using SqliteCommand tableExistsCommand = connection.CreateCommand();
        tableExistsCommand.Transaction = transaction;
        tableExistsCommand.CommandText = """
            SELECT COUNT(*)
            FROM sqlite_master
            WHERE type = 'table' AND name = $table_name;
            """;
        tableExistsCommand.Parameters.AddWithValue("$table_name", SqliteAuthSchema.SchemaVersionTable);

        long tableCount =
            (long)(await tableExistsCommand.ExecuteScalarAsync(cancellationToken).ConfigureAwait(false) ?? 0L);

        if (tableCount == 0)
        {
            return 0;
        }

        await using SqliteCommand versionCommand = connection.CreateCommand();
        versionCommand.Transaction = transaction;
        versionCommand.CommandText = """
            SELECT version
            FROM schema_version
            WHERE id = 1;
            """;

        object? version = await versionCommand.ExecuteScalarAsync(cancellationToken).ConfigureAwait(false);

        return version is null || version == DBNull.Value
            ? 0
            : Convert.ToInt32(version, CultureInfo.InvariantCulture);
    }

    // Single-shot create of the final schema (all DDL is CREATE ... IF NOT EXISTS, so it is
    // idempotent against an already-provisioned database), followed by any additive column
    // migrations that a pre-existing table would be missing. The applied version is stamped
    // separately by WriteSchemaVersionAsync.
    private static async Task ApplySchemaAsync(
        SqliteConnection connection,
        SqliteTransaction transaction,
        CancellationToken cancellationToken)
    {
        await ExecuteNonQueryAsync(
            connection,
            transaction,
            string.Join(
                "\n",
                SqliteAuthSchema.CreateSchemaVersionTable,
                SqliteAuthSchema.CreateApiKeysTable,
                SqliteAuthSchema.CreateApiKeyAuditTable,
                SqliteAuthSchema.CreateIndexes),
            cancellationToken).ConfigureAwait(false);

        // v3: add the nullable expires_utc column to a pre-v3 api_keys table. On a fresh database the
        // CREATE above already includes the column, so the existence check makes this a no-op.
        await EnsureApiKeysColumnAsync(
            connection,
            transaction,
            SqliteAuthSchema.ExpiresUtcColumn,
            SqliteAuthSchema.AddExpiresUtcColumn,
            cancellationToken).ConfigureAwait(false);
    }

    // Adds a column to api_keys only if it is not already present. SQLite has no
    // "ADD COLUMN IF NOT EXISTS", so we probe PRAGMA table_info first. The reader is fully consumed
    // and disposed before the ALTER runs on the same connection/transaction.
    private static async Task EnsureApiKeysColumnAsync(
        SqliteConnection connection,
        SqliteTransaction transaction,
        string columnName,
        string addColumnDdl,
        CancellationToken cancellationToken)
    {
        bool exists = false;

        await using (SqliteCommand pragma = connection.CreateCommand())
        {
            pragma.Transaction = transaction;
            // table_info takes an identifier, not a bindable parameter; the table name is a fixed
            // internal constant, never caller input.
            pragma.CommandText = $"PRAGMA table_info({SqliteAuthSchema.ApiKeysTable});";

            await using SqliteDataReader reader =
                await pragma.ExecuteReaderAsync(cancellationToken).ConfigureAwait(false);

            while (await reader.ReadAsync(cancellationToken).ConfigureAwait(false))
            {
                // PRAGMA table_info columns: cid(0), name(1), type(2), ...
                if (string.Equals(reader.GetString(1), columnName, StringComparison.OrdinalIgnoreCase))
                {
                    exists = true;
                    break;
                }
            }
        }

        if (!exists)
        {
            await ExecuteNonQueryAsync(connection, transaction, addColumnDdl, cancellationToken).ConfigureAwait(false);
        }
    }

    private static async Task WriteSchemaVersionAsync(
        SqliteConnection connection,
        SqliteTransaction transaction,
        CancellationToken cancellationToken)
    {
        await using SqliteCommand versionCommand = connection.CreateCommand();
        versionCommand.Transaction = transaction;
        versionCommand.CommandText = """
            INSERT INTO schema_version (id, version, applied_utc)
            VALUES (1, $version, $applied_utc)
            ON CONFLICT(id) DO UPDATE SET
                version = excluded.version,
                applied_utc = excluded.applied_utc;
            """;
        versionCommand.Parameters.AddWithValue("$version", SqliteAuthSchema.CurrentVersion);
        versionCommand.Parameters.AddWithValue("$applied_utc", DateTimeOffset.UtcNow.ToString("O"));

        await versionCommand.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
    }

    private static async Task ExecuteNonQueryAsync(
        SqliteConnection connection,
        SqliteTransaction transaction,
        string commandText,
        CancellationToken cancellationToken)
    {
        await using SqliteCommand command = connection.CreateCommand();
        command.Transaction = transaction;
        command.CommandText = commandText;

        await command.ExecuteNonQueryAsync(cancellationToken).ConfigureAwait(false);
    }
}