feat: EF Migrations Infrastructure

Author: KirillKurdyukov

src/EfCore.Ydb/src/Migrations/Internal/YdbHistoryRepository.cs

@@ -1,19 +1,27 @@
 using System;
-using System.Data;
+using System.Collections.Generic;
 using System.Diagnostics;
-using System.Linq;
+using System.Globalization;
 using System.Threading;
 using System.Threading.Tasks;
 using EfCore.Ydb.Storage.Internal;
 using Microsoft.EntityFrameworkCore.Diagnostics;
 using Microsoft.EntityFrameworkCore.Migrations;
+using Microsoft.EntityFrameworkCore.Migrations.Operations;
+using Microsoft.Extensions.Logging;
+using Ydb.Sdk;
 using Ydb.Sdk.Ado;
 
 namespace EfCore.Ydb.Migrations.Internal;
 
-// ReSharper disable once ClassNeverInstantiated.Global
-public class YdbHistoryRepository(HistoryRepositoryDependencies dependencies) : HistoryRepository(dependencies)
+public class YdbHistoryRepository(HistoryRepositoryDependencies dependencies)
+    : HistoryRepository(dependencies), IHistoryRepository
 {
+    private const string LockKey = "LockMigration";
+    private const int ReleaseMaxAttempt = 10;
+
+    private static readonly TimeSpan LockTimeout = TimeSpan.FromMinutes(2);
+
     protected override bool InterpretExistsResult(object? value)
         => throw new InvalidOperationException("Shouldn't be called");
 
@@ -25,152 +33,141 @@ public override async Task<IMigrationsDatabaseLock> AcquireDatabaseLockAsync(
     )
     {
         Dependencies.MigrationsLogger.AcquiringMigrationLock();
-        var dbLock =
-            new YdbMigrationDatabaseLock("migrationLock", this, (YdbRelationalConnection)Dependencies.Connection);
-        await dbLock.Lock(timeoutInSeconds: 60, cancellationToken);
-        return dbLock;
-    }
 
-    public override string GetCreateIfNotExistsScript()
-        => GetCreateScript().Replace("CREATE TABLE", "CREATE TABLE IF NOT EXISTS");
+        var deadline = DateTime.UtcNow + LockTimeout;
+        DateTime now;
 
-    public override LockReleaseBehavior LockReleaseBehavior => LockReleaseBehavior.Transaction;
+        do
+        {
+            now = DateTime.UtcNow;
 
-    protected override string ExistsSql
-        => throw new UnreachableException("Shouldn't be called. We check if exists using different approach");
+            await using var connection = ((IYdbRelationalConnection)Dependencies.Connection).Clone().DbConnection;
+            try
+            {
+                await connection.OpenAsync(cancellationToken);
 
-    public override bool Exists()
-        => ExistsAsync().ConfigureAwait(false).GetAwaiter().GetResult();
+                var command = connection.CreateCommand();
+                command.CommandText = GetInsertScript(new HistoryRow(LockKey,
+                    $"LockTime: {DateTime.UtcNow.ToString(CultureInfo.InvariantCulture)}, PID: {Environment.ProcessId}"));
+                await command.ExecuteNonQueryAsync(cancellationToken);
 
-    public override Task<bool> ExistsAsync(CancellationToken cancellationToken = default)
-    {
-        var connection = (YdbRelationalConnection)Dependencies.Connection;
-        var schema = (YdbConnection)connection.DbConnection;
-        var tables = schema.GetSchema("tables");
-
-        var foundTables =
-            from table in tables.AsEnumerable()
-            where table.Field<string>("table_type") == "TABLE"
-                  && table.Field<string>("table_name") == TableName
-            select table;
-        return Task.FromResult(foundTables.Count() == 1);
+                return new YdbMigrationDatabaseLock(this);
+            }
+            catch (YdbException)
+            {
+                await Task.Delay(100 + Random.Shared.Next(1000), cancellationToken);
+            }
+        } while (now < deadline);
+
+        throw new YdbException("Unable to obtain table lock - another EF instance may be running");
     }
 
-    public override string GetBeginIfNotExistsScript(string migrationId) => throw new NotImplementedException();
+    private async Task ReleaseDatabaseLockAsync()
+    {
+        for (var i = 0; i < ReleaseMaxAttempt; i++)
+        {
+            await using var connection = ((IYdbRelationalConnection)Dependencies.Connection).Clone().DbConnection;
 
-    public override string GetBeginIfExistsScript(string migrationId) => throw new NotImplementedException();
+            try
+            {
+                await connection.OpenAsync();
+                var command = connection.CreateCommand();
+                command.CommandText = GetDeleteScript(LockKey);
+                await command.ExecuteNonQueryAsync();
+                
+                return;
+            }
+            catch (YdbException e)
+            {
+                Dependencies.MigrationsLogger.Logger.LogError(e, "Failed release database lock");
+            }
+        }
+    }
 
-    public override string GetEndIfScript() => throw new NotImplementedException();
+    bool IHistoryRepository.CreateIfNotExists() => CreateIfNotExistsAsync().GetAwaiter().GetResult();
 
-    private sealed class YdbMigrationDatabaseLock(
-        string name,
-        IHistoryRepository historyRepository,
-        YdbRelationalConnection ydbConnection
-    ) : IMigrationsDatabaseLock
+    public async Task<bool> CreateIfNotExistsAsync(CancellationToken cancellationToken = default)
     {
-        private IYdbRelationalConnection Connection { get; } = ydbConnection.Clone();
-        private volatile string _pid = null!;
-        private CancellationTokenSource? _watchDogToken;
+        if (await ExistsAsync(cancellationToken))
+        {
+            return false;
+        }
 
-        public async Task Lock(int timeoutInSeconds, CancellationToken cancellationToken = default)
+        try
         {
-            if (_watchDogToken != null)
+            await Dependencies.MigrationCommandExecutor.ExecuteNonQueryAsync(
+                GetCreateIfNotExistsCommands(),
+                Dependencies.Connection,
+                new MigrationExecutionState(),
+                true,
+                cancellationToken: cancellationToken
+            ).ConfigureAwait(false);
+
+            return true;
+        }
+        catch (YdbException e)
+        {
+            if (e.Code == StatusCode.Overloaded)
             {
-                throw new InvalidOperationException("Already locked");
+                return true;
             }
 
-            await Connection.OpenAsync(cancellationToken);
-            await using (var command = Connection.DbConnection.CreateCommand())
+            throw;
+        }
+    }
+
+    private IReadOnlyList<MigrationCommand> GetCreateIfNotExistsCommands() =>
+        Dependencies.MigrationsSqlGenerator.Generate(new List<MigrationOperation>
+        {
+            new SqlOperation
             {
-                command.CommandText = """
-                                      CREATE TABLE IF NOT EXISTS shedlock (
-                                          name Text NOT NULL,
-                                          locked_at Timestamp NOT NULL,
-                                          lock_until Timestamp NOT NULL,
-                                          locked_by Text NOT NULL,
-                                          PRIMARY KEY(name)
-                                      );
-                                      """;
-                await command.ExecuteNonQueryAsync(cancellationToken);
+                Sql = GetCreateIfNotExistsScript(),
+                SuppressTransaction = true
             }
+        });
 
-            _pid = $"PID:{Environment.ProcessId}";
+    public override string GetCreateIfNotExistsScript()
+        => GetCreateScript().Replace("CREATE TABLE", "CREATE TABLE IF NOT EXISTS");
 
-            var lockAcquired = false;
-            for (var i = 0; i < 10; i++)
-            {
-                if (await UpdateLock(name, timeoutInSeconds))
-                {
-                    lockAcquired = true;
-                    break;
-                }
+    public override LockReleaseBehavior LockReleaseBehavior => LockReleaseBehavior.Transaction;
 
-                await Task.Delay(TimeSpan.FromMilliseconds(500), cancellationToken);
-            }
+    protected override string ExistsSql
+        => throw new UnreachableException("Shouldn't be called. We check if exists using different approach");
 
-            if (!lockAcquired)
-            {
-                throw new TimeoutException("Failed to acquire lock for migration`");
-            }
+    public override bool Exists()
+        => ExistsAsync().ConfigureAwait(false).GetAwaiter().GetResult();
 
-            _watchDogToken = new CancellationTokenSource();
-            _ = Task.Run((async Task () =>
-            {
-                while (true)
-                {
-                    // ReSharper disable once PossibleLossOfFraction
-                    await Task.Delay(TimeSpan.FromSeconds(timeoutInSeconds / 2), _watchDogToken.Token);
-                    await UpdateLock(name, timeoutInSeconds);
-                }
-                // ReSharper disable once FunctionNeverReturns
-            })!, _watchDogToken.Token);
-        }
+    public override async Task<bool> ExistsAsync(CancellationToken cancellationToken = default)
+    {
+        await using var ydbConnection = (YdbConnection)
+            ((IYdbRelationalConnection)Dependencies.Connection).Clone().DbConnection;
+        await ydbConnection.OpenAsync(cancellationToken);
 
-        private async Task<bool> UpdateLock(string nameLock, int timeoutInSeconds)
+        try
         {
-            var command = Connection.DbConnection.CreateCommand();
-            command.CommandText =
-                $"""
-                 UPSERT INTO shedlock (name, locked_at, lock_until, locked_by)
-                 VALUES (
-                        @name,
-                        CurrentUtcTimestamp(), 
-                        Unwrap(CurrentUtcTimestamp() + Interval("PT{timeoutInSeconds}S")),
-                        @locked_by
-                        );
-                 """;
-            command.Parameters.Add(new YdbParameter("name", DbType.String, nameLock));
-            command.Parameters.Add(new YdbParameter("locked_by", DbType.String, _pid));
+            await new YdbCommand(ydbConnection)
+                    { CommandText = $"SELECT * FROM {TableName} WHERE MigrationId = '{LockKey}';" }
+                .ExecuteNonQueryAsync(cancellationToken);
 
-            try
-            {
-                await command.ExecuteNonQueryAsync();
-                return true;
-            }
-            catch (YdbException)
-            {
-                return false;
-            }
+            return true;
         }
+        catch (YdbException)
+        {
+            return false;
+        }
+    }
 
