feat(batch): failure policy, empty-batch safety, batch discovery (#55, #58, #59)

[rcbevans]

Summary

Introduces first-class batch identity via a batches table, a consecutive-failure abort policy (AbortBatchAfter), safe empty-batch detection in wait_for_batch, and batch enumeration via list_batches. Addresses three issues: batch failure abort (#55), empty-batch/finalizer ordering safety (#58), and batch discovery (#59).

Issues addressed

• Batch-level consecutive-failure abort policy #55 — Batch failure abort: when an upstream dependency is down, consecutive terminal failures should cancel remaining batch jobs instead of dispatching doomed work.
• Fan-out-then-finalize: an empty batch reports complete, and finalizer ordering is load-bearing #58 — Empty-batch / finalizer ordering: wait_for_batch on a wrong batch_id returned silently (ambiguous with genuinely-empty batches); finalizer enqueue was not atomic with batch items.
• No way to enumerate or discover batches — batch_id is recoverable only if the caller stored it #59 — Batch discovery: callers had no way to enumerate outstanding batches without storing batch_id themselves.

Implementation

New batches table and migration

src/taskq/migrations/01.00.05_01_pre_batches.sql — Creates the batches table with columns: id (PK), queue, status (active/complete/aborted, default active), expected_size, consecutive_failures, failure_threshold (NULL = no policy), finalizer_job_id, originating_actor, created_at, completed_at, metadata. Two partial indexes: batches_queue_status_idx (queue, status WHERE active) and batches_finalizer_idx (finalizer_job_id WHERE NOT NULL). The table is opt-in — only batches with a failure_policy or finalizer create a row.

Failure policy types (src/taskq/batch_policy.py)

• BatchFailurePolicy — frozen abstract base with should_abort(consecutive_failures: int) -> bool.
• AbortBatchAfter(consecutive_failures: int) — concrete policy that aborts when the consecutive failure count reaches the threshold. Validates >= 1 at construction.

New exceptions (src/taskq/exceptions.py)

• BatchAbortedError(batch_id, consecutive_failures, threshold) — raised by wait_for_batch when the batch row has status='aborted' and all jobs are terminal.
• EmptyBatchError(batch_id, expected, actual) — raised by wait_for_batch when batch_id matches zero jobs and no batches row exists, or when expect_at_least is not satisfied.

Backend protocol additions (src/taskq/backend/_protocol.py)

Three new data carriers and 10 new Backend protocol methods:

• BatchRow — read-model of a batches row.
• BatchCounts — live job-count aggregate (mirrors BatchCompletionStatus fields).
• BatchFilter — filter params for list_batches (queue, active, batch_id, limit). Deliberately distinct from JobFilter to avoid silently ignoring job-oriented fields.

Protocol methods: enqueue_batch_atomic, create_batch, increment_batch_failures, reset_batch_failures, abort_batch, complete_batch, get_batch, list_batches, count_batch_non_terminal, prune_old_batches.

PostgresBackend (src/taskq/backend/_batch_sql.py, src/taskq/backend/postgres.py)

_batch_sql.py holds all batch SQL as module-level functions taking (conn, BatchSql, ...):

• increment_batch_failures / reset_batch_failures — single CTE statement that updates the counter and counts non-terminal jobs in one round-trip, returning (new_count, threshold, remaining).
• abort_batch — UPDATE jobs SET status='cancelled' ... WHERE metadata @> $1 AND status IN ('pending','scheduled') plus UPDATE batches SET status='aborted'. Sets cancel_requested_at and cancel_phase=2 on cancelled rows to match the normal cancel path for observability.
• list_batches — one query: batches LEFT JOIN LATERAL GROUP BY status aggregation over jobs matched by metadata @>. Returns (BatchRow, BatchCounts) pairs ordered by created_at DESC.
• enqueue_batch_atomic — acquires a connection, opens a transaction, consumes the (lazy) args iterable in chunk_size groups, inserts the batches row and finalizer as the LAST statements, commits. Rolls back on any exception.
• prune_old_batches — DELETE FROM batches WHERE completed_at < cutoff AND NOT EXISTS (matching jobs).

PostgresBackend adds thin wrapper methods delegating to these functions.

InMemoryBackend (src/taskq/testing/_batch.py, src/taskq/testing/in_memory.py)

New _batch.py companion module implements all 10 batch operations as module-level functions over InMemoryBackend._batches and _jobs. abort_batch with no batch row still cancels matching jobs but does not create a row (matching the protocol contract).

Post-terminal-write hook (src/taskq/batch.py)

apply_batch_terminal_outcome(backend, job, outcome, *, loop_conn=None) — invoked after every terminal write by worker/dispatch.py (production consumer) and testing/_runner.py (in-memory runner), so abort and completion semantics are identical on both backends:

• Non-batched jobs (no metadata.batch_id): returns immediately, zero overhead.
• succeeded: resets consecutive_failures to 0; if remaining non-terminal count hits 0, marks batch complete.
• failed: increments consecutive_failures; if threshold reached, aborts (cancels remaining pending/scheduled, sets batch aborted); if not aborted and remaining hits 0, marks complete.
• cancelled/crashed: counts non-terminal jobs; if 0, marks complete.

The hook is wrapped in try/except at both call sites so a batch-policy failure never affects the job's own terminal state.

Modified enqueue_batch (src/taskq/client/_jobs.py)

• New params: failure_policy: BatchFailurePolicy | None, finalizer: EnqueueItem | None.
• When failure_policy or finalizer is set and no caller connection: delegates to Backend.enqueue_batch_atomic for single-transaction atomicity (closes the partial-batch crash window from Fan-out-then-finalize: an empty batch reports complete, and finalizer ordering is load-bearing #58).
• When a caller connection is provided: inserts jobs, then create_batch and finalizer via enqueue_with_conn, all on the caller's transaction.
• The finalizer job is not stamped with metadata.batch_id — this prevents wait_for_batch inside the finalizer from counting the finalizer itself as a pending job (infinite Snooze deadlock). Correlation is via batches.finalizer_job_id.
• BatchHandle.finalizer_handle is set; the finalizer is also appended as the last entry of job_handles for backward compat.

New enqueue_batch_streaming (src/taskq/client/_jobs.py)

Accepts any Iterable[EnqueueItem] (including generators), chunks internally at chunk_size (1–1000), validates payloads on the fly. Same transaction semantics as enqueue_batch: autonomous atomic path via enqueue_batch_atomic when extras are set and no connection; chunked enqueue_batch calls on a caller-owned connection otherwise. Empty iterable raises ValueError.

Modified wait_for_batch (src/taskq/batch.py, src/taskq/testing/_runner.py)

New params: expect_at_least: int | None, on_empty: Literal["error", "ok"], exclude_job_id: UUID | None.

Decision table (first match wins, applied when pending == 0):

1. Batch row status='aborted' → raise BatchAbortedError.
2. expect_at_least set and total < expect_at_least → raise EmptyBatchError.
3. total == 0 with a batch row → return status (row proves deliberate creation).
4. total == 0 without a batch row, on_empty="error" → raise EmptyBatchError.
5. total == 0 without a batch row, on_empty="ok" → return status (old behavior).

exclude_job_id omits a specific job from the count query (AND id <> $n, folded into the existing GROUP BY). When not set and a batch row has finalizer_job_id, the finalizer is excluded automatically.

The in-memory wait_for_batch in _runner.py mirrors the same decision table, reading backend._batches instead of querying PG.

list_batches (src/taskq/client/_jobs.py, src/taskq/batch.py)

BatchSummary dataclass (in batch.py) pairs a BatchRow with a live BatchCompletionStatus. JobsClient.list_batches(BatchFilter) delegates to Backend.list_batches and maps each (BatchRow, BatchCounts) pair to BatchSummary. Returns batches ordered by created_at DESC.

TaskQ facade (src/taskq/client/_taskq.py)

TaskQ.enqueue_batch forwards failure_policy and finalizer. New TaskQ.enqueue_batch_streaming and TaskQ.list_batches delegates.

Leader sweeps (src/taskq/worker/_leader_sweeps.py, src/taskq/worker/_leader_shared.py)

• complete_stale_batches(conn, schema) — safety net: marks active batches with zero non-terminal jobs as complete. Runs in the sweep loop.
• prune_old_batches — called in the prune loop after terminal jobs are archived, using the maximum of all per-status retention cutoffs. A batch is only deleted after all its jobs have been archived.

Re-exports (src/taskq/__init__.py)

Added: AbortBatchAfter, BatchAbortedError, BatchFailurePolicy, BatchFilter, BatchSummary, EmptyBatchError, apply_batch_terminal_outcome.

Breaking changes

• wait_for_batch on a wrong/empty batch_id now raises EmptyBatchError by default (on_empty="error"). Previously it returned BatchCompletionStatus(total=0, is_complete=True) silently. Callers that intentionally wait on possibly-empty old-style batches can pass on_empty="ok" to restore the old behavior.
• Backend protocol gains 10 new methods. Third-party backend implementations must implement these to satisfy the protocol. BACKEND_PROTOCOL_VERSION is not bumped — changes are additive and collected for the 1.0.0 release.

Test coverage

Unit (no PG required)

File	Covers
tests/test_batch_policy.py	BatchFailurePolicy / AbortBatchAfter construction, validation, should_abort logic, frozen dataclass
tests/test_batch_exceptions.py	BatchAbortedError, EmptyBatchError — isinstance, message content, carried fields
tests/test_batch_protocol.py	BatchRow, BatchCounts, BatchFilter — frozen dataclass, construction, validation
tests/test_in_memory_batch.py	All 10 InMemoryBackend batch operations: create, get, increment, reset, abort (with and without batch row), complete, count, list, prune, enqueue_batch_atomic
tests/test_batch_enqueue.py	enqueue_batch with failure_policy (creates batch row) and without (no row), finalizer not stamped with batch_id, finalizer_job_id set on batch row
tests/test_batch_abort_integration.py	End-to-end abort via run_until_drained: threshold triggers abort + cancellation, success resets counter, no-abort without policy, batch marked complete on all-terminal
tests/test_batch_streaming.py	Large/small iterables, failure_policy with streaming, empty iterable raises, invalid chunk_size raises
tests/test_batch_discovery.py	list_batches with queue/active filters, empty result, filtering by queue
tests/test_batch_prune.py	InMemoryBackend prune_old_batches (completed batch pruned after retention, active not pruned, batch with remaining jobs not pruned); complete_stale_batches integration against PG
tests/test_taskq_facade_batch.py	TaskQ facade forwards failure_policy/finalizer, enqueue_batch_streaming, list_batches
tests/test_wait_for_batch.py	expect_at_least satisfied/not-satisfied, empty batch with row (expected_size=0) OK, empty without row raises by default / ok with on_empty="ok", aborted batch raises BatchAbortedError, exclude_job_id omits caller, finalizer auto-excluded, abort property invariant (Hypothesis)

PG integration

File	Covers
tests/test_batch_pg.py	create_batch/get_batch, increment_batch_failures returns threshold + remaining, reset_batch_failures, abort_batch cancels jobs + marks batch aborted + sets cancel_requested_at/cancel_phase, complete_batch, count_batch_non_terminal, list_batches with filters, prune_old_batches
tests/test_migrations_batches.py	Batches table columns, defaults, indexes after migration

E2E (full worker + PG)

File	Covers
tests/e2e/test_batch_abort.py	10-job all-fail batch with AbortBatchAfter(3) — exactly 3 fail, 7 cancelled, batch aborted
tests/e2e/test_batch_finalizer.py	5 chunk jobs + finalizer in one enqueue_batch — finalizer snoozes via wait_for_batch until children complete, then runs
tests/e2e/test_batch_discovery.py	list_batches finds active and completed batches with live counts
tests/e2e/test_batch_empty_safety.py	Wrong batch_id raises EmptyBatchError; expect_at_least=5 with 3 jobs raises EmptyBatchError

Migration

New forward-only migration: src/taskq/migrations/01.00.05_01_pre_batches.sql. Creates the batches table and two partial indexes. No post counterpart needed. wait_for_batch catches UndefinedTableError on the batches-row fetch so job counting keeps working during the rolling-deploy window before the migration is applied.

---

Closes #55, closes #58, closes #59