agentic-coding-workflow/internal/jobs/repository.go

package jobs

import (
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"time"

	"github.com/google/uuid"
	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgtype"
	"github.com/jackc/pgx/v5/pgxpool"

	"github.com/yan1h/agent-coding-workflow/internal/infra/errs"
	jobssqlc "github.com/yan1h/agent-coding-workflow/internal/jobs/sqlc"
)

// Repository is the persistence contract for jobs. It is implemented by
// PgRepository for production and may be faked for unit tests of higher-level
// components (worker, scheduler).
type Repository interface {
	Enqueue(ctx context.Context, typ JobType, payload json.RawMessage, scheduledAt time.Time, maxAttempts int32) (*Job, error)
	Lease(ctx context.Context, leasedBy string) (*Job, error)
	// Complete, FailWithRetry and DeadLetter take the leasing worker id and apply
	// a fencing predicate (id + leased_by + status='running'); if the lease has
	// been lost (e.g. the reaper re-leased the job) they return ErrLeaseLost so
	// the caller drops the write instead of clobbering the new run's state.
	Complete(ctx context.Context, id uuid.UUID, leasedBy string) error
	FailWithRetry(ctx context.Context, id uuid.UUID, leasedBy string, scheduledAt time.Time, lastErr string) error
	DeadLetter(ctx context.Context, id uuid.UUID, leasedBy string, lastErr string) error
	Get(ctx context.Context, id uuid.UUID) (*Job, error)
	ReapStuckRunningJobs(ctx context.Context, leasedBefore time.Time) ([]uuid.UUID, error)
	PurgeCompleted(ctx context.Context, completedBefore time.Time) (int64, error)
}

// PgRepository wraps sqlc-generated queries with a pgxpool.Pool. It is the only
// production implementation of Repository.
type PgRepository struct {
	pool *pgxpool.Pool
	q    *jobssqlc.Queries
}

// NewPostgresRepository builds a PgRepository on the given pool.
func NewPostgresRepository(pool *pgxpool.Pool) *PgRepository {
	return &PgRepository{pool: pool, q: jobssqlc.New(pool)}
}

var _ Repository = (*PgRepository)(nil)

// ===== type-conversion helpers =====

func toPgUUID(u uuid.UUID) pgtype.UUID { return pgtype.UUID{Bytes: u, Valid: true} }

func fromPgUUID(p pgtype.UUID) uuid.UUID {
	if !p.Valid {
		return uuid.Nil
	}
	return uuid.UUID(p.Bytes)
}

// ===== Repository methods =====

// Enqueue persists a new pending job. Caller-supplied payload is rejected if it
// exceeds MaxPayloadBytes so a runaway producer can't bloat the table.
func (r *PgRepository) Enqueue(ctx context.Context, typ JobType, payload json.RawMessage, scheduledAt time.Time, maxAttempts int32) (*Job, error) {
	if len(payload) > MaxPayloadBytes {
		return nil, errs.New(errs.CodeJobPayloadTooLarge, fmt.Sprintf("payload exceeds %d bytes", MaxPayloadBytes))
	}
	row, err := r.q.EnqueueJob(ctx, jobssqlc.EnqueueJobParams{
		ID:          toPgUUID(uuid.New()),
		Type:        string(typ),
		Payload:     payload,
		ScheduledAt: pgtype.Timestamptz{Time: scheduledAt, Valid: true},
		MaxAttempts: maxAttempts,
	})
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "enqueue job")
	}
	return rowToJob(row), nil
}

// Lease atomically transitions the oldest due pending job to running and
// returns it. Returns (nil, nil) when no job is ready — callers treat that as
// "queue empty" and back off.
func (r *PgRepository) Lease(ctx context.Context, leasedBy string) (*Job, error) {
	row, err := r.q.LeaseNextJob(ctx, &leasedBy)
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, nil
		}
		return nil, errs.Wrap(err, errs.CodeInternal, "lease job")
	}
	return rowToJob(row), nil
}

// Complete marks a job as completed and clears its lease/error fields. The
// update is fenced on the leasing worker id and status='running'; if no row
// matches (lease lost to the reaper) it returns ErrLeaseLost as a no-op signal.
func (r *PgRepository) Complete(ctx context.Context, id uuid.UUID, leasedBy string) error {
	n, err := r.q.CompleteJob(ctx, jobssqlc.CompleteJobParams{
		ID:       toPgUUID(id),
		LeasedBy: &leasedBy,
	})
	if err != nil {
		return errs.Wrap(err, errs.CodeInternal, "complete job")
	}
	if n == 0 {
		return ErrLeaseLost
	}
	return nil
}