-        public void Dispose()
-            => DisposeInternalAsync().GetAwaiter().GetResult();
+    public override string GetBeginIfNotExistsScript(string migrationId) => throw new NotSupportedException();
 
-        public async ValueTask DisposeAsync()
-            => await DisposeInternalAsync();
+    public override string GetBeginIfExistsScript(string migrationId) => throw new NotSupportedException();
 
-        private async Task DisposeInternalAsync()
-        {
-            if (_watchDogToken != null)
-            {
-                await _watchDogToken.CancelAsync();
-            }
+    public override string GetEndIfScript() => throw new NotSupportedException();
 
-            _watchDogToken = null;
-            await using var connection = Connection.DbConnection.CreateCommand();
-            connection.CommandText = "DELETE FROM shedlock WHERE name = '{_name}' AND locked_by = '{PID}';";
-            await connection.ExecuteNonQueryAsync();
-        }
+    private sealed class YdbMigrationDatabaseLock(YdbHistoryRepository historyRepository) : IMigrationsDatabaseLock
+    {
+        public void Dispose() => historyRepository.ReleaseDatabaseLockAsync().GetAwaiter().GetResult();
+
+        public async ValueTask DisposeAsync() => await historyRepository.ReleaseDatabaseLockAsync();
 
         public IHistoryRepository HistoryRepository { get; } = historyRepository;
     }

