diff --git a/experimental/store/postgres/queue.go b/experimental/store/postgres/queue.go index 775cb2f..e9c7555 100644 --- a/experimental/store/postgres/queue.go +++ b/experimental/store/postgres/queue.go @@ -12,8 +12,37 @@ import ( "github.com/deepnoodle-ai/workflow/experimental/worker" ) -// Enqueue implements worker.QueueStore. +// Enqueue implements worker.QueueStore. The insert runs in its own +// connection. When the insert must be atomic with writes to adjacent +// tables (credit ledger, idempotency keys, audit records, …), use +// EnqueueTx inside a caller-owned transaction instead. func (s *Store) Enqueue(ctx context.Context, run worker.NewRun) error { + tx, err := s.pool.Begin(ctx) + if err != nil { + return fmt.Errorf("postgres: begin enqueue tx: %w", err) + } + defer func() { _ = tx.Rollback(ctx) }() + if err := s.EnqueueTx(ctx, tx, run); err != nil { + return err + } + if err := tx.Commit(ctx); err != nil { + return fmt.Errorf("postgres: commit enqueue %s: %w", run.ID, err) + } + return nil +} + +// EnqueueTx inserts a queued run inside a caller-provided pgx +// transaction. The caller owns the tx lifecycle (Begin, Commit, +// Rollback). Use this when the run insert must be atomic with writes +// to tables outside the store's schema — e.g., debiting a credit +// ledger and creating the run in one commit. +// +// The tx must be against the same database as the Store's pool; the +// library does not verify this. +func (s *Store) EnqueueTx(ctx context.Context, tx pgx.Tx, run worker.NewRun) error { + if tx == nil { + return fmt.Errorf("postgres: EnqueueTx requires a non-nil tx") + } if run.ID == "" { return fmt.Errorf("postgres: NewRun.ID is required") } @@ -29,7 +58,7 @@ func (s *Store) Enqueue(ctx context.Context, run worker.NewRun) error { ) VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11) `, s.t("workflow_runs")) - if _, err := s.pool.Exec(ctx, query, + if _, err := tx.Exec(ctx, query, run.ID, run.Spec, string(worker.StatusQueued), nullableString(run.OrgID), nullableString(run.ProjectID), @@ -188,6 +217,39 @@ func (s *Store) Complete(ctx context.Context, claim *worker.Claim, outcome worke return nil } +// UpdateRunSpec replaces the spec on a running claim. It fences on +// (claim_id, worker_id, attempt) and status = running, and returns +// ErrLeaseLost if the fence fails — matching Heartbeat and Complete. +// +// Use this during long-running activities that mutate the run spec +// incrementally (e.g., a KB-apply loop persisting progress between +// steps) and need the update durable without waiting for the next +// checkpoint. The caller retains responsibility for producing a +// valid spec; the store does not inspect it. +func (s *Store) UpdateRunSpec(ctx context.Context, claim *worker.Claim, spec []byte) error { + if claim == nil { + return fmt.Errorf("postgres: UpdateRunSpec requires a non-nil claim") + } + query := fmt.Sprintf(` + UPDATE %s + SET spec = $1 + WHERE id = $2 + AND claimed_by = $3 + AND attempt = $4 + AND status = $5 + `, s.t("workflow_runs")) + tag, err := s.pool.Exec(ctx, query, + spec, claim.ID, claim.WorkerID, claim.Attempt, string(worker.StatusRunning), + ) + if err != nil { + return fmt.Errorf("postgres: update run spec %s: %w", claim.ID, err) + } + if tag.RowsAffected() == 0 { + return worker.ErrLeaseLost + } + return nil +} + // ReclaimStale implements worker.QueueStore. func (s *Store) ReclaimStale(ctx context.Context, staleBefore time.Time, maxAttempts int, excludeIDs []string) (int, error) { query := fmt.Sprintf(` diff --git a/experimental/store/postgres/step_progress.go b/experimental/store/postgres/step_progress.go index e1cff0f..3ee502b 100644 --- a/experimental/store/postgres/step_progress.go +++ b/experimental/store/postgres/step_progress.go @@ -57,6 +57,64 @@ func (s *Store) UpdateStepProgress(ctx context.Context, executionID string, p wo return nil } +// GetStepProgress returns every step progress row recorded for an +// execution, ordered by started_at (NULLS LAST) then step_name. One +// row per (step_name, branch_id). Returns an empty slice if no rows +// exist. Use this on the read side to render per-step status for a +// run whose identity came back from runquery.Store.GetRun, which +// intentionally does not carry step progress. +func (s *Store) GetStepProgress(ctx context.Context, executionID string) ([]workflow.StepProgress, error) { + query := fmt.Sprintf(` + SELECT step_name, branch_id, status, activity, attempt, + detail, started_at, finished_at, error + FROM %s + WHERE execution_id = $1 + ORDER BY started_at ASC NULLS LAST, step_name ASC, branch_id ASC + `, s.t("workflow_step_progress")) + + rows, err := s.pool.Query(ctx, query, executionID) + if err != nil { + return nil, fmt.Errorf("postgres: query step progress %s: %w", executionID, err) + } + defer rows.Close() + + var out []workflow.StepProgress + for rows.Next() { + var ( + p workflow.StepProgress + status string + detail []byte + startedAt *time.Time + finishedAt *time.Time + ) + if err := rows.Scan( + &p.StepName, &p.BranchID, &status, &p.ActivityName, &p.Attempt, + &detail, &startedAt, &finishedAt, &p.Error, + ); err != nil { + return nil, fmt.Errorf("postgres: scan step progress: %w", err) + } + p.Status = workflow.StepStatus(status) + if len(detail) > 0 { + pd := &workflow.ProgressDetail{} + if err := json.Unmarshal(detail, pd); err != nil { + return nil, fmt.Errorf("postgres: unmarshal progress detail for %s/%s: %w", executionID, p.StepName, err) + } + p.Detail = pd + } + if startedAt != nil { + p.StartedAt = *startedAt + } + if finishedAt != nil { + p.FinishedAt = *finishedAt + } + out = append(out, p) + } + if err := rows.Err(); err != nil { + return nil, fmt.Errorf("postgres: iterate step progress %s: %w", executionID, err) + } + return out, nil +} + func nullTime(t time.Time) any { if t.IsZero() { return nil diff --git a/experimental/store/postgres/store_test.go b/experimental/store/postgres/store_test.go index 3a45916..5b2e695 100644 --- a/experimental/store/postgres/store_test.go +++ b/experimental/store/postgres/store_test.go @@ -2,6 +2,7 @@ package postgres_test import ( "context" + "errors" "os" "testing" "time" @@ -11,6 +12,7 @@ import ( "github.com/deepnoodle-ai/workflow" "github.com/deepnoodle-ai/workflow/experimental/store/postgres" "github.com/deepnoodle-ai/workflow/experimental/worker" + "github.com/deepnoodle-ai/workflow/experimental/worker/runquery" ) // These tests require a real Postgres instance. Set WORKFLOW_PG_DSN @@ -269,3 +271,209 @@ func TestStore_StepProgressAndActivityLog(t *testing.T) { t.Fatalf("unexpected history: %+v", hist) } } + +func TestStore_GetStepProgress(t *testing.T) { + store, _ := openTestStore(t) + ctx := context.Background() + + // Empty executions return an empty slice, not an error. + got, err := store.GetStepProgress(ctx, "nobody") + if err != nil { + t.Fatalf("empty get: %v", err) + } + if len(got) != 0 { + t.Fatalf("expected empty slice, got %+v", got) + } + + t0 := time.Date(2026, 4, 12, 10, 0, 0, 0, time.UTC) + rows := []workflow.StepProgress{ + { + StepName: "second", + BranchID: "main", + Status: workflow.StepStatusCompleted, + ActivityName: "print", + Attempt: 1, + StartedAt: t0.Add(5 * time.Second), + FinishedAt: t0.Add(6 * time.Second), + }, + { + StepName: "first", + BranchID: "main", + Status: workflow.StepStatusCompleted, + ActivityName: "print", + Attempt: 1, + StartedAt: t0, + FinishedAt: t0.Add(1 * time.Second), + Detail: &workflow.ProgressDetail{ + Message: "halfway", + Data: map[string]any{"pct": float64(50)}, + }, + }, + { + StepName: "pending-step", + BranchID: "main", + Status: workflow.StepStatusPending, + ActivityName: "print", + Attempt: 0, + }, + } + for _, p := range rows { + if err := store.UpdateStepProgress(ctx, "run-sp", p); err != nil { + t.Fatalf("update: %v", err) + } + } + + // Progress for a different execution must not leak. + if err := store.UpdateStepProgress(ctx, "other-run", workflow.StepProgress{ + StepName: "x", BranchID: "main", Status: workflow.StepStatusRunning, ActivityName: "print", Attempt: 1, + StartedAt: t0, + }); err != nil { + t.Fatalf("update other: %v", err) + } + + got, err = store.GetStepProgress(ctx, "run-sp") + if err != nil { + t.Fatalf("get: %v", err) + } + if len(got) != 3 { + t.Fatalf("expected 3 rows, got %d: %+v", len(got), got) + } + // Ordered by started_at NULLS LAST, step_name, branch_id: + // "first" (t0) → "second" (t0+5s) → "pending-step" (NULL). + if got[0].StepName != "first" || got[1].StepName != "second" || got[2].StepName != "pending-step" { + t.Fatalf("unexpected order: [%s, %s, %s]", got[0].StepName, got[1].StepName, got[2].StepName) + } + // Detail round-trips. + if got[0].Detail == nil || got[0].Detail.Message != "halfway" || got[0].Detail.Data["pct"] != float64(50) { + t.Fatalf("detail round-trip mismatch: %+v", got[0].Detail) + } + // Pending row has zero started_at. + if !got[2].StartedAt.IsZero() || !got[2].FinishedAt.IsZero() { + t.Fatalf("expected zero times on pending row: %+v", got[2]) + } + // Status parses back correctly. + if got[0].Status != workflow.StepStatusCompleted || got[2].Status != workflow.StepStatusPending { + t.Fatalf("status mismatch: %+v", got) + } +} + +func TestStore_EnqueueTxRollbackAndCommit(t *testing.T) { + store, pool := openTestStore(t) + ctx := context.Background() + + // Rollback path: an EnqueueTx inside a tx that the caller + // aborts must leave no row behind. + tx, err := pool.Begin(ctx) + if err != nil { + t.Fatalf("begin: %v", err) + } + if err := store.EnqueueTx(ctx, tx, worker.NewRun{ID: "rb-run", Spec: []byte(`{}`)}); err != nil { + _ = tx.Rollback(ctx) + t.Fatalf("enqueue tx: %v", err) + } + if err := tx.Rollback(ctx); err != nil { + t.Fatalf("rollback: %v", err) + } + if _, err := store.GetRun(ctx, "", "rb-run"); !errors.Is(err, runquery.ErrRunNotFound) { + t.Fatalf("expected not found after rollback, got err=%v", err) + } + + // Commit path: EnqueueTx + another write in the same tx are + // both visible after commit. + tx2, err := pool.Begin(ctx) + if err != nil { + t.Fatalf("begin 2: %v", err) + } + if err := store.EnqueueTx(ctx, tx2, worker.NewRun{ + ID: "cm-run", Spec: []byte(`{}`), OrgID: "org-x", + }); err != nil { + _ = tx2.Rollback(ctx) + t.Fatalf("enqueue tx 2: %v", err) + } + // Write to an adjacent row inside the same tx to simulate a + // credit ledger debit: both rows must commit atomically. + if _, err := tx2.Exec(ctx, ` + INSERT INTO workflow_credit_ledger (id, org_id, run_id, amount, reason, created_at) + VALUES ('led-1', 'org-x', 'cm-run', -5, 'debit', NOW()) + `); err != nil { + _ = tx2.Rollback(ctx) + t.Fatalf("ledger insert: %v", err) + } + if err := tx2.Commit(ctx); err != nil { + t.Fatalf("commit: %v", err) + } + + run, err := store.GetRun(ctx, "org-x", "cm-run") + if err != nil { + t.Fatalf("GetRun: %v", err) + } + if run.ID != "cm-run" || run.OrgID != "org-x" { + t.Fatalf("unexpected run: %+v", run) + } + var ledgerCount int + if err := pool.QueryRow(ctx, `SELECT count(*) FROM workflow_credit_ledger WHERE run_id = $1`, "cm-run").Scan(&ledgerCount); err != nil { + t.Fatalf("ledger count: %v", err) + } + if ledgerCount != 1 { + t.Fatalf("expected 1 ledger row, got %d", ledgerCount) + } + + // Nil tx is rejected. + if err := store.EnqueueTx(ctx, nil, worker.NewRun{ID: "nil-tx", Spec: []byte(`{}`)}); err == nil { + t.Fatalf("expected error for nil tx") + } +} + +func TestStore_UpdateRunSpecFencing(t *testing.T) { + store, pool := openTestStore(t) + ctx := context.Background() + + if err := store.Enqueue(ctx, worker.NewRun{ID: "spec-run", Spec: []byte(`{"v":1}`)}); err != nil { + t.Fatalf("enqueue: %v", err) + } + claim, err := store.ClaimQueued(ctx, "w1") + if err != nil || claim == nil { + t.Fatalf("claim: %v / %+v", err, claim) + } + + // Happy path: same lease updates the spec. + newSpec := []byte(`{"v":2}`) + if err := store.UpdateRunSpec(ctx, claim, newSpec); err != nil { + t.Fatalf("update spec: %v", err) + } + var stored []byte + if err := pool.QueryRow(ctx, `SELECT spec FROM workflow_runs WHERE id = $1`, "spec-run").Scan(&stored); err != nil { + t.Fatalf("read spec: %v", err) + } + if string(stored) != `{"v":2}` { + t.Fatalf("spec not updated: %s", stored) + } + + // Wrong worker rejected. + wrong := *claim + wrong.WorkerID = "w2" + if err := store.UpdateRunSpec(ctx, &wrong, []byte(`{"v":3}`)); err != worker.ErrLeaseLost { + t.Fatalf("wrong worker: expected ErrLeaseLost, got %v", err) + } + + // Wrong attempt rejected. + badAttempt := *claim + badAttempt.Attempt = 99 + if err := store.UpdateRunSpec(ctx, &badAttempt, []byte(`{"v":3}`)); err != worker.ErrLeaseLost { + t.Fatalf("wrong attempt: expected ErrLeaseLost, got %v", err) + } + + // After completion, the run is no longer running, so updates + // must fail even with the right lease. + if err := store.Complete(ctx, claim, worker.Outcome{Status: worker.StatusCompleted}); err != nil { + t.Fatalf("complete: %v", err) + } + if err := store.UpdateRunSpec(ctx, claim, []byte(`{"v":4}`)); err != worker.ErrLeaseLost { + t.Fatalf("after complete: expected ErrLeaseLost, got %v", err) + } + + // Nil claim rejected. + if err := store.UpdateRunSpec(ctx, nil, []byte(`{"v":5}`)); err == nil { + t.Fatalf("expected error for nil claim") + } +}