pgqueue.go

package pgqueue

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

	sq "github.com/Masterminds/squirrel"
	"github.com/goware/pgkit/v2"
	"github.com/jackc/pgx/v5"
)

var (
	ErrTaskNotFound       = errors.New("pgqueue: task not found")
	ErrTaskNotRecoverable = errors.New("pgqueue: task is not in a recoverable state")
)

const tableName = "pgqueue_tasks"

// psql is a squirrel statement builder with PostgreSQL dollar placeholders.
var psql = sq.StatementBuilder.PlaceholderFormat(sq.Dollar)

// Queue provides access to the task table for enqueueing, querying, and recovery.
type Queue struct {
	tasks *pgkit.Table[Task, *Task, int64]
	log   *slog.Logger
}

// New creates a Queue backed by the given pgkit.DB.
func New(db *pgkit.DB, opts ...QueueOption) *Queue {
	cfg := queueConfig{}
	for _, o := range opts {
		o(&cfg)
	}
	q := &Queue{
		tasks: &pgkit.Table[Task, *Task, int64]{
			DB:       db,
			Name:     tableName,
			IDColumn: "id",
		},
		log: slog.Default(),
	}
	if cfg.logger != nil {
		q.log = cfg.logger
	}
	return q
}

// Get returns a task by ID.
func (q *Queue) Get(ctx context.Context, id int64) (*Task, error) {
	task, err := q.tasks.GetByID(ctx, id)
	if err != nil {
		if errors.Is(err, pgkit.ErrNoRows) {
			return nil, ErrTaskNotFound
		}
		return nil, fmt.Errorf("pgqueue: get task: %w", err)
	}
	return task, nil
}

// Find returns a task by queue name and hash.
func (q *Queue) Find(ctx context.Context, queue string, hash string) (*Task, error) {
	task, err := q.tasks.Get(ctx, sq.Eq{"queue": queue, "hash": hash}, nil)
	if err != nil {
		if errors.Is(err, pgkit.ErrNoRows) {
			return nil, ErrTaskNotFound
		}
		return nil, fmt.Errorf("pgqueue: find task: %w", err)
	}
	return task, nil
}

// Cancel hard-deletes a pending task. Returns ErrTaskNotFound if the task
// doesn't exist or is not in pending status.
func (q *Queue) Cancel(ctx context.Context, id int64) error {
	tag, err := q.tasks.Query.Exec(ctx,
		psql.Delete(q.tasks.Name).
			Where(sq.Eq{"id": id, "status": TaskStatusPending}),
	)
	if err != nil {
		return fmt.Errorf("pgqueue: cancel task: %w", err)
	}
	if tag.RowsAffected() == 0 {
		return ErrTaskNotFound
	}
	return nil
}

// Enable resets a failed or disabled task to pending. Only works on tasks
// that are not running or claimed. Resets try count and clears lease fields.
func (q *Queue) Enable(ctx context.Context, id int64) error {
	return q.enableWithRunAt(ctx, id, time.Now().UTC())
}

// Requeue resets a failed or disabled task to pending with an explicit run_at.
func (q *Queue) Requeue(ctx context.Context, id int64, runAt time.Time) error {
	return q.enableWithRunAt(ctx, id, runAt)
}

func (q *Queue) enableWithRunAt(ctx context.Context, id int64, runAt time.Time) error {
	tag, err := q.tasks.Query.Exec(ctx,
		psql.Update(q.tasks.Name).
			Set("status", TaskStatusPending).
			Set("try", 0).
			Set("run_at", runAt).
			Set("claimed_at", nil).
			Set("lease_until", nil).
			Set("claim_token", nil).
			Where(sq.And{
				sq.Eq{"id": id},
				sq.Eq{"status": []TaskStatus{TaskStatusFailed, TaskStatusDisabled}},
				sq.Eq{"claim_token": nil},
			}),
	)
	if err != nil {
		return fmt.Errorf("pgqueue: enable task: %w", err)
	}
	if tag.RowsAffected() == 0 {
		return ErrTaskNotRecoverable
	}
	return nil
}

// ReplacePayload overwrites the payload on a failed or disabled task.
// Validates that the payload is valid JSON. Does not modify hash.
func (q *Queue) ReplacePayload(ctx context.Context, id int64, payload []byte) error {
	if len(payload) > 0 && !json.Valid(payload) {
		return fmt.Errorf("pgqueue: payload is not valid JSON")
	}
	tag, err := q.tasks.Query.Exec(ctx,
		psql.Update(q.tasks.Name).
			Set("payload", payload).
			Where(sq.And{
				sq.Eq{"id": id},
				sq.Eq{"status": []TaskStatus{TaskStatusFailed, TaskStatusDisabled}},
				sq.Eq{"claim_token": nil},
			}),
	)
	if err != nil {
		return fmt.Errorf("pgqueue: replace payload: %w", err)
	}
	if tag.RowsAffected() == 0 {
		return ErrTaskNotRecoverable
	}
	return nil
}