src/EfCore.Ydb/src/Storage/Internal/YdbDatabaseCreator.cs

@@ -1,24 +1,26 @@
 using System;
+using System.Data;
+using System.Linq;
 using System.Threading;
 using System.Threading.Tasks;
+using Microsoft.EntityFrameworkCore.Migrations.Operations;
 using Microsoft.EntityFrameworkCore.Storage;
 using Ydb.Sdk.Ado;
 
 namespace EfCore.Ydb.Storage.Internal;
 
-public class YdbDatabaseCreator(
-    RelationalDatabaseCreatorDependencies dependencies
-) : RelationalDatabaseCreator(dependencies)
+public class YdbDatabaseCreator(RelationalDatabaseCreatorDependencies dependencies)
+    : RelationalDatabaseCreator(dependencies)
 {
-    public override bool Exists()
-        => ExistsInternal().GetAwaiter().GetResult();
+    public override bool Exists() => ExistsInternal().GetAwaiter().GetResult();
 
-    public override Task<bool> ExistsAsync(CancellationToken cancellationToken = new())
-        => ExistsInternal(cancellationToken);
+    public override Task<bool> ExistsAsync(CancellationToken cancellationToken = default) =>
+        ExistsInternal(cancellationToken);
 
     private async Task<bool> ExistsInternal(CancellationToken cancellationToken = default)
     {
-        await using var connection = Dependencies.Connection;
+        await using var connection = ((IYdbRelationalConnection)Dependencies.Connection).Clone();
+
         try
         {
             await connection.OpenAsync(cancellationToken, errorsExpected: true);
@@ -30,9 +32,55 @@ private async Task<bool> ExistsInternal(CancellationToken cancellationToken = de
         }
     }
 
-    public override bool HasTables() => false;
+    public override bool HasTables() => HasTablesAsync().GetAwaiter().GetResult();
+
+    public override async Task<bool> HasTablesAsync(CancellationToken cancellationToken = default)
+    {
+        await using var connection = ((IYdbRelationalConnection)Dependencies.Connection).Clone();
+
+        try
+        {
+            await connection.OpenAsync(cancellationToken, errorsExpected: true);
 
-    public override void Create() => throw new NotSupportedException("YDB does not support database creation");
+            var dataTable = await connection
+                .DbConnection
+                .GetSchemaAsync("Tables", [null, "TABLE"], cancellationToken);
+
+            return dataTable.Rows.Count > 0;
+        }
+        catch (YdbException)
+        {
+            return false;
+        }
+    }
 
-    public override void Delete() => throw new NotSupportedException("YDB does not support database deletion");
+    public override void Create() => CreateAsync().GetAwaiter().GetResult();
+
+    public override async Task CreateAsync(CancellationToken cancellationToken = default)
+    {
+        if (await ExistsAsync(cancellationToken))
+        {
+            return;
+        }
+
+        throw new NotSupportedException("YDB does not support database creation");
+    }
+
+    public override void Delete() => DeleteAsync().GetAwaiter().GetResult();
+
+    public override async Task DeleteAsync(CancellationToken cancellationToken = default)
+    {
+        await using var connection = Dependencies.Connection;
+        await connection.OpenAsync(cancellationToken, errorsExpected: true);
+
+        var dataTable = await connection
+            .DbConnection
+            .GetSchemaAsync("Tables", [null, "TABLE"], cancellationToken);
+
+        var dropTableOperations = (from DataRow row in dataTable.Rows
+            select new DropTableOperation { Name = row["table_name"].ToString() }).ToList();
+
+        await Dependencies.MigrationCommandExecutor.ExecuteNonQueryAsync(Dependencies.MigrationsSqlGenerator
+            .Generate(dropTableOperations), connection, cancellationToken);
+    }
 }