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(ctx context.Context, id uuid.UUID) error
	FailWithRetry(ctx context.Context, id uuid.UUID, scheduledAt time.Time, lastErr string) error
	DeadLetter(ctx context.Context, id uuid.UUID, 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.
func (r *PgRepository) Complete(ctx context.Context, id uuid.UUID) error {
	if err := r.q.CompleteJob(ctx, toPgUUID(id)); err != nil {
		return errs.Wrap(err, errs.CodeInternal, "complete job")
	}
	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, scheduledAt time.Time, lastErr string) error {
	if err := r.q.FailJobWithRetry(ctx, jobssqlc.FailJobWithRetryParams{
		ID:          toPgUUID(id),
		ScheduledAt: pgtype.Timestamptz{Time: scheduledAt, Valid: true},
		LastError:   &lastErr,
	}); err != nil {
		return errs.Wrap(err, errs.CodeInternal, "fail job with retry")
	}
	return nil
}

// DeadLetter terminally marks a job as dead and stores the last error.
func (r *PgRepository) DeadLetter(ctx context.Context, id uuid.UUID, lastErr string) error {
	if err := r.q.DeadLetterJob(ctx, jobssqlc.DeadLetterJobParams{
		ID:        toPgUUID(id),
		LastError: &lastErr,
	}); err != nil {
		return errs.Wrap(err, errs.CodeInternal, "dead-letter job")
	}
	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
}