// List returns tasks matching the optional filters with pagination.
func (q *Queue) List(ctx context.Context, queue *string, status *TaskStatus, page *pgkit.Page) ([]*Task, *pgkit.Page, error) {
	where := sq.And{}
	if queue != nil {
		where = append(where, sq.Eq{"queue": *queue})
	}
	if status != nil {
		where = append(where, sq.Eq{"status": *status})
	}
	tasks, nextPage, err := q.tasks.ListPaged(ctx, where, page)
	if err != nil {
		return nil, nil, fmt.Errorf("pgqueue: list tasks: %w", err)
	}
	return tasks, nextPage, nil
}

// Count returns the number of tasks matching the optional filters.
func (q *Queue) Count(ctx context.Context, queues []string, status *TaskStatus) (uint64, error) {
	where := sq.And{}
	if len(queues) > 0 {
		where = append(where, sq.Eq{"queue": queues})
	}
	if status != nil {
		where = append(where, sq.Eq{"status": *status})
	}
	count, err := q.tasks.Count(ctx, where)
	if err != nil {
		return 0, fmt.Errorf("pgqueue: count tasks: %w", err)
	}
	return count, nil
}

// Enqueue adds a new task to the queue. Generic free function because Go
// doesn't support generic methods on non-generic receiver types.
func Enqueue[P any](ctx context.Context, q *Queue, spec JobSpec[P], payload P, opts ...EnqueueOption) (int64, error) {
	return enqueue(ctx, q.tasks, spec.Queue, spec.HashFn, payload, opts...)
}

// EnqueueTx adds a new task within an existing transaction.
func EnqueueTx[P any](ctx context.Context, q *Queue, tx pgx.Tx, spec JobSpec[P], payload P, opts ...EnqueueOption) (int64, error) {
	return enqueue(ctx, q.tasks.WithTx(tx), spec.Queue, spec.HashFn, payload, opts...)
}

func enqueue[P any](ctx context.Context, tasks *pgkit.Table[Task, *Task, int64], queue string, hashFn func(P) *string, payload P, opts ...EnqueueOption) (int64, error) {
	cfg := enqueueConfig{}
	for _, o := range opts {
		o(&cfg)
	}

	data, err := json.Marshal(payload)
	if err != nil {
		return 0, fmt.Errorf("pgqueue: marshal payload: %w", err)
	}

	task := &Task{
		Queue:   queue,
		Status:  TaskStatusPending,
		Payload: data,
	}

	if cfg.runAt != nil {
		task.RunAt = *cfg.runAt
	}
	if cfg.status != nil {
		task.Status = *cfg.status
	}
	if hashFn != nil {
		task.Hash = hashFn(payload)
	}

	if cfg.ignoreDuplicates {
		if task.Hash == nil {
			return 0, fmt.Errorf("pgqueue: WithIgnoreDuplicates requires a non-nil dedup hash (spec.HashFn must be set and return non-nil)")
		}
		if err := tasks.Upsert(ctx, []string{"queue", "hash"}, nil, task); err != nil {
			return 0, fmt.Errorf("pgqueue: upsert task: %w", err)
		}
		// Upsert DO NOTHING doesn't return the ID. Look up by (queue, hash).
		existing, err := tasks.Get(ctx, sq.Eq{"queue": queue, "hash": *task.Hash}, nil)
		if err != nil {
			return 0, fmt.Errorf("pgqueue: lookup after upsert: %w", err)
		}
		return existing.ID, nil
	}

	if err := tasks.Insert(ctx, task); err != nil {
		return 0, fmt.Errorf("pgqueue: insert task: %w", err)
	}
	return task.ID, nil
}

// ReplacePayloadJSON marshals a typed payload and calls ReplacePayload.
// Does not modify hash — safe for tickers and non-deduped jobs.
func ReplacePayloadJSON[P any](ctx context.Context, q *Queue, id int64, payload P) error {
	data, err := json.Marshal(payload)
	if err != nil {
		return fmt.Errorf("pgqueue: marshal payload: %w", err)
	}
	return q.ReplacePayload(ctx, id, data)
}

// ReplacePayloadAndHash marshals a typed payload, recomputes hash via spec.HashFn,
// and updates both atomically. For recovering deduped jobs.
func ReplacePayloadAndHash[P any](ctx context.Context, q *Queue, id int64, spec JobSpec[P], payload P) error {
	data, err := json.Marshal(payload)
	if err != nil {
		return fmt.Errorf("pgqueue: marshal payload: %w", err)
	}
	if !json.Valid(data) {
		return fmt.Errorf("pgqueue: payload is not valid JSON")
	}

	update := psql.Update(q.tasks.Name).
		Set("payload", data).
		Where(sq.And{
			sq.Eq{"id": id},
			sq.Eq{"status": []TaskStatus{TaskStatusFailed, TaskStatusDisabled}},
			sq.Eq{"claim_token": nil},
		})

	if spec.HashFn != nil {
		update = update.Set("hash", spec.HashFn(payload))
	}

	tag, err := q.tasks.Query.Exec(ctx, update)
	if err != nil {
		return fmt.Errorf("pgqueue: replace payload and hash: %w", err)
	}
	if tag.RowsAffected() == 0 {
		return ErrTaskNotRecoverable
	}
	return nil
}