journal/Journal.Core/Services/Database/DatabaseSessionService.cs

using Microsoft.Data.Sqlite;

namespace Journal.Core.Services.Database;

public sealed class DatabaseSessionService(IJournalDatabaseService database) : IDatabaseSessionService, IDisposable
{
    private readonly IJournalDatabaseService _database = database;
    private readonly Lock _lock = new();
    private string? _password;
    private string? _dataDirectory;
    private SqliteConnection? _connection;

    public bool IsUnlocked
    {
        get
        {
            lock (_lock) { return _password is not null; }
        }
    }

    public void SetPassword(string password, string? dataDirectory = null)
    {
        if (string.IsNullOrWhiteSpace(password))
            throw new ArgumentException("Password cannot be empty.", nameof(password));

        lock (_lock)
        {
            // If password or directory changed, close the old connection
            if (_connection is not null &&
                (_password != password || _dataDirectory != dataDirectory))
            {
                _connection.Dispose();
                _connection = null;
            }

            _password = password;
            _dataDirectory = dataDirectory;
        }
    }

    public SqliteConnection GetConnection()
    {
        lock (_lock)
        {
            if (_password is null)
                throw new InvalidOperationException(
                    "Database is locked. Authenticate first (e.g. vault.load_all or db.hydrate_workspace).");

            if (_connection is not null)
                return _connection;

            _connection = OpenEncryptedConnection(_password, _dataDirectory);
            EnsureSchema(_connection, _database);
            return _connection;
        }
    }

    public void Dispose()
    {
        lock (_lock)
        {
            _connection?.Dispose();
            _connection = null;
        }
    }

    private SqliteConnection OpenEncryptedConnection(string password, string? dataDirectory)
    {
        var dbPath = _database.GetDatabasePath(dataDirectory);
        var directory = Path.GetDirectoryName(dbPath);
        if (!string.IsNullOrWhiteSpace(directory))
            Directory.CreateDirectory(directory);

        var connection = new SqliteConnection(
            $"Data Source={dbPath};Mode=ReadWriteCreate;Pooling=False");
        connection.Open();

        using var keyCmd = connection.CreateCommand();
        keyCmd.CommandText = _database.BuildPragmaKeyStatement(password) + ";";
        keyCmd.ExecuteNonQuery();

        // Verify the key is correct
        using var verifyCmd = connection.CreateCommand();
        verifyCmd.CommandText = "SELECT count(*) FROM sqlite_master;";
        _ = verifyCmd.ExecuteScalar();

        return connection;
    }

    private static void EnsureSchema(SqliteConnection connection, IJournalDatabaseService database)
    {
        foreach (var statement in database.GetSchemaStatements().Values)
        {
            using var cmd = connection.CreateCommand();
            cmd.CommandText = statement;
            cmd.ExecuteNonQuery();
        }
    }
}