knowledge-governance/internal/platform/postgres/store.go

package postgres

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"time"

	"github.com/jackc/pgx/v5"
	"github.com/jackc/pgx/v5/pgconn"

	"knowledge-governance/internal/jobs"
)

var ErrLeaseLost = errors.New("job lease lost")

type database interface {
	QueryRow(ctx context.Context, sql string, args ...any) pgx.Row
	Exec(ctx context.Context, sql string, args ...any) (pgconn.CommandTag, error)
	Query(ctx context.Context, sql string, args ...any) (pgx.Rows, error)
	Begin(ctx context.Context) (pgx.Tx, error)
}

type Store struct {
	db database
}

func NewStore(db database) *Store {
	return &Store{db: db}
}

func (s *Store) Enqueue(ctx context.Context, request jobs.EnqueueRequest, now time.Time) (*jobs.Job, error) {
	job, err := scanJob(s.db.QueryRow(
		ctx,
		enqueueJobSQL,
		request.ID,
		request.Type,
		request.Payload,
		request.Priority,
		request.AvailableAt,
		request.MaxAttempts,
		request.IdempotencyKey,
		now,
	))
	if err != nil {
		return nil, fmt.Errorf("enqueue job: %w", err)
	}
	return job, nil
}

func (s *Store) Claim(ctx context.Context, workerID string, allowedTypes []string, now time.Time, leaseDuration time.Duration) (*jobs.Job, error) {
	// Liveness heartbeat: refresh even on an empty claim so a healthy idle
	// worker still looks alive to /readyz and the monitor matrix.
	if _, err := s.db.Exec(ctx, workerHeartbeatSQL, workerID, strings.Join(allowedTypes, ","), now); err != nil {
		return nil, fmt.Errorf("record worker heartbeat: %w", err)
	}
	job, err := scanJob(s.db.QueryRow(ctx, claimJobSQL, workerID, now, now.Add(leaseDuration), allowedTypes))
	if errors.Is(err, pgx.ErrNoRows) {
		return nil, nil
	}
	if err != nil {
		return nil, fmt.Errorf("claim job: %w", err)
	}
	return job, nil
}

// ListByType lists only the requested job type, so a busy queue cannot hide the
// caller's own recent jobs behind unrelated traffic.
func (s *Store) ListByType(ctx context.Context, jobType string, limit int) ([]jobs.Job, error) {
	if strings.TrimSpace(jobType) == "" || limit < 1 || limit > 200 {
		return nil, errors.New("list jobs by type requires a type and a bounded limit")
	}
	rows, err := s.db.Query(ctx, listJobsByTypeSQL, jobType, limit)
	if err != nil {
		return nil, fmt.Errorf("list jobs by type: %w", err)
	}
	defer rows.Close()
	result := make([]jobs.Job, 0, limit)
	for rows.Next() {
		job, err := scanJob(rows)
		if err != nil {
			return nil, fmt.Errorf("scan listed job: %w", err)
		}
		result = append(result, *job)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("iterate jobs: %w", err)
	}
	return result, nil
}

func (s *Store) List(ctx context.Context, limit int) ([]jobs.Job, error) {
	rows, err := s.db.Query(ctx, listJobsSQL, limit)
	if err != nil {
		return nil, fmt.Errorf("list jobs: %w", err)
	}
	defer rows.Close()

	result := make([]jobs.Job, 0)
	for rows.Next() {
		job, err := scanJob(rows)
		if err != nil {
			return nil, fmt.Errorf("scan listed job: %w", err)
		}
		result = append(result, *job)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("iterate jobs: %w", err)
	}
	return result, nil
}

func (s *Store) Complete(ctx context.Context, id, owner string, now time.Time) error {
	// completeJobSQL is one statement: mark the job completed AND retire any
	// open dead-letter rows for it. The EXISTS projection is the lease check.
	var completed bool
	if err := s.db.QueryRow(ctx, completeJobSQL, id, owner, now).Scan(&completed); err != nil {
		return fmt.Errorf("complete job: %w", err)
	}
	if !completed {
		return ErrLeaseLost
	}
	return nil
}

func (s *Store) Heartbeat(ctx context.Context, id, owner string, now time.Time, leaseDuration time.Duration) error {
	tag, err := s.db.Exec(ctx, heartbeatJobSQL, id, owner, now, now.Add(leaseDuration))
	if err != nil {
		return fmt.Errorf("heartbeat job: %w", err)
	}
	if tag.RowsAffected() != 1 {
		return ErrLeaseLost
	}
	return nil
}

func (s *Store) Fail(ctx context.Context, id, owner, message, code string, now, retryAt time.Time) error {
	tag, err := s.db.Exec(ctx, failJobSQL, id, owner, message, now, code, retryAt)
	if err != nil {
		return fmt.Errorf("fail job: %w", err)
	}
	if tag.RowsAffected() != 1 {
		return ErrLeaseLost
	}
	return nil
}

// FailTerminal implements jobs.TerminalFailer (P6): park the job in 'failed'
// immediately with a stable error code and project the dead-letter row.
func (s *Store) FailTerminal(ctx context.Context, id, owner, message, code string, now time.Time) error {
	tag, err := s.db.Exec(ctx, failTerminalJobSQL, id, owner, message, code, now)
	if err != nil {
		return fmt.Errorf("fail terminal job: %w", err)
	}
	if tag.RowsAffected() != 1 {
		return ErrLeaseLost
	}
	return nil
}

type rowScanner interface {
	Scan(dest ...any) error
}

func scanJob(row rowScanner) (*jobs.Job, error) {
	job := &jobs.Job{}
	var state string
	var leaseOwner *string
	var lastError *string
	if err := row.Scan(
		&job.ID,
		&job.Type,
		&job.Payload,
		&state,
		&job.Priority,
		&job.AvailableAt,
		&leaseOwner,
		&job.LeaseUntil,
		&job.Attempt,
		&job.MaxAttempts,
		&lastError,
		&job.IdempotencyKey,
		&job.CreatedAt,
		&job.UpdatedAt,
	); err != nil {
		return nil, err
	}
	job.State = jobs.State(state)
	if leaseOwner != nil {
		job.LeaseOwner = *leaseOwner
	}
	if lastError != nil {
		job.LastError = *lastError
	}
	return job, nil
}