InMemorySessionFsSqliteHandler.cs

/*---------------------------------------------------------------------------------------------
 *  Copyright (c) Microsoft Corporation. All rights reserved.
 *--------------------------------------------------------------------------------------------*/

using System.Collections.Concurrent;
using GitHub.Copilot;
using GitHub.Copilot.Rpc;
using Microsoft.Data.Sqlite;

namespace GitHub.Copilot.Test.E2E;

internal record SqliteCall(string SessionId, string QueryType, string Query);

/// <summary>
/// A SessionFsProvider that implements <see cref="ISessionFsSqliteProvider"/> with a real
/// in-memory SQLite database, and uses a simple <see cref="ConcurrentDictionary{TKey,TValue}"/>
/// for file operations instead of touching disk.
/// </summary>
internal sealed class InMemorySessionFsSqliteHandler(string sessionId, List<SqliteCall> sqliteCalls)
    : SessionFsProvider, ISessionFsSqliteProvider, ISessionFsSqliteTransactionProvider
{
    internal ConcurrentDictionary<string, string> Files { get; } = new();
    private readonly ConcurrentDictionary<string, byte> _directories = new();
    private SqliteConnection? _db;

    private SqliteConnection GetOrCreateDb()
    {
        if (_db is not null)
        {
            return _db;
        }

        _db = new SqliteConnection("Data Source=:memory:");
        _db.Open();
        using var cmd = _db.CreateCommand();
        cmd.CommandText = "PRAGMA busy_timeout = 5000";
        cmd.ExecuteNonQuery();
        return _db;
    }

    // ---- ISessionFsSqliteProvider ----

    public Task<SessionFsSqliteResult?> QueryAsync(
        SessionFsSqliteQueryType queryType,
        string query,
        IDictionary<string, object?>? bindParams,
        CancellationToken cancellationToken)
    {
        return Task.FromResult(RunStatement(GetOrCreateDb(), null, queryType, query, bindParams));
    }

    public Task<IList<SessionFsSqliteResult>> TransactionAsync(
        IList<SessionFsSqliteStatement> statements,
        CancellationToken cancellationToken)
    {
        var db = GetOrCreateDb();
        using var transaction = db.BeginTransaction();
        try
        {
            IList<SessionFsSqliteResult> results = statements
                .Select(statement => RunStatement(db, transaction, statement.QueryType, statement.Query, statement.Params)
                    ?? new SessionFsSqliteResult())
                .ToList();
            try
            {
                transaction.Commit();
            }
            catch (Exception ex)
            {
                throw new SessionFsSqliteTransactionException(
                    ex.Message,
                    SessionFsSqliteTransactionErrorClass.PostCommitAmbiguous,
                    ex);
            }
            return Task.FromResult(results);
        }
        catch (SessionFsSqliteTransactionException)
        {
            throw;
        }
        catch (SqliteException ex)
        {
            transaction.Rollback();
            var errorClass = ex.SqliteErrorCode is 5 or 6
                ? SessionFsSqliteTransactionErrorClass.BusyOrLocked
                : SessionFsSqliteTransactionErrorClass.Fatal;
            throw new SessionFsSqliteTransactionException(ex.Message, errorClass, ex);
        }
        catch (Exception ex)
        {
            transaction.Rollback();
            throw new SessionFsSqliteTransactionException(ex.Message, SessionFsSqliteTransactionErrorClass.Fatal, ex);
        }
    }

    private SessionFsSqliteResult? RunStatement(
        SqliteConnection db,
        SqliteTransaction? transaction,
        SessionFsSqliteQueryType queryType,
        string query,
        IDictionary<string, object?>? bindParams)
    {
        sqliteCalls.Add(new SqliteCall(sessionId, queryType.Value, query));

        var trimmed = query.Trim();
        if (trimmed.Length == 0)
        {
            return null;
        }

        if (queryType == SessionFsSqliteQueryType.Exec)
        {
            using var cmd = db.CreateCommand();
            cmd.Transaction = transaction;
            cmd.CommandText = trimmed;
            cmd.ExecuteNonQuery();
            return null;
        }

        if (queryType == SessionFsSqliteQueryType.Query)
        {
            using var cmd = db.CreateCommand();
            cmd.Transaction = transaction;
            cmd.CommandText = trimmed;
            AddParams(cmd, bindParams);

            using var reader = cmd.ExecuteReader();
            var columns = new List<string>();
            for (var i = 0; i < reader.FieldCount; i++)
            {
                columns.Add(reader.GetName(i));
            }

            var rows = new List<IDictionary<string, object>>();
            while (reader.Read())
            {
                var row = new Dictionary<string, object>(reader.FieldCount);
                for (var i = 0; i < reader.FieldCount; i++)
                {
                    row[columns[i]] = reader.IsDBNull(i) ? null! : reader.GetValue(i);
                }
                rows.Add(row);
            }

            return new SessionFsSqliteResult
            {
                Columns = columns,
                Rows = rows,
                RowsAffected = 0,
            };
        }

        if (queryType == SessionFsSqliteQueryType.Run)
        {
            using var cmd = db.CreateCommand();
            cmd.Transaction = transaction;
            cmd.CommandText = trimmed;
            AddParams(cmd, bindParams);

            var rowsAffected = cmd.ExecuteNonQuery();

            using var rowidCmd = db.CreateCommand();
            rowidCmd.Transaction = transaction;
            rowidCmd.CommandText = "SELECT last_insert_rowid()";
            var lastRowid = rowidCmd.ExecuteScalar();

            return new SessionFsSqliteResult
            {
                Columns = [],
                Rows = [],
                RowsAffected = rowsAffected,
                LastInsertRowid = lastRowid is long l ? l : null,
            };
        }

        throw new ArgumentException($"Unknown queryType: {queryType}");
    }

    public Task<bool> ExistsAsync(CancellationToken cancellationToken)
    {
        return Task.FromResult(_db is not null);
    }

    private static void AddParams(SqliteCommand cmd, IDictionary<string, object?>? bindParams)
    {
        if (bindParams is null) return;
        foreach (var (key, value) in bindParams)
        {
            cmd.Parameters.AddWithValue(key.StartsWith(':') || key.StartsWith('$') || key.StartsWith('@') ? key : $":{key}", value ?? DBNull.Value);
        }
    }

    // ---- File operations (in-memory) ----

    private string Resolve(string path) => $"/{sessionId}{(path.StartsWith('/') ? path : "/" + path)}";

    protected override Task<string> ReadFileAsync(string path, CancellationToken cancellationToken)
    {
        var key = Resolve(path);
        if (!Files.TryGetValue(key, out var content))
            throw new FileNotFoundException($"File not found: {path}");
        return Task.FromResult(content);
    }

    protected override Task WriteFileAsync(string path, string content, int? mode, CancellationToken cancellationToken)
    {
        Files[Resolve(path)] = content;
        return Task.CompletedTask;
    }

    protected override Task AppendFileAsync(string path, string content, int? mode, CancellationToken cancellationToken)
    {
        Files.AddOrUpdate(Resolve(path), content, (_, existing) => existing + content);
        return Task.CompletedTask;
    }

    protected override Task<bool> ExistsAsync(string path, CancellationToken cancellationToken)
    {
        var key = Resolve(path);
        return Task.FromResult(Files.ContainsKey(key) || _directories.ContainsKey(key));
    }

    protected override Task<SessionFsStatResult> StatAsync(string path, CancellationToken cancellationToken)
    {
        var key = Resolve(path);
        if (Files.TryGetValue(key, out var fileContent))
            return Task.FromResult(new SessionFsStatResult { IsFile = true, IsDirectory = false, Size = fileContent.Length });
        if (_directories.ContainsKey(key))
            return Task.FromResult(new SessionFsStatResult { IsFile = false, IsDirectory = true, Size = 0 });
        throw new FileNotFoundException($"Path does not exist: {path}");
    }

    protected override Task MakeDirectoryAsync(string path, bool recursive, int? mode, CancellationToken cancellationToken)
    {
        _directories[Resolve(path)] = 0;
        return Task.CompletedTask;
    }

    protected override Task<IList<string>> ReadDirectoryAsync(string path, CancellationToken cancellationToken)
        => Task.FromResult<IList<string>>([]);

    protected override Task<IList<SessionFsReaddirWithTypesEntry>> ReadDirectoryWithTypesAsync(string path, CancellationToken cancellationToken)
        => Task.FromResult<IList<SessionFsReaddirWithTypesEntry>>([]);

    protected override Task RemoveAsync(string path, bool recursive, bool force, CancellationToken cancellationToken)
    {
        var key = Resolve(path);
        Files.TryRemove(key, out _);
        _directories.TryRemove(key, out _);
        return Task.CompletedTask;
    }

    protected override Task RenameAsync(string src, string dest, CancellationToken cancellationToken)
    {
        var srcKey = Resolve(src);
        var destKey = Resolve(dest);
        if (Files.TryRemove(srcKey, out var content))
            Files[destKey] = content;
        return Task.CompletedTask;
    }
}