// FailWithRetry returns a job to pending with a new scheduled_at and stores
// the last error. Attempts is preserved (it was incremented on lease) so the
// worker can compare against MaxAttempts to decide whether to retry or
// dead-letter on the next failure.
func (r *PgRepository) FailWithRetry(ctx context.Context, id uuid.UUID, leasedBy string, scheduledAt time.Time, lastErr string) error {
	n, err := r.q.FailJobWithRetry(ctx, jobssqlc.FailJobWithRetryParams{
		ID:          toPgUUID(id),
		ScheduledAt: pgtype.Timestamptz{Time: scheduledAt, Valid: true},
		LastError:   &lastErr,
		LeasedBy:    &leasedBy,
	})
	if err != nil {
		return errs.Wrap(err, errs.CodeInternal, "fail job with retry")
	}
	if n == 0 {
		return ErrLeaseLost
	}
	return nil
}

// DeadLetter terminally marks a job as dead and stores the last error. Fenced
// on the leasing worker id and status='running'; returns ErrLeaseLost on a
// 0-row match so a stale worker does not overwrite a re-leased run.
func (r *PgRepository) DeadLetter(ctx context.Context, id uuid.UUID, leasedBy string, lastErr string) error {
	n, err := r.q.DeadLetterJob(ctx, jobssqlc.DeadLetterJobParams{
		ID:        toPgUUID(id),
		LastError: &lastErr,
		LeasedBy:  &leasedBy,
	})
	if err != nil {
		return errs.Wrap(err, errs.CodeInternal, "dead-letter job")
	}
	if n == 0 {
		return ErrLeaseLost
	}
	return nil
}

// Get fetches a job by id. Returns CodeJobNotFound when no row exists.
func (r *PgRepository) Get(ctx context.Context, id uuid.UUID) (*Job, error) {
	row, err := r.q.GetJob(ctx, toPgUUID(id))
	if err != nil {
		if errors.Is(err, pgx.ErrNoRows) {
			return nil, errs.New(errs.CodeJobNotFound, "job not found")
		}
		return nil, errs.Wrap(err, errs.CodeInternal, "get job")
	}
	return rowToJob(row), nil
}

// ReapStuckRunningJobs resets jobs whose lease started before leasedBefore
// back to pending so another worker can pick them up. Returns the ids of all
// reaped jobs for logging/metrics.
func (r *PgRepository) ReapStuckRunningJobs(ctx context.Context, leasedBefore time.Time) ([]uuid.UUID, error) {
	rows, err := r.q.ReapStuckRunningJobs(ctx, pgtype.Timestamptz{Time: leasedBefore, Valid: true})
	if err != nil {
		return nil, errs.Wrap(err, errs.CodeInternal, "reap stuck running jobs")
	}
	out := make([]uuid.UUID, 0, len(rows))
	for _, p := range rows {
		out = append(out, fromPgUUID(p))
	}
	return out, nil
}

// PurgeCompleted deletes completed and dead jobs whose completed_at is older
// than completedBefore, returning the row count for metrics.
func (r *PgRepository) PurgeCompleted(ctx context.Context, completedBefore time.Time) (int64, error) {
	rows, err := r.q.PurgeCompletedJobs(ctx, pgtype.Timestamptz{Time: completedBefore, Valid: true})
	if err != nil {
		return 0, errs.Wrap(err, errs.CodeInternal, "purge completed jobs")
	}
	return rows, nil
}

// rowToJob converts a sqlc-generated row into the domain Job. LeasedBy and
// LastError are *string in the sqlc layer (emit_pointers_for_null_types), so
// we copy the value to avoid aliasing the row's pointer; LeasedAt and
// CompletedAt are pgtype.Timestamptz, so we check Valid.
func rowToJob(r jobssqlc.Job) *Job {
	out := &Job{
		ID:          fromPgUUID(r.ID),
		Type:        JobType(r.Type),
		Payload:     r.Payload,
		Status:      JobStatus(r.Status),
		ScheduledAt: r.ScheduledAt.Time,
		Attempts:    r.Attempts,
		MaxAttempts: r.MaxAttempts,
		CreatedAt:   r.CreatedAt.Time,
		UpdatedAt:   r.UpdatedAt.Time,
	}
	if r.LeasedAt.Valid {
		t := r.LeasedAt.Time
		out.LeasedAt = &t
	}
	if r.LeasedBy != nil {
		s := *r.LeasedBy
		out.LeasedBy = &s
	}
	if r.LastError != nil {
		s := *r.LastError
		out.LastError = &s
	}
	if r.CompletedAt.Valid {
		t := r.CompletedAt.Time
		out.CompletedAt = &t
	}
	return out
}