MDSQLiteAdapter.cs

namespace PowerSync.Common.MDSQLite;

using System;
using System.Collections.Generic;
using System.Threading.Channels;
using System.Threading.Tasks;

using Microsoft.Data.Sqlite;

using Nito.AsyncEx;

using PowerSync.Common.DB;
using PowerSync.Common.Utils;

public class MDSQLiteAdapterOptions()
{
    public string Name { get; set; } = null!;

    public MDSQLiteOptions? SqliteOptions;
}

public class MDSQLiteAdapter : IDBAdapter
{
    public string Name => options.Name;

    public DBAdapterEvents Events { get; } = new();

    // One writer
    private MDSQLiteConnection writeConnection = null!;
    private readonly AsyncLock writeMutex = new();

    // Many readers
    private MDSQLiteConnectionPool readPool = null!;

    private readonly Task initialized;

    protected MDSQLiteAdapterOptions options;

    protected RequiredMDSQLiteOptions resolvedOptions;
    private CancellationTokenSource? tablesUpdatedCts;
    private Task? tablesUpdatedTask;

    public MDSQLiteAdapter(MDSQLiteAdapterOptions options)
    {
        this.options = options;
        resolvedOptions = ResolveMDSQLiteOptions(options.SqliteOptions);
        initialized = Init();
    }

    private RequiredMDSQLiteOptions ResolveMDSQLiteOptions(MDSQLiteOptions? options)
    {
        var defaults = RequiredMDSQLiteOptions.DEFAULT_SQLITE_OPTIONS;
        return new RequiredMDSQLiteOptions
        {
            JournalMode = options?.JournalMode ?? defaults.JournalMode,
            Synchronous = options?.Synchronous ?? defaults.Synchronous,
            JournalSizeLimit = options?.JournalSizeLimit ?? defaults.JournalSizeLimit,
            CacheSizeKb = options?.CacheSizeKb ?? defaults.CacheSizeKb,
            TemporaryStorage = options?.TemporaryStorage ?? defaults.TemporaryStorage,
            LockTimeoutMs = options?.LockTimeoutMs ?? defaults.LockTimeoutMs,
            EncryptionKey = options?.EncryptionKey ?? defaults.EncryptionKey,
            Extensions = options?.Extensions ?? defaults.Extensions,
            LoadPowerSyncExtension = options?.LoadPowerSyncExtension ?? defaults.LoadPowerSyncExtension,
            ReadPoolSize = options?.ReadPoolSize ?? defaults.ReadPoolSize,
        };
    }

    private async Task Init()
    {
        string[] baseStatements =
        [
            $"PRAGMA busy_timeout = {resolvedOptions.LockTimeoutMs}",
            $"PRAGMA cache_size = -{resolvedOptions.CacheSizeKb}",
            $"PRAGMA temp_store = {resolvedOptions.TemporaryStorage}"
        ];

        string[] writeConnectionStatements =
        [
            .. baseStatements,
            $"PRAGMA journal_mode = {resolvedOptions.JournalMode}",
            $"PRAGMA journal_size_limit = {resolvedOptions.JournalSizeLimit}",
            $"PRAGMA synchronous = {resolvedOptions.Synchronous}",
        ];

        string[] readConnectionStatements =
        [
            .. baseStatements,
            "PRAGMA query_only = true",
        ];


        // Prepare write connection
        writeConnection = await OpenConnection(options.Name);
        foreach (var statement in writeConnectionStatements)
        {
            await writeConnection!.Execute(statement);
        }

        // Prepare read pool and create connection factory
        Func<Task<MDSQLiteConnection>> readConnectionFactory = async () =>
        {
            var readConnection = await OpenConnection(options.Name);
            foreach (var statement in readConnectionStatements)
            {
                await readConnection.Execute(statement);
            }
            return readConnection;
        };
        readPool = new MDSQLiteConnectionPool(resolvedOptions.ReadPoolSize, readConnectionFactory);
        await readPool.Init();

        // Register TablesUpdated listener
        tablesUpdatedCts = new CancellationTokenSource();
        tablesUpdatedTask = Task.Run(async () =>
        {
            await foreach (var notification in writeConnection.ListenAsync(tablesUpdatedCts.Token))
            {
                if (notification.TablesUpdated != null)
                {
                    Events.Emit(notification);
                }
            }
        });
    }

