MoShiro/backend/src/MoShiro.Backend.Storage.Sqlite/SqliteTestStore.cs

using Fleet.Core;
using Microsoft.Data.Sqlite;
using MoShiro.Backend.Application;
using MoShiro.Protocol;

namespace MoShiro.Backend.Storage.Sqlite;

public sealed class SqliteTestStore : ITestStore
{
    private const string Columns = "id, agent_id, plan, params, trigger, status, error_code, phase, phase_count, requested_at, started_at, finished_at, probe_id, ip_version";

    private readonly SqliteDatabase _db;

    public SqliteTestStore(SqliteDatabase db)
    {
        _db = db;
    }

    public ValueTask CreateAsync(TestRecord test, CancellationToken ct) =>
        _db.ExecuteAsync(
            c => c.Exec(
                "INSERT INTO tests (id, agent_id, probe_id, plan, params, trigger, status, error_code, phase, phase_count, requested_at, ip_version) VALUES (@id, @agent, @probe, @plan, @params, @trigger, @status, @err, @phase, @count, @at, @ip)",
                ("@id", test.Id.Value), ("@agent", test.AgentId.Value), ("@probe", test.ProbeId.Value), ("@plan", test.Plan), ("@params", test.ParamsJson), ("@trigger", test.Trigger),
                ("@status", test.Status), ("@err", test.ErrorCode), ("@phase", test.Phase), ("@count", test.PhaseCount), ("@at", test.RequestedAtMs), ("@ip", test.IpVersion)),
            ct);

    public ValueTask<TestRecord?> GetAsync(TestId id, CancellationToken ct) =>
        _db.ReadAsync(c => c.QueryOne($"SELECT {Columns} FROM tests WHERE id = @id", Read, ("@id", id.Value)), ct);

    public ValueTask<IReadOnlyList<TestRecord>> ListAsync(NodeId? agentId, ProbeId? probeId, int limit, CancellationToken ct) =>
        ListAsync(agentId, probeId, limit, null, ct);

    public ValueTask<IReadOnlyList<TestRecord>> ListAsync(NodeId? agentId, ProbeId? probeId, int limit, string? ipVersion, CancellationToken ct) =>
        _db.ReadAsync<IReadOnlyList<TestRecord>>(
            c => c.Query(
                $"SELECT {Columns} FROM tests WHERE (@agent IS NULL OR agent_id = @agent) AND (@probe IS NULL OR probe_id = @probe) AND (@ip IS NULL OR ip_version = @ip) ORDER BY requested_at DESC, id DESC LIMIT @limit",
                Read, ("@agent", agentId?.Value), ("@probe", probeId?.Value), ("@limit", limit), ("@ip", ipVersion)),
            ct);

    public ValueTask<bool> UpdateStatusAsync(TestId id, string status, int? phase, string? errorCode, long nowMs, CancellationToken ct) =>
        _db.WriteAsync(
            c => c.Exec(
                "UPDATE tests SET status = @status, phase = COALESCE(@phase, phase), error_code = COALESCE(@err, error_code), " +
                "started_at = CASE WHEN @status = 'running' AND started_at IS NULL THEN @now ELSE started_at END, " +
                "finished_at = CASE WHEN @status IN ('succeeded', 'failed', 'aborted') THEN @now ELSE finished_at END WHERE id = @id",
                ("@id", id.Value), ("@status", status), ("@phase", phase), ("@err", errorCode), ("@now", nowMs)) == 1,
            ct);

    public ValueTask SetPhaseCountAsync(TestId id, int phaseCount, CancellationToken ct) =>
        _db.ExecuteAsync(c => c.Exec("UPDATE tests SET phase_count = @n WHERE id = @id", ("@id", id.Value), ("@n", phaseCount)), ct);

    public ValueTask<bool> HasActiveAsync(NodeId agentId, ProbeId probeId, CancellationToken ct) =>
        _db.ReadAsync(c => c.Scalar<long>("SELECT COUNT(*) FROM tests WHERE agent_id = @a AND probe_id = @p AND status IN ('queued', 'running')", ("@a", agentId.Value), ("@p", probeId.Value)) > 0, ct);

    public ValueTask SaveResultAsync(TestResultRecord result, CancellationToken ct) =>
        _db.ExecuteAsync(
            c => c.Exec(
                "INSERT INTO test_results (test_id, metrics, series) VALUES (@id, @m, @s) ON CONFLICT(test_id) DO UPDATE SET metrics = excluded.metrics, series = excluded.series",
                ("@id", result.TestId.Value), ("@m", result.MetricsJson), ("@s", result.SeriesJson)),
            ct);

    public ValueTask<TestResultRecord?> GetResultAsync(TestId id, CancellationToken ct) =>
        _db.ReadAsync(
            c => c.QueryOne("SELECT test_id, metrics, series FROM test_results WHERE test_id = @id", r => new TestResultRecord(new TestId(r.GetString(0)), r.Bytes(1), r.Bytes(2)), ("@id", id.Value)),
            ct);

    public ValueTask<IReadOnlyList<TestRecord>> ListNeedingResultAsync(int limit, CancellationToken ct) =>
        _db.ReadAsync<IReadOnlyList<TestRecord>>(
            c => c.Query(
                $"SELECT {Columns} FROM tests WHERE status = 'succeeded' AND id NOT IN (SELECT test_id FROM test_results) ORDER BY finished_at LIMIT @limit",
                Read, ("@limit", limit)),
            ct);

    public ValueTask<int> FailStaleAsync(long olderThanMs, string errorCode, long nowMs, CancellationToken ct) =>
        _db.WriteAsync(
            c => c.Exec(
                "UPDATE tests SET status = 'failed', error_code = @err, finished_at = @now WHERE status IN ('queued', 'running') AND requested_at < @cutoff",
                ("@err", errorCode), ("@now", nowMs), ("@cutoff", olderThanMs)),
            ct);

    private static TestRecord Read(SqliteDataReader r) => new(
        new TestId(r.GetString(0)),
        new NodeId(r.GetString(1)),
        new ProbeId(r.GetString(12)),
        r.GetString(2),
        r.GetString(3),
        r.GetString(4),
        r.GetString(5),
        r.NullableString(6),
        r.IsDBNull(7) ? null : r.GetInt32(7),
        r.GetInt32(8),
        r.GetInt64(9),
        r.NullableLong(10),
        r.NullableLong(11),
        r.GetString(13));
}