Skip to content

Commit b0e3e36

Browse files
myzieclaude
andauthored
feat(postgres): GetStepProgress, EnqueueTx, UpdateRunSpec (#41)
Three additive methods on *postgres.Store that unblock consumers migrating off hand-rolled Postgres adapters onto the experimental store. None of them touch the worker package, the runquery interface, or existing method signatures. - GetStepProgress(ctx, executionID) reads the step_progress rows back out. Without it, consumers whose runs come through runquery.Store.GetRun have no API to render per-step status in their UI. Ordered by started_at NULLS LAST, step_name, branch_id. Lives on *Store, not runquery.Store, to keep the backend-neutral read interface minimal. - EnqueueTx(ctx, tx, run) takes a caller-owned pgx.Tx so the run insert can be atomic with writes to adjacent tables (credit ledger debits, idempotency keys, audit rows) that live outside the library's schema. Enqueue is now a one-line wrapper around EnqueueTx with its own transaction. pgx.Tx is the right type here — experimental/store/postgres already exposes *pgxpool.Pool publicly, so this does not leak a new dependency. - UpdateRunSpec(ctx, claim, spec) replaces the spec on a running claim with the same (claim_id, worker_id, attempt, status=running) fencing as Heartbeat and Complete, returning ErrLeaseLost on fence failure. For long-running activities that mutate the spec incrementally (e.g. a KB-apply loop persisting progress between steps) and need it durable without waiting for the next checkpoint. Tests exercise ordering and detail round-trip for GetStepProgress, rollback+commit semantics and cross-table atomicity for EnqueueTx, and all three fence failure modes for UpdateRunSpec (wrong worker, wrong attempt, completed run). Nil-argument guards on EnqueueTx and UpdateRunSpec. Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
1 parent 46f649f commit b0e3e36

3 files changed

Lines changed: 330 additions & 2 deletions

File tree

‎experimental/store/postgres/queue.go‎

Lines changed: 64 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,8 +12,37 @@ import (
1212
"github.com/deepnoodle-ai/workflow/experimental/worker"
1313
)
1414

15-
// Enqueue implements worker.QueueStore.
15+
// Enqueue implements worker.QueueStore. The insert runs in its own
16+
// connection. When the insert must be atomic with writes to adjacent
17+
// tables (credit ledger, idempotency keys, audit records, …), use
18+
// EnqueueTx inside a caller-owned transaction instead.
1619
func (s *Store) Enqueue(ctx context.Context, run worker.NewRun) error {
20+
tx, err := s.pool.Begin(ctx)
21+
if err != nil {
22+
return fmt.Errorf("postgres: begin enqueue tx: %w", err)
23+
}
24+
defer func() { _ = tx.Rollback(ctx) }()
25+
if err := s.EnqueueTx(ctx, tx, run); err != nil {
26+
return err
27+
}
28+
if err := tx.Commit(ctx); err != nil {
29+
return fmt.Errorf("postgres: commit enqueue %s: %w", run.ID, err)
30+
}
31+
return nil
32+
}
33+
34+
// EnqueueTx inserts a queued run inside a caller-provided pgx
35+
// transaction. The caller owns the tx lifecycle (Begin, Commit,
36+
// Rollback). Use this when the run insert must be atomic with writes
37+
// to tables outside the store's schema — e.g., debiting a credit
38+
// ledger and creating the run in one commit.
39+
//
40+
// The tx must be against the same database as the Store's pool; the
41+
// library does not verify this.
42+
func (s *Store) EnqueueTx(ctx context.Context, tx pgx.Tx, run worker.NewRun) error {
43+
if tx == nil {
44+
return fmt.Errorf("postgres: EnqueueTx requires a non-nil tx")
45+
}
1746
if run.ID == "" {
1847
return fmt.Errorf("postgres: NewRun.ID is required")
1948
}
@@ -29,7 +58,7 @@ func (s *Store) Enqueue(ctx context.Context, run worker.NewRun) error {
2958
)
3059
VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9, $10, $11)
3160
`, s.t("workflow_runs"))
32-
if _, err := s.pool.Exec(ctx, query,
61+
if _, err := tx.Exec(ctx, query,
3362
run.ID, run.Spec, string(worker.StatusQueued),
3463
nullableString(run.OrgID),
3564
nullableString(run.ProjectID),
@@ -188,6 +217,39 @@ func (s *Store) Complete(ctx context.Context, claim *worker.Claim, outcome worke
188217
return nil
189218
}
190219

220+
// UpdateRunSpec replaces the spec on a running claim. It fences on
221+
// (claim_id, worker_id, attempt) and status = running, and returns
222+
// ErrLeaseLost if the fence fails — matching Heartbeat and Complete.
223+
//
224+
// Use this during long-running activities that mutate the run spec
225+
// incrementally (e.g., a KB-apply loop persisting progress between
226+
// steps) and need the update durable without waiting for the next
227+
// checkpoint. The caller retains responsibility for producing a
228+
// valid spec; the store does not inspect it.
229+
func (s *Store) UpdateRunSpec(ctx context.Context, claim *worker.Claim, spec []byte) error {
230+
if claim == nil {
231+
return fmt.Errorf("postgres: UpdateRunSpec requires a non-nil claim")
232+
}
233+
query := fmt.Sprintf(`
234+
UPDATE %s
235+
SET spec = $1
236+
WHERE id = $2
237+
AND claimed_by = $3
238+
AND attempt = $4
239+
AND status = $5
240+
`, s.t("workflow_runs"))
241+
tag, err := s.pool.Exec(ctx, query,
242+
spec, claim.ID, claim.WorkerID, claim.Attempt, string(worker.StatusRunning),
243+
)
244+
if err != nil {
245+
return fmt.Errorf("postgres: update run spec %s: %w", claim.ID, err)
246+
}
247+
if tag.RowsAffected() == 0 {
248+
return worker.ErrLeaseLost
249+
}
250+
return nil
251+
}
252+
191253
// ReclaimStale implements worker.QueueStore.
192254
func (s *Store) ReclaimStale(ctx context.Context, staleBefore time.Time, maxAttempts int, excludeIDs []string) (int, error) {
193255
query := fmt.Sprintf(`

‎experimental/store/postgres/step_progress.go‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -57,6 +57,64 @@ func (s *Store) UpdateStepProgress(ctx context.Context, executionID string, p wo
5757
return nil
5858
}
5959

60+
// GetStepProgress returns every step progress row recorded for an
61+
// execution, ordered by started_at (NULLS LAST) then step_name. One
62+
// row per (step_name, branch_id). Returns an empty slice if no rows
63+
// exist. Use this on the read side to render per-step status for a
64+
// run whose identity came back from runquery.Store.GetRun, which
65+
// intentionally does not carry step progress.
66+
func (s *Store) GetStepProgress(ctx context.Context, executionID string) ([]workflow.StepProgress, error) {
67+
query := fmt.Sprintf(`
68+
SELECT step_name, branch_id, status, activity, attempt,
69+
detail, started_at, finished_at, error
70+
FROM %s
71+
WHERE execution_id = $1
72+
ORDER BY started_at ASC NULLS LAST, step_name ASC, branch_id ASC
73+
`, s.t("workflow_step_progress"))
74+
75+
rows, err := s.pool.Query(ctx, query, executionID)
76+
if err != nil {
77+
return nil, fmt.Errorf("postgres: query step progress %s: %w", executionID, err)
78+
}
79+
defer rows.Close()
80+
81+
var out []workflow.StepProgress
82+
for rows.Next() {
83+
var (
84+
p workflow.StepProgress
85+
status string
86+
detail []byte
87+
startedAt *time.Time
88+
finishedAt *time.Time
89+
)
90+
if err := rows.Scan(
91+
&p.StepName, &p.BranchID, &status, &p.ActivityName, &p.Attempt,
92+
&detail, &startedAt, &finishedAt, &p.Error,
93+
); err != nil {
94+
return nil, fmt.Errorf("postgres: scan step progress: %w", err)
95+
}
96+
p.Status = workflow.StepStatus(status)
97+
if len(detail) > 0 {
98+
pd := &workflow.ProgressDetail{}
99+
if err := json.Unmarshal(detail, pd); err != nil {
100+
return nil, fmt.Errorf("postgres: unmarshal progress detail for %s/%s: %w", executionID, p.StepName, err)
101+
}
102+
p.Detail = pd
103+
}
104+
if startedAt != nil {
105+
p.StartedAt = *startedAt
106+
}
107+
if finishedAt != nil {
108+
p.FinishedAt = *finishedAt
109+
}
110+
out = append(out, p)
111+
}
112+
if err := rows.Err(); err != nil {
113+
return nil, fmt.Errorf("postgres: iterate step progress %s: %w", executionID, err)
114+
}
115+
return out, nil
116+
}
117+
60118
func nullTime(t time.Time) any {
61119
if t.IsZero() {
62120
return nil

‎experimental/store/postgres/store_test.go‎

Lines changed: 208 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@ package postgres_test
22

33
import (
44
"context"
5+
"errors"
56
"os"
67
"testing"
78
"time"
@@ -11,6 +12,7 @@ import (
1112
"github.com/deepnoodle-ai/workflow"
1213
"github.com/deepnoodle-ai/workflow/experimental/store/postgres"
1314
"github.com/deepnoodle-ai/workflow/experimental/worker"
15+
"github.com/deepnoodle-ai/workflow/experimental/worker/runquery"
1416
)
1517

1618
// These tests require a real Postgres instance. Set WORKFLOW_PG_DSN
@@ -269,3 +271,209 @@ func TestStore_StepProgressAndActivityLog(t *testing.T) {
269271
t.Fatalf("unexpected history: %+v", hist)
270272
}
271273
}
274+
275+
func TestStore_GetStepProgress(t *testing.T) {
276+
store, _ := openTestStore(t)
277+
ctx := context.Background()
278+
279+
// Empty executions return an empty slice, not an error.
280+
got, err := store.GetStepProgress(ctx, "nobody")
281+
if err != nil {
282+
t.Fatalf("empty get: %v", err)
283+
}
284+
if len(got) != 0 {
285+
t.Fatalf("expected empty slice, got %+v", got)
286+
}
287+
288+
t0 := time.Date(2026, 4, 12, 10, 0, 0, 0, time.UTC)
289+
rows := []workflow.StepProgress{
290+
{
291+
StepName: "second",
292+
BranchID: "main",
293+
Status: workflow.StepStatusCompleted,
294+
ActivityName: "print",
295+
Attempt: 1,
296+
StartedAt: t0.Add(5 * time.Second),
297+
FinishedAt: t0.Add(6 * time.Second),
298+
},
299+
{
300+
StepName: "first",
301+
BranchID: "main",
302+
Status: workflow.StepStatusCompleted,
303+
ActivityName: "print",
304+
Attempt: 1,
305+
StartedAt: t0,
306+
FinishedAt: t0.Add(1 * time.Second),
307+
Detail: &workflow.ProgressDetail{
308+
Message: "halfway",
309+
Data: map[string]any{"pct": float64(50)},
310+
},
311+
},
312+
{
313+
StepName: "pending-step",
314+
BranchID: "main",
315+
Status: workflow.StepStatusPending,
316+
ActivityName: "print",
317+
Attempt: 0,
318+
},
319+
}
320+
for _, p := range rows {
321+
if err := store.UpdateStepProgress(ctx, "run-sp", p); err != nil {
322+
t.Fatalf("update: %v", err)
323+
}
324+
}
325+
326+
// Progress for a different execution must not leak.
327+
if err := store.UpdateStepProgress(ctx, "other-run", workflow.StepProgress{
328+
StepName: "x", BranchID: "main", Status: workflow.StepStatusRunning, ActivityName: "print", Attempt: 1,
329+
StartedAt: t0,
330+
}); err != nil {
331+
t.Fatalf("update other: %v", err)
332+
}
333+
334+
got, err = store.GetStepProgress(ctx, "run-sp")
335+
if err != nil {
336+
t.Fatalf("get: %v", err)
337+
}
338+
if len(got) != 3 {
339+
t.Fatalf("expected 3 rows, got %d: %+v", len(got), got)
340+
}
341+
// Ordered by started_at NULLS LAST, step_name, branch_id:
342+
// "first" (t0) → "second" (t0+5s) → "pending-step" (NULL).
343+
if got[0].StepName != "first" || got[1].StepName != "second" || got[2].StepName != "pending-step" {
344+
t.Fatalf("unexpected order: [%s, %s, %s]", got[0].StepName, got[1].StepName, got[2].StepName)
345+
}
346+
// Detail round-trips.
347+
if got[0].Detail == nil || got[0].Detail.Message != "halfway" || got[0].Detail.Data["pct"] != float64(50) {
348+
t.Fatalf("detail round-trip mismatch: %+v", got[0].Detail)
349+
}
350+
// Pending row has zero started_at.
351+
if !got[2].StartedAt.IsZero() || !got[2].FinishedAt.IsZero() {
352+
t.Fatalf("expected zero times on pending row: %+v", got[2])
353+
}
354+
// Status parses back correctly.
355+
if got[0].Status != workflow.StepStatusCompleted || got[2].Status != workflow.StepStatusPending {
356+
t.Fatalf("status mismatch: %+v", got)
357+
}
358+
}
359+
360+
func TestStore_EnqueueTxRollbackAndCommit(t *testing.T) {
361+
store, pool := openTestStore(t)
362+
ctx := context.Background()
363+
364+
// Rollback path: an EnqueueTx inside a tx that the caller
365+
// aborts must leave no row behind.
366+
tx, err := pool.Begin(ctx)
367+
if err != nil {
368+
t.Fatalf("begin: %v", err)
369+
}
370+
if err := store.EnqueueTx(ctx, tx, worker.NewRun{ID: "rb-run", Spec: []byte(`{}`)}); err != nil {
371+
_ = tx.Rollback(ctx)
372+
t.Fatalf("enqueue tx: %v", err)
373+
}
374+
if err := tx.Rollback(ctx); err != nil {
375+
t.Fatalf("rollback: %v", err)
376+
}
377+
if _, err := store.GetRun(ctx, "", "rb-run"); !errors.Is(err, runquery.ErrRunNotFound) {
378+
t.Fatalf("expected not found after rollback, got err=%v", err)
379+
}
380+
381+
// Commit path: EnqueueTx + another write in the same tx are
382+
// both visible after commit.
383+
tx2, err := pool.Begin(ctx)
384+
if err != nil {
385+
t.Fatalf("begin 2: %v", err)
386+
}
387+
if err := store.EnqueueTx(ctx, tx2, worker.NewRun{
388+
ID: "cm-run", Spec: []byte(`{}`), OrgID: "org-x",
389+
}); err != nil {
390+
_ = tx2.Rollback(ctx)
391+
t.Fatalf("enqueue tx 2: %v", err)
392+
}
393+
// Write to an adjacent row inside the same tx to simulate a
394+
// credit ledger debit: both rows must commit atomically.
395+
if _, err := tx2.Exec(ctx, `
396+
INSERT INTO workflow_credit_ledger (id, org_id, run_id, amount, reason, created_at)
397+
VALUES ('led-1', 'org-x', 'cm-run', -5, 'debit', NOW())
398+
`); err != nil {
399+
_ = tx2.Rollback(ctx)
400+
t.Fatalf("ledger insert: %v", err)
401+
}
402+
if err := tx2.Commit(ctx); err != nil {
403+
t.Fatalf("commit: %v", err)
404+
}
405+
406+
run, err := store.GetRun(ctx, "org-x", "cm-run")
407+
if err != nil {
408+
t.Fatalf("GetRun: %v", err)
409+
}
410+
if run.ID != "cm-run" || run.OrgID != "org-x" {
411+
t.Fatalf("unexpected run: %+v", run)
412+
}
413+
var ledgerCount int
414+
if err := pool.QueryRow(ctx, `SELECT count(*) FROM workflow_credit_ledger WHERE run_id = $1`, "cm-run").Scan(&ledgerCount); err != nil {
415+
t.Fatalf("ledger count: %v", err)
416+
}
417+
if ledgerCount != 1 {
418+
t.Fatalf("expected 1 ledger row, got %d", ledgerCount)
419+
}
420+
421+
// Nil tx is rejected.
422+
if err := store.EnqueueTx(ctx, nil, worker.NewRun{ID: "nil-tx", Spec: []byte(`{}`)}); err == nil {
423+
t.Fatalf("expected error for nil tx")
424+
}
425+
}
426+
427+
func TestStore_UpdateRunSpecFencing(t *testing.T) {
428+
store, pool := openTestStore(t)
429+
ctx := context.Background()
430+
431+
if err := store.Enqueue(ctx, worker.NewRun{ID: "spec-run", Spec: []byte(`{"v":1}`)}); err != nil {
432+
t.Fatalf("enqueue: %v", err)
433+
}
434+
claim, err := store.ClaimQueued(ctx, "w1")
435+
if err != nil || claim == nil {
436+
t.Fatalf("claim: %v / %+v", err, claim)
437+
}
438+
439+
// Happy path: same lease updates the spec.
440+
newSpec := []byte(`{"v":2}`)
441+
if err := store.UpdateRunSpec(ctx, claim, newSpec); err != nil {
442+
t.Fatalf("update spec: %v", err)
443+
}
444+
var stored []byte
445+
if err := pool.QueryRow(ctx, `SELECT spec FROM workflow_runs WHERE id = $1`, "spec-run").Scan(&stored); err != nil {
446+
t.Fatalf("read spec: %v", err)
447+
}
448+
if string(stored) != `{"v":2}` {
449+
t.Fatalf("spec not updated: %s", stored)
450+
}
451+
452+
// Wrong worker rejected.
453+
wrong := *claim
454+
wrong.WorkerID = "w2"
455+
if err := store.UpdateRunSpec(ctx, &wrong, []byte(`{"v":3}`)); err != worker.ErrLeaseLost {
456+
t.Fatalf("wrong worker: expected ErrLeaseLost, got %v", err)
457+
}
458+
459+
// Wrong attempt rejected.
460+
badAttempt := *claim
461+
badAttempt.Attempt = 99
462+
if err := store.UpdateRunSpec(ctx, &badAttempt, []byte(`{"v":3}`)); err != worker.ErrLeaseLost {
463+
t.Fatalf("wrong attempt: expected ErrLeaseLost, got %v", err)
464+
}
465+
466+
// After completion, the run is no longer running, so updates
467+
// must fail even with the right lease.
468+
if err := store.Complete(ctx, claim, worker.Outcome{Status: worker.StatusCompleted}); err != nil {
469+
t.Fatalf("complete: %v", err)
470+
}
471+
if err := store.UpdateRunSpec(ctx, claim, []byte(`{"v":4}`)); err != worker.ErrLeaseLost {
472+
t.Fatalf("after complete: expected ErrLeaseLost, got %v", err)
473+
}
474+
475+
// Nil claim rejected.
476+
if err := store.UpdateRunSpec(ctx, nil, []byte(`{"v":5}`)); err == nil {
477+
t.Fatalf("expected error for nil claim")
478+
}
479+
}

0 commit comments

Comments
 (0)