    protected async Task<MDSQLiteConnection> OpenConnection(string dbFilename)
    {
        var db = OpenDatabase(dbFilename);
        LoadExtensions(db);

        var connection = new MDSQLiteConnection(new MDSQLiteConnectionOptions(db));
        try
        {
            await connection.Execute("SELECT powersync_init()");
        }
        catch (SqliteException ex)
        {
            // SQLite will throw a very unhelpful "SQLite Error 1: 'The specified
            // module could not be found.'" error if uncaught.
            throw new SqliteException(
                "Failed to initialize PowerSync: powersync_init() is not registered. " +
                "Ensure the PowerSync core SQLite extension is loaded. Either set " +
                "MDSQLiteOptions.LoadPowerSyncExtension to true (default), or supply " +
                "a PowerSync-compatible extension via MDSQLiteOptions.Extensions.",
                ex.SqliteErrorCode,
                ex.SqliteExtendedErrorCode);
        }

        return connection;
    }

    private static SqliteConnection OpenDatabase(string dbFilename)
    {
        string connectionString = $"Data Source={dbFilename};Pooling=False;";
        var connection = new SqliteConnection(connectionString);
        connection.Open();
        return connection;
    }

    protected virtual void LoadExtensions(SqliteConnection db)
    {
        db.EnableExtensions(true);
        if (resolvedOptions.LoadPowerSyncExtension)
        {
            LoadDefaultPowerSyncExtension(db);
        }
        foreach (var extension in resolvedOptions.Extensions)
        {
            db.LoadExtension(extension.Path, extension.EntryPoint);
        }
    }

    /// <summary>
    /// Loads the bundled PowerSync core SQLite extension. Override on
    /// platform-specific adapters (e.g. MAUI iOS/Android) where the native library
    /// lives outside the desktop runtime path.
    /// </summary>
    protected virtual void LoadDefaultPowerSyncExtension(SqliteConnection db)
    {
        var path = PowerSyncPathResolver.GetNativeLibraryPath(AppContext.BaseDirectory);
        db.LoadExtension(path, "sqlite3_powersync_init");
    }

    public async Task Close()
    {
        tablesUpdatedCts?.Cancel();
        try { tablesUpdatedTask?.Wait(2000); } catch { /* expected */ }
        writeConnection?.Close();
        await readPool.Close();
        Events.Close();
    }

    public async Task<NonQueryResult> Execute(string query, object?[]? parameters = null)
    {
        return await WriteLock((ctx) => ctx.Execute(query, parameters));
    }

    public async Task<NonQueryResult> ExecuteBatch(string query, object?[][]? parameters = null)
    {
        return await WriteTransaction((ctx) => ctx.ExecuteBatch(query, parameters));
    }

    public async Task<T[]> GetAll<T>(string sql, object?[]? parameters = null)
    {
        return await ReadLock((ctx) => ctx.GetAll<T>(sql, parameters));
    }

    public async Task<dynamic[]> GetAll(string sql, object?[]? parameters = null)
    {
        return await ReadLock((ctx) => ctx.GetAll(sql, parameters));
    }

    public async Task<T?> GetOptional<T>(string sql, object?[]? parameters = null)
    {
        return await ReadLock((ctx) => ctx.GetOptional<T>(sql, parameters));
    }

    public async Task<dynamic?> GetOptional(string sql, object?[]? parameters = null)
    {
        return await ReadLock((ctx) => ctx.GetOptional(sql, parameters));
    }

    public async Task<T> Get<T>(string sql, object?[]? parameters = null)
    {
        return await ReadLock((ctx) => ctx.Get<T>(sql, parameters));
    }

    public async Task<dynamic> Get(string sql, object?[]? parameters = null)
    {
        return await ReadLock((ctx) => ctx.Get(sql, parameters));
    }

    public async Task<T> ReadTransaction<T>(Func<ITransaction, Task<T>> fn, DBLockOptions? options = null)
    {
        return await ReadLock((ctx) => InternalTransaction(new MDSQLiteTransaction((MDSQLiteConnection)ctx), fn));
    }

    public async Task<T> ReadLock<T>(Func<ILockContext, Task<T>> fn, DBLockOptions? options = null)
    {
        await initialized;
        return await readPool.Lease(fn);
    }

    public async Task WriteLock(Func<ILockContext, Task> fn, DBLockOptions? options = null)
    {
        await initialized;

        using (await writeMutex.LockAsync())
        {
            await fn(writeConnection);
        }

        writeConnection.FlushUpdates();

    }

    public async Task<T> WriteLock<T>(Func<ILockContext, Task<T>> fn, DBLockOptions? options = null)
    {
        await initialized;

        T result;
        using (await writeMutex.LockAsync())
        {
            result = await fn(writeConnection);
        }

        writeConnection.FlushUpdates();

        return result;
    }

    public async Task WriteTransaction(Func<ITransaction, Task> fn, DBLockOptions? options = null)
    {
        await WriteLock(ctx => InternalTransaction(new MDSQLiteTransaction(writeConnection), fn));
    }

    public async Task<T> WriteTransaction<T>(Func<ITransaction, Task<T>> fn, DBLockOptions? options = null)
    {
        return await WriteLock((ctx) => InternalTransaction(new MDSQLiteTransaction(writeConnection), fn));
    }

