KArtSell.Aegis/src/KArtSell.Modules.ModelOperations/Infrastructure/DapperModelScheduleRepository.cs

using Dapper;
using KArtSell.BuildingBlocks.Data;
using KArtSell.Modules.ModelOperations.Application;

namespace KArtSell.Modules.ModelOperations.Infrastructure;

public sealed class DapperModelScheduleRepository(IDbConnectionFactory connectionFactory) : IModelScheduleRepository
{
    private const string AcquireSql = """
        with due as (
            select schedule_id
            from evaluation.model_operation_schedule
            where enabled = true
              and next_due_at <= @Now
              and (lease_until is null or lease_until < @Now)
            order by next_due_at, operation_code, scope_key
            for update skip locked
            limit @Limit
        )
        update evaluation.model_operation_schedule s
        set lease_owner = @LeaseOwner,
            lease_until = @LeaseUntil,
            dispatch_revision = dispatch_revision + 1,
            updated_at = @Now
        from due
        where s.schedule_id = due.schedule_id
        returning s.schedule_id as ScheduleId,
                  s.operation_code as OperationCode,
                  s.scope_key as ScopeKey,
                  s.cadence as Cadence,
                  s.automation_mode as AutomationMode,
                  s.queue_name as Queue,
                  concat(s.operation_code, ':', s.scope_key, ':', s.schedule_version, ':', to_char(s.next_due_at at time zone 'UTC', 'YYYYMMDDHH24MISS')) as IdempotencyKey,
                  s.schedule_version as ScheduleVersion,
                  s.dispatch_revision as DispatchRevision,
                  s.next_due_at as ScheduledFor,
                  s.catch_up_policy as CatchUpPolicy,
                  s.max_catch_up as MaxCatchUp;
        """;

    private const string DispatchedSql = """
        update evaluation.model_operation_schedule
        set last_dispatched_at = @DispatchedAt,
            last_background_job_id = @BackgroundJobId,
            next_due_at = @NextDueAt,
            lease_owner = null,
            lease_until = null,
            last_error_code = null,
            updated_at = @DispatchedAt
        where schedule_id = @ScheduleId
          and lease_owner = @LeaseOwner
          and dispatch_revision = @ExpectedDispatchRevision;
        """;

    private const string ReleaseSql = """
        update evaluation.model_operation_schedule
        set lease_owner = null,
            lease_until = null,
            last_error_code = @ReasonCode,
            next_due_at = greatest(next_due_at, @ReleasedAt) + interval '1 hour',
            updated_at = @ReleasedAt
        where schedule_id = @ScheduleId
          and lease_owner = @LeaseOwner
          and dispatch_revision = @ExpectedDispatchRevision;
        """;

    private const string AdvanceSql = """
        update evaluation.model_operation_schedule
        set next_due_at = @NextDueAt,
            lease_owner = null,
            lease_until = null,
            last_error_code = null,
            updated_at = @AdvancedAt
        where schedule_id = @ScheduleId
          and lease_owner = @LeaseOwner
          and dispatch_revision = @ExpectedDispatchRevision;
        """;

    public async Task<IReadOnlyList<DueModelOperation>> AcquireDueAsync(DateTimeOffset now, string leaseOwner, TimeSpan leaseDuration, int limit, CancellationToken cancellationToken)
    {
        await using var connection = await connectionFactory.OpenAsync(cancellationToken);
        var items = await connection.QueryAsync<DueModelOperation>(new CommandDefinition(AcquireSql, new { Now = now, LeaseOwner = leaseOwner, LeaseUntil = now.Add(leaseDuration), Limit = limit }, cancellationToken: cancellationToken));
        return items.AsList();
    }

    public async Task MarkDispatchedAsync(Guid scheduleId, string leaseOwner, int expectedDispatchRevision, string backgroundJobId, DateTimeOffset dispatchedAt, DateTimeOffset nextDueAt, CancellationToken cancellationToken)
    {
        await using var connection = await connectionFactory.OpenAsync(cancellationToken);
        var affected = await connection.ExecuteAsync(new CommandDefinition(DispatchedSql, new { ScheduleId = scheduleId, LeaseOwner = leaseOwner, ExpectedDispatchRevision = expectedDispatchRevision, BackgroundJobId = backgroundJobId, DispatchedAt = dispatchedAt, NextDueAt = nextDueAt }, cancellationToken: cancellationToken));
        if (affected != 1) throw new InvalidOperationException("Schedule lease was lost before dispatch completion.");
    }

    public async Task AdvanceWithoutDispatchAsync(Guid scheduleId, string leaseOwner, int expectedDispatchRevision, DateTimeOffset advancedAt, DateTimeOffset nextDueAt, CancellationToken cancellationToken)
    {
        await using var connection = await connectionFactory.OpenAsync(cancellationToken);
        var affected = await connection.ExecuteAsync(new CommandDefinition(AdvanceSql, new { ScheduleId = scheduleId, LeaseOwner = leaseOwner, ExpectedDispatchRevision = expectedDispatchRevision, AdvancedAt = advancedAt, NextDueAt = nextDueAt }, cancellationToken: cancellationToken));
        if (affected != 1) throw new InvalidOperationException("Schedule lease was lost before advance completion.");
    }

    public async Task ReleaseAsync(Guid scheduleId, string leaseOwner, int expectedDispatchRevision, string reasonCode, DateTimeOffset releasedAt, CancellationToken cancellationToken)
    {
        await using var connection = await connectionFactory.OpenAsync(cancellationToken);
        await connection.ExecuteAsync(new CommandDefinition(ReleaseSql, new { ScheduleId = scheduleId, LeaseOwner = leaseOwner, ExpectedDispatchRevision = expectedDispatchRevision, ReasonCode = reasonCode, ReleasedAt = releasedAt }, cancellationToken: cancellationToken));
    }
}