feat(migrate): let migrations opt out of the transaction wrapper (-- taskq:no-transaction)

[rcbevans]

Closes #29

Summary

apply_pending wrapped every migration in a single transaction (migrate.py:189 on main), which structurally forbade CREATE INDEX CONCURRENTLY / DROP INDEX CONCURRENTLY (Postgres rejects them inside a transaction block). Every index migration was therefore a blocking build holding a SHARE lock against writes for a full-table scan, harmless on tiny tables, fleet-stalling on jobs/job_events. This PR adds the per-migration opt-out the issue proposes.

How it works

• A migration places -- taskq:no-transaction in its leading comment block (-- line comments only, before the first SQL token). discover() parses it into a new Migration.use_transaction field (default True). The directive lives inside the SQL template, so toggling it is covered by the existing checksum-drift warning.
• Directive parsing is deliberately forgiving: prefix match and case-insensitive, so a trailing note after the token or mixed case still counts, while \b keeps no-transactional from matching. A header line that mentions taskq: without matching logs a migration-directive-unrecognized warning naming the file, so a typo cannot silently run transactional.
• apply_pending executes flagged migrations statement by statement with no wrapping transaction (Alembic's autocommit_block semantics). Splitting is required: a multi-statement string over Postgres' simple query protocol runs as one implicit transaction, so conn.execute(whole_file) would still defeat CONCURRENTLY. split_statements() handles quoted strings, quoted identifiers, line and nested block comments, and dollar-quoted bodies, and a $ inside an identifier (a$b$c) is not misread as a dollar-quote opener.
• The ledger records completion only after every statement succeeds. A failed non-transactional migration leaves earlier statements in place and is not recorded, so it re-runs on the next migrate up. Hence the documented contract that such migrations be idempotent and re-runnable (DROP INDEX CONCURRENTLY IF EXISTS + CREATE INDEX CONCURRENTLY IF NOT EXISTS, which is also the standard remedy for an INVALID index left by an interrupted build; the docstring example spells out why the DROP line is not redundant).
• The distinction is surfaced in the ledger: schema_migrations.use_transaction records how each migration ran, and taskq migrate status annotates no-transaction migrations with (no transaction). The runner owns the ledger's shape (like Rails' schema_migrations / Alembic's alembic_version) and adds the column itself, once up front for existing ledgers and lazily on the first record for fresh installs. Pre-existing rows backfill to true.
• Guardrail (strong_migrations-style): transaction-control statements in a flagged file are rejected before anything executes. That covers BEGIN/COMMIT/ROLLBACK and aliases, SAVEPOINT/RELEASE, and SET LOCAL/SET TRANSACTION, which are silent no-ops outside a transaction (an author would believe e.g. statement_timeout was disabled for a long build when it was not). A small tokenizer skips whitespace and comments, including nested block comments, so comment-separated forms like SET /* tune */ LOCAL ... cannot slip past. Plain SET/SET SESSION and CHECKPOINT stay allowed.

Failure handling: users never inspect the database by hand

A failed transactional migration rolls back on its own. A failed no-transaction migration is not recorded and is idempotent by contract, so re-running taskq migrate up (or restarting the worker) heals it. The tool now says so itself:

• taskq migrate up catches apply failures and prints a short report instead of a traceback: which migration failed, whether it rolled back or left partial effects, any INVALID indexes it found in the schema, and the single action to take. The failing migration is tagged on the exception, so --phase runs cannot misattribute the report to an earlier pending migration.
• The startup path (apply_pending_locked) does the same and aborts with one actionable line. Connection failures get the short report too. Nothing prints a traceback.

CI proves migrations are safe on a populated database

New tests/test_migrations_populated.py: a stepwise harness that seeds a realistic dataset (tens of thousands of jobs across all statuses, plus events, attempts, archives, cron schedules, rate-limit rows) and applies every bundled migration one at a time on a real Postgres container, asserting after each step that exactly that migration applied, no INVALID indexes exist, the ledger records how it ran, and the seeded data is intact. It finishes with a real-API smoke test (enqueue, dispatch, idempotency dedupe, cross-scope idempotency keys) and a migrate up no-op run on the populated database. Future migrations join the harness automatically, so CI fails naming the offender if any migration ever damages a populated database.

Prior art consulted (as the issue requests)

• Rails disable_ddl_transaction!: per-migration macro; the migrator skips the wrapping transaction for that migration only. Our use_transaction field mirrors Rails' internal use_transaction? check.
• Django AddIndexConcurrently / atomic = False: concurrent index ops require a non-atomic migration; docs defer to the PostgreSQL INVALID-index caveats. Same contract here: author-owned idempotency.
• Alembic autocommit_block(): "each statement is committed immediately" in autocommit mode, with transaction_per_migration recommended for the rest, exactly TaskQ's existing per-migration-transaction design plus the new per-statement escape hatch. (Why not Alembic itself: it requires SQLAlchemy as a dependency, has no pre/post phase guard for rolling deploys, and fleet-safe startup application under an advisory lock would be hand-rolled around it anyway.)

Verification (real output)

$ uv run ruff check .
All checks passed!

$ uv run ruff format --check .
470 files already formatted

$ uv run pyright src/taskq tests
0 errors (1 pre-existing hvac stub warning in tests/test_vault.py)

$ uv run pytest -n 2 --cov=taskq --cov-fail-under=90 -m "not redis" -q
4788 passed, 6 skipped, coverage 93.07%

(One flake in the full run, tests/test_cancel_hook_integration.py::test_force_cancel, passed on immediate rerun and is unrelated to this PR. The 6 skips are pre-existing extras/PG-only skips.)

Test coverage added

• Unit: directive parsing (trailing notes, mixed case, no-transactional rejection, near-miss warning), statement splitter (dollar sign inside identifiers, U&'...' regression pins, dollar quotes after punctuation), transaction-control guard (new keywords, comment-separated and nested-comment bypass forms, allowlist pins), failure-diagnosis renderer (exact line lists for every variant), exception tagging, ApplyFailureDiagnosis invariant.
• Integration, real Postgres 18: the original directive/CONCURRENTLY suite; end-to-end parse through the real discover(); the populated-database stepwise harness; CLI failure report through CliRunner; apply_pending_locked failure self-diagnosis; list_invalid_indexes round-trip on staged debris.

Follow-ups (deliberately deferred)

• No bundled migration is retrofitted (guarded by test_bundled_migrations_are_all_transactional). The directive is there for the index work on hot tables to adopt when ready.
• The job_attempts to job_attempts_archive SELECT ja.* issue from the closing note was handled on PR feat: scope idempotency_key uniqueness via idempotency_scope #27's branch, left alone here.
• Pre-existing nits noted by reviewers, out of scope here: list_applied fetches a checksum column it never uses; apply_pending could document that the advisory lock is opt-in (apply_pending_locked) for concurrent runners.

[XBeg9]

comments inline

[XBeg9]

fullmatch + no IGNORECASE means this only fires if the line is character-exact. Poked at it:

-- taskq:no-transaction                  -> works
-- taskq:no-transaction  needed for CIC  -> silently transactional
-- TaskQ:No-Transaction                  -> silently transactional
/* taskq:no-transaction */               -> silently transactional

Sticking a note after the directive explaining why you opted out is probably the first thing anyone writes, and it kills it.

With CIC you find out straight away since PG refuses. What worries me more is someone going no-transaction for lock duration instead — that just quietly runs wrapped forever and nothing ever tells them.

match + IGNORECASE would cover most of it, but honestly the bigger win is warning when a header line has taskq: in it and doesn't parse.

[rcbevans]

Fixed. The directive now prefix-matches and is case-insensitive, bounded by \b so no-transactional still does not count, and a leading-comment line that mentions taskq: without parsing logs a migration-directive-unrecognized warning naming the file. The trailing-note case works, and the silent transactional failure mode you flagged now surfaces in logs. Unit tests plus a new end-to-end test through discover() pin it.

[XBeg9]

Verified. Trailing note and mixed case both opt out now, no-transactional correctly doesn't, and the near-miss warning fires on it. Ran all four against HEAD.

[XBeg9]

These all sail through the guard:

SAVEPOINT sp1;
RELEASE SAVEPOINT sp1;
SET TRANSACTION ISOLATION LEVEL SERIALIZABLE;
SET LOCAL work_mem = '1GB';
CHECKPOINT;

Most are harmless no-ops. SET LOCAL isn't — SET LOCAL statement_timeout = 0; before a long CIC is a totally normal thing to write, and outside a transaction it does nothing except log a warning nobody reads. You think you've disabled the timeout, you haven't.

(I did try to break this with nested block comments and CR line endings first. Both get caught properly, the regex is fine there.)

[rcbevans]

Extended. The guard now rejects SAVEPOINT and RELEASE up front, and SET LOCAL / SET TRANSACTION for exactly the reason you gave: outside a transaction they are silent no-ops, so an author would believe statement_timeout was disabled for a long build when it was not. Plain SET / SET SESSION and CHECKPOINT stay allowed: session-scoped, and they behave the same either way. The regex is gone in favor of a small tokenizer that skips comments in both positions, including nested block comments, so forms like SET /* tune */ LOCAL cannot slip past the new rule either.

[XBeg9]

Verified — SET LOCAL, SAVEPOINT, RELEASE, SET TRANSACTION all rejected, CHECKPOINT/SET/SET SESSION still allowed. The tokenizer also catches SET /* x */ LOCAL, which the old regex couldn't, so this is better than what I asked for.

One gap left: SET CONSTRAINTS ALL DEFERRED is allowed. Outside a transaction it warns and does nothing — the exact rationale in your own docstring for rejecting SET LOCAL. Probably wants "constraints" in the second-word tuple.

[XBeg9]

$ inside an identifier gets read as a dollar-quote opener:

split_statements("SELECT a$b$c; SELECT 2;")  ->  ['SELECT a$b$c; SELECT 2;']

One chunk, second statement swallowed. And a$b$c is a real identifier as far as PG is concerned, I checked — SELECT a$b$c FROM (SELECT 1 AS a$b$c) q runs fine. U&'...' does the same thing.

Not scary in practice: the merged chunk hits PG as multi-statement, so anything CONCURRENTLY in there blows up loudly. Two boring statements would just quietly end up sharing a transaction.

A tag can't follow an identifier char, so checking the previous char before entering the $ branch should sort it.

[rcbevans]

Fixed with the prev-char check you suggested: a dollar-quote tag can no longer follow an identifier character, so a$b$c stays one identifier and the following statement splits normally. One correction: U&'...' already splits correctly, since backslash only matters before a quote in the splitter and PG rejects a backslash-escape inside U& strings regardless. Regression pins for it added anyway next to the a$b$c test.

[XBeg9]

a$b$c fixed, verified — 2 chunks, and real $$/$tag$ quoting still parses as one.

You're right about U&'...' and I was wrong. My test input was malformed — PG rejects SELECT U&'\''; with unterminated quoted string, so returning one chunk was correct behaviour, not a bug. Well-formed U& cases all split fine. Sorry for the noise.

[XBeg9]

That DROP INDEX CONCURRENTLY IF EXISTS is doing real work and looks like it isn't. On 17:

after interrupted CIC:             indisvalid = false
re-run CIC IF NOT EXISTS:          NOTICE: ... already exists, skipping
                                   indisvalid = false    <- still busted
DROP CONCURRENTLY + recreate:      indisvalid = true

IF NOT EXISTS sees the invalid index and skips, so without the DROP the retry never actually fixes anything — you're left with an index the planner won't use but writers still maintain.

Someone's going to delete that line as redundant one day.

[rcbevans]

Done. The docstring example now carries a NOT-redundant comment explaining the INVALID-index debris case, mirrored in the upgrading guide, so nobody deletes the line as dead code.

[XBeg9]

Good — and putting it in CONTRIBUTING.md as well as the docstring is the right call, that's where someone writing a new migration will actually look.

[XBeg9]

This is still the old runbook — stop workers, restore from backup, pin the version. None of that is what you do when a no-transaction migration dies halfway. Then it's: see what actually landed, check for an INVALID index, fix forward, re-run.

Bit ironic that the section about things going wrong is the one that didn't get updated for the mode where things don't roll back.

docs/architecture.md:385 still claims migrations are always wrapped and CONCURRENTLY is impossible, too.

[rcbevans]

Rewritten, and the manual inspection is gone entirely: the runbook no longer asks anyone to query catalogs by hand. migrate up now diagnoses itself on failure and prints which migration failed, whether it rolled back or left partial effects, any INVALID indexes found, and the single action to take. The section says that in a few lines, and architecture.md:385 is updated to match.

[XBeg9]

The rewrite is right, and building diagnose_apply_failure rather than just fixing the prose is a bigger swing than I expected.

The subsystem has some problems though — left them as separate comments rather than burying them here. Short version: it can hang holding the advisory lock, and it prints a couple of things that aren't true.

[XBeg9]

These construct Migration(use_transaction=False) by hand, so they test the runner but never the parsing that decides the flag.

Made _uses_transaction() always return True — feature completely off, everything forced transactional — and integration went 975/975 green. Three unit assertions on discover() were the only thing that noticed.

Which is how the fullmatch thing above would get through: happy path covered, near-misses not, and integration blind to the whole question.

A single integration test that writes a real .sql with the header and goes through discover() would cover both.

Flip side, since I was mutating things anyway — split_statements, _reject_transaction_control and _ensure_ledger_use_transaction_column all fail loudly when broken (18, 9, and basically the entire integration suite). Those are fine.

[rcbevans]

Added test_discover_directive_parsing_applies_end_to_end: writes a real .sql carrying the directive with a trailing note (so it covers the fullmatch issue above too), runs it through the real discover() into apply_pending against Postgres, and asserts the concurrent index lands valid and the ledger records use_transaction=false. Repeating your _uses_transaction always-True mutation now fails loudly with ActiveSQLTransactionError.

[XBeg9]

Closed. test_discover_directive_parsing_applies_end_to_end goes through real discover() over a real file, which is what was missing. Confirmed the mixed-case pin goes red if the regex is reverted.