    protected static async Task InternalTransaction(
        MDSQLiteTransaction ctx,
        Func<ITransaction, Task> fn)
    {
        await RunTransaction(ctx, () => fn(ctx));
    }

    protected static async Task<T> InternalTransaction<T>(
        MDSQLiteTransaction ctx,
        Func<ITransaction, Task<T>> fn)
    {
        T result = default!;
        await RunTransaction(ctx, async () =>
        {
            result = await fn(ctx);
        });

        return result;
    }

    private static async Task RunTransaction(
        MDSQLiteTransaction ctx,
        Func<Task> action)
    {
        try
        {
            await ctx.Begin();
            await action();
            await ctx.Commit();
        }
        catch (Exception)
        {
            // In rare cases, a rollback may fail. Safe to ignore.
            try { await ctx.Rollback(); }
            catch
            {
                // Ignore rollback errors
            }
            throw;
        }
    }

    public async Task RefreshSchema()
    {
        await initialized;
        await writeConnection.RefreshSchema();
        await readPool.LeaseAll(async (connections) =>
        {
            foreach (var conn in connections) await conn.RefreshSchema();
        });
    }
}

class MDSQLiteConnectionPool
{
    private readonly Channel<MDSQLiteConnection> _channel;
    private readonly int _poolSize;
    private readonly Func<Task<MDSQLiteConnection>> _connectionFactory;

    private readonly Task _initialized;

    public MDSQLiteConnectionPool(int poolSize, Func<Task<MDSQLiteConnection>> connectionFactory)
    {
        _channel = Channel.CreateBounded<MDSQLiteConnection>(poolSize);
        _poolSize = poolSize;
        _connectionFactory = connectionFactory;
        _initialized = Initialize();
    }

    public async Task Init() => await _initialized;

    private async Task Initialize()
    {
        for (int i = 0; i < _poolSize; i++)
        {
            var connection = await _connectionFactory();
            await _channel.Writer.WriteAsync(connection);
        }
    }

    public async Task<T> Lease<T>(Func<MDSQLiteConnection, Task<T>> callback)
    {
        await _initialized;
        var connection = await _channel.Reader.ReadAsync();
        try
        {
            return await callback(connection);
        }
        finally
        {
            await _channel.Writer.WriteAsync(connection);
        }
    }

    public async Task LeaseAll(Func<List<MDSQLiteConnection>, Task> callback)
    {
        await _initialized;
        var connections = new List<MDSQLiteConnection>(_poolSize);
        for (int i = 0; i < _poolSize; i++)
        {
            connections.Add(await _channel.Reader.ReadAsync());
        }

        try
        {
            await callback(connections);
        }
        finally
        {
            foreach (var conn in connections)
            {
                _channel.Writer.TryWrite(conn);
            }
        }
    }

    public async Task Close()
    {
        await LeaseAll((connections) =>
        {
            foreach (var conn in connections) conn.Close();
            return Task.CompletedTask;
        });
        _channel.Writer.TryComplete();
    }
}

public class MDSQLiteTransaction(MDSQLiteConnection connection) : ITransaction
{
    private readonly MDSQLiteConnection connection = connection;
    private bool finalized = false;

    public async Task Begin()
    {
        if (finalized) return;
        await connection.Execute("BEGIN");
    }

    public async Task Commit()
    {
        if (finalized) return;
        finalized = true;
        await connection.Execute("COMMIT");
    }

    public async Task Rollback()
    {
        if (finalized) return;
        finalized = true;
        await connection.Execute("ROLLBACK");
    }

    public Task<NonQueryResult> Execute(string query, object?[]? parameters = null)
    {
        return connection.Execute(query, parameters);
    }

    public Task<NonQueryResult> ExecuteBatch(string query, object?[][]? parameters = null)
    {
        return connection.ExecuteBatch(query, parameters);
    }

    public Task<T[]> GetAll<T>(string sql, object?[]? parameters = null)
    {
        return connection.GetAll<T>(sql, parameters);
    }

    public Task<dynamic[]> GetAll(string sql, object?[]? parameters = null)
    {
        return connection.GetAll(sql, parameters);
    }

    public Task<T?> GetOptional<T>(string sql, object?[]? parameters = null)
    {
        return connection.GetOptional<T>(sql, parameters);
    }

    public Task<dynamic?> GetOptional(string sql, object?[]? parameters = null)
    {
        return connection.GetOptional(sql, parameters);
    }

    public Task<T> Get<T>(string sql, object?[]? parameters = null)
    {
        return connection.Get<T>(sql, parameters);
    }

    public Task<dynamic> Get(string sql, object?[]? parameters = null)
    {
        return connection.Get(sql, parameters);
    }
}