annotations/Azaion.Common/Database/DbFactory.cs

using System.Data.SQLite;
using System.Diagnostics;
using System.IO;
using Azaion.Common.DTO;
using Azaion.Common.DTO.Config;
using Azaion.Common.Extensions;
using LinqToDB;
using LinqToDB.DataProvider.SQLite;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using Serilog;

namespace Azaion.Common.Database;

public interface IDbFactory
{
    Task<T> Run<T>(Func<AnnotationsDb, Task<T>> func);
    Task RunWrite(Func<AnnotationsDb, Task> func);
    Task<T> RunWrite<T>(Func<AnnotationsDb, Task<T>> func);
    Task DeleteAnnotations(List<string> annotationNames, CancellationToken cancellationToken = default);
}

public class DbFactory : IDbFactory
{
    private readonly ILogger<DbFactory> _logger;
    private readonly AnnotationConfig _annConfig;

    private string MemoryConnStr => "Data Source=:memory:";
    private readonly SQLiteConnection _memoryConnection;
    private readonly DataOptions _memoryDataOptions;

    private string FileConnStr   => $"Data Source={_annConfig.AnnotationsDbFile}";
    private readonly SQLiteConnection _fileConnection;
    private readonly DataOptions _fileDataOptions;

    private static readonly SemaphoreSlim WriteSemaphore = new(1, 1);
    private static readonly Guid SaveTaskId = Guid.NewGuid();

    public DbFactory(IOptions<AnnotationConfig> annConfig, ILogger<DbFactory> logger)
    {
        _logger = logger;
        _annConfig = annConfig.Value;

        _memoryConnection = new SQLiteConnection(MemoryConnStr);
        _memoryConnection.Open();
        _memoryDataOptions = new DataOptions()
                .UseDataProvider(SQLiteTools.GetDataProvider())
                .UseConnection(_memoryConnection)
                .UseMappingSchema(AnnotationsDbSchemaHolder.MappingSchema)
                .UseTracing(TraceLevel.Info, t => logger.LogInformation(t.SqlText));


        _fileConnection = new SQLiteConnection(FileConnStr);
        _fileDataOptions = new DataOptions()
            .UseDataProvider(SQLiteTools.GetDataProvider())
            .UseConnection(_fileConnection)
            .UseMappingSchema(AnnotationsDbSchemaHolder.MappingSchema);

        if (!File.Exists(_annConfig.AnnotationsDbFile))
            SQLiteConnection.CreateFile(_annConfig.AnnotationsDbFile);
        RecreateTables();
        
        _fileConnection.Open();
        using var db = new AnnotationsDb(_fileDataOptions);
        var entityTypes = typeof(AnnotationsDb)
            .GetProperties()
            .Where(p => p.PropertyType.IsGenericType && p.PropertyType.GetGenericTypeDefinition() == typeof(ITable<>))
            .Select(p => p.PropertyType.GetGenericArguments()[0])
            .ToArray();
        SchemaMigrator.EnsureSchemaUpdated(db, entityTypes);
        _fileConnection.BackupDatabase(_memoryConnection, "main", "main", -1, null, -1);
    }

    private void RecreateTables()
    {
        using var db = new AnnotationsDb(_fileDataOptions);
        var schema = db.DataProvider.GetSchemaProvider().GetSchema(db);
        var existingTables = schema.Tables.Select(x => x.TableName).ToHashSet();
        if (!existingTables.Contains(Constants.ANNOTATIONS_TABLENAME))
            db.CreateTable<Annotation>();
        if (!existingTables.Contains(Constants.DETECTIONS_TABLENAME))
            db.CreateTable<Detection>();
        if (!existingTables.Contains(Constants.ANNOTATIONS_QUEUE_TABLENAME))
            db.CreateTable<AnnotationQueueRecord>();
    }

    public async Task<T> Run<T>(Func<AnnotationsDb, Task<T>> func)
    {
        await using var db = new AnnotationsDb(_memoryDataOptions);
        return await func(db);
    }

    public async Task RunWrite(Func<AnnotationsDb, Task> func)
    {
        await WriteSemaphore.WaitAsync();
        try
        {
            await using var db = new AnnotationsDb(_memoryDataOptions);
            await func(db);
            ThrottleExt.Throttle(async () =>
            {
                _memoryConnection.BackupDatabase(_fileConnection, "main", "main", -1, null, -1);
                await Task.CompletedTask;
            }, SaveTaskId, TimeSpan.FromSeconds(5), true);
        }
        catch (Exception e)
        {
            _logger.LogError(e, e.Message);
            throw;
        }
        finally
        {
            WriteSemaphore.Release();
        }
    }

    public async Task<T> RunWrite<T>(Func<AnnotationsDb, Task<T>> func)
    {
        await WriteSemaphore.WaitAsync();
        try
        {
            await using var db = new AnnotationsDb(_memoryDataOptions);
            var result = await func(db);
            ThrottleExt.Throttle(async () =>
            {
                _memoryConnection.BackupDatabase(_fileConnection, "main", "main", -1, null, -1);
                await Task.CompletedTask;
            }, SaveTaskId, TimeSpan.FromSeconds(5), true);
            return result;
        }
        catch (Exception e)
        {
            _logger.LogError(e, e.Message);
            throw;
        }
        finally
        {
            WriteSemaphore.Release();
        }
    }

    public async Task DeleteAnnotations(List<string> annotationNames, CancellationToken cancellationToken = default)
    {
        await RunWrite(async db =>
        {
            var detDeleted = await db.Detections.DeleteAsync(x => annotationNames.Contains(x.AnnotationName), token: cancellationToken);
            var annDeleted = await db.Annotations.DeleteAsync(x => annotationNames.Contains(x.Name), token: cancellationToken);
            _logger.LogInformation($"Deleted {detDeleted} detections, {annDeleted} annotations");
        });
    }
}