diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 3cd9721..2ce2ea5 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -196,6 +196,113 @@ uv run ruff format --check . uv run ruff format --check . ``` +## Authoring Database Migrations + +TaskQ's schema evolves through SQL migration files bundled with the package and applied by the runner in `src/taskq/migrate.py`. This section is for TaskQ developers adding or changing those files. End users never author migrations — they only apply them — so operator-facing guidance (upgrading, failure recovery) lives in [docs/guides/upgrading.md](docs/guides/upgrading.md) and is out of scope here. + +### File Naming and Discovery + +Migration files live in `src/taskq/migrations/` and follow one naming convention: + +``` +{ver}_{nn}_{pre|post}_{description}.sql # e.g. 01.00.03_01_pre_idempotency_scope.sql +``` + +- `{ver}` is three two-digit groups joined by dots (`01.00.03`), `{nn}` is a two-digit sequence, and `{description}` is lowercase `[a-z0-9_]+`. A filename that does not match this pattern fails discovery with a `ValueError`. +- `discover()` finds every `*.sql` file in the package automatically — there is no registry or index to update — and returns them sorted by version, with `pre` before `post` at the same version. The zero-padding is what makes the lexicographic sort chronological, so keep the two-digit groups. +- Migration versions are numbered independently of the package version: releases v0.1.0–v0.2.2 all shipped the same two migrations (`01.00.00_01` and `01.00.01_01`). Pick the next free `{ver}_{nn}`; do not try to mirror the package's semver. + +### The Default Transaction Wrapper + +By default each migration file runs inside its own transaction, so a failure rolls the whole file back. Keep migrations transactional unless Postgres makes that impossible. + +### Non-Transactional Migrations (`-- taskq:no-transaction`) + +Postgres forbids some statements inside a transaction block — `CREATE INDEX CONCURRENTLY`, `DROP INDEX CONCURRENTLY`, several `ALTER TYPE`/`VACUUM` forms. A migration that needs one of these opts out of the wrapper with a header directive: + +```sql +-- taskq:no-transaction +``` + +The usual motivation is lock duration: a plain `CREATE INDEX` holds a lock that blocks readers and writers for the whole build, which is not acceptable on hot tables such as `jobs` or `job_events`. The `CONCURRENTLY` forms avoid that lock but must run outside a transaction. + +Directive parsing rules (implemented by `_uses_transaction` in `src/taskq/migrate.py`): + +- **Leading comment block only.** Only blank lines and `--` line comments before the first SQL token are scanned; a directive later in the file is ignored, so a stray mention cannot silently flip a migration's semantics. +- **Prefix match.** A trailing note after the token is fine and encouraged: `-- taskq:no-transaction — CIC cannot run inside a transaction`. +- **Case-insensitive**, like SQL itself. +- **Token-bounded.** A `\b` word boundary stops prefix drift, so `-- taskq:no-transactional` does NOT opt out. +- **Typos surface.** A `taskq:` header line that doesn't match the directive (a typo, or a lookalike such as `no-transactional`) logs a `migration-directive-unrecognized` warning naming the file, instead of silently running the migration transactional — the worst outcome for a lock-duration-motivated migration. +- Files are decoded as `utf-8-sig`, so a BOM from a Windows editor cannot silently disable the directive. + +A non-transactional file is executed statement by statement (a multi-statement string would run as one implicit transaction and defeat the opt-out). Two rules keep that safe: + +1. **The migration must be idempotent and re-runnable.** Nothing rolls back: a mid-file failure leaves earlier statements in place, the ledger records the migration only after every statement succeeds, and the migration is re-executed on the next `migrate up`. Write every statement to tolerate being re-run (`IF NOT EXISTS`, guarded inserts, ...). +2. **Drop interrupted-build debris first.** An interrupted `CREATE INDEX CONCURRENTLY` leaves an INVALID index behind, and `IF NOT EXISTS` alone would then silently skip rebuilding it. Pair the statements: + + ```sql + -- taskq:no-transaction + -- NOT redundant with IF NOT EXISTS below: an interrupted CREATE INDEX + -- CONCURRENTLY leaves an INVALID index that IF NOT EXISTS alone would + -- silently skip rebuilding, so drop the debris first. + DROP INDEX CONCURRENTLY IF EXISTS "{schema}".jobs_foo_idx; + CREATE INDEX CONCURRENTLY IF NOT EXISTS jobs_foo_idx ON "{schema}".jobs (foo); + ``` + +Transaction-control statements are rejected in non-transactional files: `BEGIN`, `COMMIT`, `ROLLBACK` (plus their Postgres synonyms `END`, `ABORT`, `START`), `SAVEPOINT`, `RELEASE`, and `SET LOCAL` / `SET TRANSACTION`. The transaction-control group would silently re-open a transaction — defeating `CONCURRENTLY` and, on failure, poisoning the caller's connection. `SET LOCAL`/`SET TRANSACTION` are rejected for the opposite reason: outside a transaction they are silent no-ops, so you would believe e.g. `statement_timeout` was disabled for a long build when it was not. The guard rejects the whole file before any statement executes. Plain `SET` / `SET SESSION` and `CHECKPOINT` are deliberately allowed: they are session-scoped (or transaction-agnostic) and behave identically either way. + +### The `{schema}` Placeholder + +Always qualify objects with the literal `"{schema}"` token — the runner substitutes the configured schema name at apply time (after validating it as an identifier), which is what lets multiple TaskQ instances share one Postgres cluster. Never hardcode a schema name. Substitution uses `str.format`, so a literal curly brace in the SQL is written doubled (`{{` / `}}`). + +### Pre and Post Phases for Rolling Deploys + +Destructive changes are split across a deploy boundary: + +- **`pre`** adds structures that BOTH old and new code tolerate (e.g. a new index alongside the old one). Safe to apply before or during the code rollout. +- **`post`** removes the old structures after every worker in the fleet has been upgraded. + +The runner enforces the ordering: a post-phase migration is refused until its same-version pre-phase counterpart is applied (or applies earlier in the same run). See `01.00.03_01_pre_idempotency_scope.sql` for the canonical example, including the "PHASE OBLIGATIONS" header that spells out the deployment sequence. + +### Never Edit a Released Migration + +The ledger (`{schema}.schema_migrations`) records each applied migration under its `{ver}_{nn}:{phase}` key together with a SHA-256 checksum of the rendered SQL. Editing a file after it shipped makes the checksum drift, and the runner logs a `migration-checksum-drift` warning on every subsequent apply. Treat released migrations as frozen — releases v0.1.0 through v0.2.2 shipped only `01.00.00_01` and `01.00.01_01`, and those files have not changed since. Fix forward with a new migration instead. + +### File Header Convention + +Every migration file opens with a `--` header comment block, modeled on the existing files: + +1. One short paragraph stating what the migration does and why. +2. The forward-only reminder: "Forward-only; there is no down migration. To revert, restore from backup." +3. The substitution note: 'The literal "{schema}" token is substituted at apply time by the migration runner.' + +Migrations with operational impact add named ops-note sections under banner comments — see the `-- ── Maintenance-window caveat ... ──` block in `01.00.02_01_pre_job_events_outbox.sql`. State the lock impact plainly: which lock the statement takes, on which table, how long it can be held (build time is proportional to row count), what it blocks (readers/writers, hot paths), and what operators of large deployments should do instead (e.g. build the index manually with `CONCURRENTLY` during a maintenance window). Phase-coupled migrations additionally document the deployment sequence and the failure modes of applying out of order (see "PHASE OBLIGATIONS" in `01.00.03_01_pre_idempotency_scope.sql`). + +### What CI Proves for You + +- **tests/test_migrations.py** applies the full bundled set to a real PostgreSQL and pins that a second `apply_pending` is a no-op. +- **tests/test_migrations_populated.py** applies every bundled migration one step at a time onto seeded data. New migrations join this harness automatically — unknown keys get the generic per-step invariants — and its seeder intersects live columns from `information_schema`, so a future NOT NULL-without-default column fails there as the alarm. Review `_MIGRATION_SPECIFIC_CHECKS` when adding a migration whose populated-DB effect deserves a sharper assertion. +- **`test_bundled_migrations_are_all_transactional`** (tests/test_migrations_unit.py) pins that no bundled migration carries the no-transaction directive yet. It becomes an allowlist once PRs #25/#27 land the first bundled no-transaction migration. +- **Directive-parsing, transaction-control-guard, and statement-splitter unit tests** (tests/test_migrations_unit.py) pin every rule quoted above. +- **`test_discover_directive_parsing_applies_end_to_end`** (tests/test_migrate_no_transaction.py) exercises the real `discover()` parse against a real database, directive-with-trailing-note included. +- **`test_migrate_up_cli_reports_failed_no_transaction_migration`** (tests/test_migrate_no_transaction.py) proves the `migrate up` failure report end to end: what failed, what state it left the schema in, and the one action to take. +- **`test_interrupted_concurrent_build_remedy_drop_and_rebuild`** (tests/test_migrate_no_transaction.py) stages a real interrupted `CREATE INDEX CONCURRENTLY` and proves the drop-and-rebuild remedy replaces the INVALID index with a valid one. + +### Local Loop + +```bash +uv run pytest tests/test_migrations.py tests/test_migrations_populated.py tests/test_migrations_unit.py tests/test_migrate_no_transaction.py -q +``` + +The integration files need Docker, like the rest of the integration suite. + +### Non-Goals + +- **No downgrade machinery.** There is no `down` operation and none is planned; reverting means restoring from a database backup. +- **No auto-retry in the runner.** Re-running `migrate up` (or restarting the worker) IS the heal — migrations are idempotent by contract, so a failed apply is fixed forward, not retried by the framework. +- **No `migrate status` health probing.** `migrate status` lists applied and pending migrations; it deliberately does not validate or repair schema state. +- **End users never author migrations.** This section covers the bundled files only; there is no user-defined migration hook. + ## Pull Request Process ### Before Submitting diff --git a/docs/architecture.md b/docs/architecture.md index b3522e5..2a1af82 100644 --- a/docs/architecture.md +++ b/docs/architecture.md @@ -391,12 +391,18 @@ Crashes are rare so the added wakeup cost is low, but the channel's meaning has broadened. The index added by migration `01.00.02_01_pre_job_events_outbox.sql` uses -`CREATE INDEX` (not `CONCURRENTLY`, since `migrate.py` applies each migration -inside a transaction and Postgres forbids `CONCURRENTLY` there) and takes an -exclusive lock on `job_events` for the duration of the build, which can stall -writes to that heavily-written table — operators with a large/populated table -should run the equivalent `CREATE INDEX CONCURRENTLY` manually during a -maintenance window (see the migration file for details). +`CREATE INDEX` (not `CONCURRENTLY`) and takes an exclusive lock on +`job_events` for the duration of the build, which can stall writes to that +heavily-written table. The migration runner supports a per-migration opt-out +from its default transaction wrapper — the `-- taskq:no-transaction` header +directive, which unlocks `CONCURRENTLY` forms (see +[Non-transactional migrations](guides/upgrading.md#non-transactional-migrations)) — +but bundled migrations deliberately remain transactional (pinned by +`test_bundled_migrations_are_all_transactional`), so operators with a +large/populated table should still run the equivalent `CREATE INDEX +CONCURRENTLY` manually during a maintenance window (see the migration file +for details). Future index migrations on hot tables can adopt the directive +instead. --- @@ -961,8 +967,10 @@ These invariants must remain true across all changes. documentation; the `_single_threaded()` guard is a no-op. 7. **Migration files are append-only** — never modify an applied migration. - The migration runner stores a checksum of each applied file in - `schema_migrations` and rejects re-runs with a checksum mismatch. + The migration runner stores a SHA-256 checksum of each applied file's + rendered SQL in `schema_migrations` and logs a `migration-checksum-drift` + warning when an applied file no longer matches, so tampering surfaces in + logs (drift is warned on, not rejected — applied migrations never re-run). 8. **`BACKEND_PROTOCOL_VERSION` is checked at import time** — both `PostgresBackend` and `InMemoryBackend` assert the version constant at module diff --git a/docs/guides/cli.md b/docs/guides/cli.md index 441d389..f0b6569 100644 --- a/docs/guides/cli.md +++ b/docs/guides/cli.md @@ -148,6 +148,8 @@ taskq migrate up [OPTIONS] The command is idempotent: each migration is recorded in `{schema}.schema_migrations` and is skipped on subsequent runs. Running `taskq migrate up` with no options applies all pending migrations. +On failure (exit code 1) the command never prints a traceback. When a migration fails during apply, the command diagnoses itself: it names the failed migration, reports the state the schema was left in — a clean rollback for transactional migrations, or for `-- taskq:no-transaction` migrations the statements that remain applied plus any INVALID indexes found — and prints the single action to take. When the database connection itself cannot be established, there is nothing to diagnose, so it prints a short report naming the connection error and the same re-run action. + **Example: apply all pending:** ```shell diff --git a/docs/guides/upgrading.md b/docs/guides/upgrading.md index ec4d8b9..2a7edf5 100644 --- a/docs/guides/upgrading.md +++ b/docs/guides/upgrading.md @@ -54,10 +54,85 @@ This is a deliberate tradeoff, not a missing feature: `{schema}.schema_migrations` are skipped. See [cli.md](cli.md#taskq-migrate-up) for the full option reference (`--phase`, `--target`, `--max-steps`). +## Non-transactional migrations + +By default every migration file runs inside its own transaction, so a failure +rolls the whole file back. PostgreSQL forbids some statements inside a +transaction block — notably `CREATE INDEX CONCURRENTLY` and +`DROP INDEX CONCURRENTLY`, the only forms that build or drop an index without +blocking writes on the table. On hot tables (`jobs`, `job_events`) a plain +`CREATE INDEX` takes a `SHARE` lock that blocks `INSERT`/`UPDATE`/`DELETE` +for the duration of a full-table scan and stalls the worker fleet, so index +migrations on those tables should use the concurrent forms. + +A migration opts out of the transaction wrapper with a header directive in its +leading comment block (`--` line comments only, before the first SQL token): + +```sql +-- taskq:no-transaction +-- NOT redundant with IF NOT EXISTS below: an interrupted CREATE INDEX +-- CONCURRENTLY leaves an INVALID index that IF NOT EXISTS alone would +-- silently skip rebuilding, so drop the debris first. +DROP INDEX CONCURRENTLY IF EXISTS "{schema}".jobs_queue_idx; +CREATE INDEX CONCURRENTLY IF NOT EXISTS jobs_queue_idx ON "{schema}".jobs (queue); +``` + +The runner then executes the file **statement by statement, each in its own +implicit transaction** (the same semantics as Alembic's `autocommit_block` or +Rails' `disable_ddl_transaction!`). This changes the failure contract, so +three rules apply: + +- **The migration must be idempotent and re-runnable.** Nothing rolls back: if + the third statement fails, the first two stay applied. The ledger records + the migration only after *every* statement succeeds, so the next + `migrate up` re-executes the whole file — every statement must tolerate + being re-run (`IF NOT EXISTS`, guarded inserts, etc.). +- **An interrupted `CREATE INDEX CONCURRENTLY` leaves an `INVALID` index + behind.** The standard remedy is drop-and-rebuild, written into the + migration itself as shown above: the `DROP INDEX CONCURRENTLY IF EXISTS` + line removes debris from an interrupted attempt before rebuilding. A plain + `CREATE INDEX CONCURRENTLY IF NOT EXISTS` alone would silently skip the + rebuild while the invalid index keeps its name. You never have to find + these by hand: when a run fails, `taskq migrate up` lists any INVALID + indexes in its failure report. +- **No transaction-control statements.** `BEGIN`/`COMMIT`/`ROLLBACK` (and + aliases) are rejected before anything executes — they would silently + re-open a transaction, defeating the directive. The statement splitter + assumes the server default `standard_conforming_strings=on`. + +Operators can see the distinction two ways: `taskq migrate status` annotates +non-transactional migrations with `(no transaction)`, and the +`{schema}.schema_migrations` ledger records how each migration ran in its +`use_transaction` column (`false` = ran outside a transaction). The runner +adds that column when recording the next migration, so deployments upgraded +from older TaskQ versions need no manual step; rows applied before the column +existed read `true`. + ## If a migration goes wrong -- Stop workers pointed at the affected schema to avoid further writes. -- Restore the database from the pre-migration backup. -- Pin `taskq-py` back to the previous version until the issue is resolved, - since the previous version's code may not be compatible with the new - schema. +`taskq migrate up` diagnoses its own failures: it tells you which migration +failed, what state the schema is in, and the one action to take. You never +need to inspect catalog state by hand. + +### Transactional migration (the default) + +The whole file rolled back automatically, so the schema is exactly as it was +before the attempt. Fix the cause of the error, then re-run `taskq migrate +up`. + +### Non-transactional migration (`-- taskq:no-transaction`) + +Nothing rolls back: statements before the failure remain applied, and the +migration is **not** recorded. Re-run `taskq migrate up` — the migration is +idempotent, and the command's failure report lists any INVALID indexes the +interrupted attempt left behind; the drop-and-rebuild already written into +the migration cleans them up on the re-run. Only pin `taskq-py` back to the +previous version if the migration SQL itself is wrong and you need time to +ship a correction. + +### A migration applied successfully but broke older workers + +Restoring from backup is for this scenario: the migration itself succeeded, +but not-yet-upgraded workers cannot run against the new schema. Stop the +workers pointed at the affected schema, restore the pre-migration backup, +and pin `taskq-py` back until every worker is upgraded. diff --git a/pyproject.toml b/pyproject.toml index db397dc..086a680 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -238,6 +238,9 @@ ignore = [ "tests/test_idempotency_scope_migrations.py" = ["S608"] # Same rationale — schema name validated against _IDENT_RE; all # user-supplied values use $N parameter binding. +"tests/test_migrations_populated.py" = ["S608"] +# Same rationale — schema name validated against _IDENT_RE; all +# user-supplied values use $N parameter binding. # SIM117: conn.transaction() depends on conn from pool.acquire(); # async-with nesting cannot be flattened. "tests/test_singleton.py" = ["S608", "SIM117"] diff --git a/src/taskq/cli.py b/src/taskq/cli.py index 865db7e..7278fa8 100644 --- a/src/taskq/cli.py +++ b/src/taskq/cli.py @@ -303,7 +303,8 @@ async def _status(settings: TaskQSettings) -> None: typer.echo(f"applied: {len(applied)}") for migration in migrate_mod.discover(): marker = "✔" if migration.key in applied else " " - typer.echo(f" [{marker}] {migration.filename}") + suffix = "" if migration.use_transaction else " (no transaction)" + typer.echo(f" [{marker}] {migration.filename}{suffix}") async def _up( @@ -313,8 +314,12 @@ async def _up( target: str | None, max_steps: int | None, ) -> None: - conn = await asyncpg.connect(str(settings.pg_dsn)) + # conn mirrors apply_pending_locked's conn-or-None pattern: connect + # failures land in the same guarded region as apply failures, so both + # get the report — and the close is skipped when no conn was acquired. + conn: asyncpg.Connection | None = None try: + conn = await asyncpg.connect(str(settings.pg_dsn)) applied = await migrate_mod.apply_pending( conn, schema=settings.schema_name, @@ -322,10 +327,18 @@ async def _up( target=target, max_steps=max_steps, ) + except Exception as exc: + # Why diagnose here: both apply paths leave the connection reusable + # (a transactional failure rolls back; the no-transaction path never + # opened one), so the CLI reports what failed and the schema state + # on the same conn instead of escaping a raw traceback. + await _report_up_failure(conn, settings.schema_name, exc) + raise typer.Exit(code=1) from None finally: - # Why bounded: same dead-PG wedge risk as _status above (#38 - # follow-up); terminate-on-timeout, never raises. - await close_conn_bounded(conn, "migrate-up", CLOSE_TIMEOUT_SECS) + if conn is not None: + # Why bounded: same dead-PG wedge risk as _status above (#38 + # follow-up); terminate-on-timeout, never raises. + await close_conn_bounded(conn, "migrate-up", CLOSE_TIMEOUT_SECS) if not applied: typer.echo("no pending migrations") return @@ -650,6 +663,36 @@ async def _actor_config_diff( _print_actor_diff(name, registry.get(name), stored_by_actor.get(name)) +async def _report_up_failure(conn: asyncpg.Connection | None, schema: str, exc: Exception) -> None: + """Print a self-diagnosing ``migrate up`` failure report to stderr. + + TaskQ users must never inspect catalog state by hand, so this reports — + gathered on the still-open connection — what failed, what state the + schema is in (INVALID indexes included), and the single action to take. + The diagnosis lives in :mod:`taskq.migrate` (shared with the + worker/startup path). + + The report must NEVER mask the original error: when the conn was never + acquired (connect itself failed) or the diagnosis itself raises, the + fallback is the generic two-line report — the original error's headline + and the fix-and-re-run action. + """ + diagnosis: migrate_mod.ApplyFailureDiagnosis | None = None + if conn is not None: + with contextlib.suppress(Exception): + diagnosis = await migrate_mod.diagnose_apply_failure(conn, schema, exc) + if diagnosis is None: + diagnosis = migrate_mod.ApplyFailureDiagnosis( + headline=migrate_mod._exception_headline(exc), # pyright: ignore[reportPrivateUsage] # Why: the headline rule (first line, else type name) must match diagnose_apply_failure exactly; sharing the helper keeps the two from drifting. + failed_filename=None, + use_transaction=None, + invalid_indexes=(), + schema=schema, + ) + for line in migrate_mod.render_apply_failure_lines(diagnosis): + typer.echo(line, err=True) + + async def _health_request(settings: WorkerSettings, path: str) -> int: try: reader, writer = await asyncio.wait_for( diff --git a/src/taskq/migrate.py b/src/taskq/migrate.py index ecf6d4e..ffcc3dc 100644 --- a/src/taskq/migrate.py +++ b/src/taskq/migrate.py @@ -9,6 +9,45 @@ recording a SHA-256 checksum of the rendered SQL after each successful apply. There is no ``down`` operation. To revert, restore from a database backup. + +Non-transactional migrations +---------------------------- + +By default each migration file runs inside its own transaction, so a failure +rolls the whole file back. Postgres forbids some statements inside a +transaction block — ``CREATE INDEX CONCURRENTLY``, ``DROP INDEX +CONCURRENTLY``, several ``ALTER TYPE``/``VACUUM`` forms — which makes them +inexpressible under the default wrapper. A migration can opt out by placing +the header directive ``-- taskq:no-transaction`` in its leading comment +block (``--`` line comments only, before the first SQL token); the runner +then executes the file statement by statement with no wrapping transaction +(Alembic's autocommit-block semantics). Two rules keep that safe: + +* The migration **must be idempotent and re-runnable** — nothing rolls back, + so a mid-file failure leaves earlier statements in place and the migration + is re-executed on the next run. The ledger records completion only after + every statement succeeds. +* An interrupted ``CREATE INDEX CONCURRENTLY`` leaves an ``INVALID`` index + behind; the standard remedy is drop-and-rebuild, written into the + migration itself:: + + -- taskq:no-transaction + -- NOT redundant with IF NOT EXISTS below: an interrupted CREATE INDEX + -- CONCURRENTLY leaves an INVALID index that IF NOT EXISTS alone would + -- silently skip rebuilding, so drop the debris first. + DROP INDEX CONCURRENTLY IF EXISTS "{schema}".jobs_foo_idx; + CREATE INDEX CONCURRENTLY IF NOT EXISTS jobs_foo_idx ON "{schema}".jobs (foo); + +Transaction-control statements (``BEGIN``/``COMMIT``/``ROLLBACK``/...) are +rejected in non-transactional files — they would silently re-open a +transaction and, on failure, poison the caller's connection. Statement +splitting assumes ``standard_conforming_strings=on`` (the Postgres default +since 9.1). + +The ledger column ``schema_migrations.use_transaction`` records how each +migration ran (``false`` = outside a transaction) so operators can see which +migrations were/are safe to run online; the runner adds the column on first +use, so pre-upgrade ledgers need no dedicated migration. """ import asyncio @@ -29,13 +68,18 @@ ) __all__ = [ + "ApplyFailureDiagnosis", "Migration", "Phase", "apply_pending", "apply_pending_locked", + "diagnose_apply_failure", "discover", "list_applied", + "list_invalid_indexes", "render", + "render_apply_failure_lines", + "split_statements", ] logger = structlog.get_logger("taskq.migrate") @@ -46,6 +90,47 @@ r"^(?P\d{2}\.\d{2}\.\d{2})_(?P\d{2})_(?Ppre|post)_(?P[a-z0-9_]+)\.sql$" ) +# Prefix match (not fullmatch): a trailing note after the token is the common +# real-world form ("-- taskq:no-transaction — CIC cannot run inside a +# transaction"). Case-insensitive like SQL itself. \b stops prefix drift, so +# "-- taskq:no-transactional" does NOT match — but see the near-miss warning +# in _uses_transaction, which keeps such typos from failing silently. +_NO_TRANSACTION_DIRECTIVE_RE = re.compile(r"--\s*taskq:no-transaction\b", re.IGNORECASE) + +_DOLLAR_TAG_RE = re.compile( + r"\$[^\W\d][\w]*\$|\$\$" +) # tags follow identifier rules (Unicode-aware, no digits first) + + +def _uses_transaction(sql_template: str, filename: str) -> bool: + """Parse the ``-- taskq:no-transaction`` header directive. + + The directive is only honored in the file's *leading comment block* — + the blank lines and ``--`` comments before the first SQL token — so a + stray mention later in the file cannot silently flip a migration's + semantics. + + A ``--`` line that mentions ``taskq:`` without matching the directive + (a typo, or drift like ``taskq:no-transactional``) is nearly always an + attempted opt-out that will silently run transactional — the worst + outcome for a lock-duration-motivated migration. Warn so the author + notices; the warning names the file because discovery parses many. + """ + for line in sql_template.splitlines(): + stripped = line.strip() + if stripped == "" or stripped.startswith("--"): + if _NO_TRANSACTION_DIRECTIVE_RE.match(stripped): + return False + if "taskq:" in stripped.lower(): + logger.warning( + "migration-directive-unrecognized", + filename=filename, + line=stripped, + ) + continue + break + return True + @dataclass(frozen=True, slots=True) class Migration: @@ -59,6 +144,12 @@ class Migration: filename: str sql_template: str + use_transaction: bool = True + """When False (``-- taskq:no-transaction`` header directive), apply the + file statement by statement with no wrapping transaction. Postgres + requires this for ``CREATE INDEX CONCURRENTLY`` and friends — but nothing + rolls back on failure, so such migrations must be idempotent.""" + @property def key(self) -> str: """Identity stored in ``schema_migrations.version``: ``{version}:{phase}``.""" @@ -83,13 +174,17 @@ def discover() -> list[Migration]: raise ValueError(f"migration filename does not match convention: {entry.name!r}") version = f"{match.group('ver')}_{match.group('seq')}" phase: Phase = match.group("phase") # type: ignore[assignment] # Why: regex group "phase" is constrained to "pre|post" by _NAME_RE; re.match guarantees the value matches the Literal["pre", "post"] alias but str cannot be narrowed to it statically. + sql_template = entry.read_text( + encoding="utf-8-sig" + ) # utf-8-sig: a BOM (Windows editors) must not silently disable the header directive found.append( Migration( version=version, phase=phase, description=match.group("desc"), filename=entry.name, - sql_template=entry.read_text(encoding="utf-8"), + sql_template=sql_template, + use_transaction=_uses_transaction(sql_template, entry.name), ) ) found.sort(key=lambda m: (m.version, 0 if m.phase == "pre" else 1)) @@ -107,6 +202,256 @@ def render(template: str, schema: str) -> str: return template.format(schema=schema) +def split_statements(sql: str) -> list[str]: + """Split a SQL script into individual statements, without their ``;``. + + Non-transactional migrations are executed statement by statement: a + multi-statement string sent through Postgres' simple query protocol runs + as ONE implicit transaction, which would defeat the point (``CREATE + INDEX CONCURRENTLY`` would still be "inside a transaction block"). + + Splitting understands single-quoted strings (including ``E'...'`` + backslash escapes and ``''`` doubling), ``"..."``-quoted identifiers, + ``--`` line comments, nested ``/* ... */`` block comments, and + dollar-quoted bodies (``$$...$$`` / ``$tag$...$tag$``). Leading comments + stay attached to the statement that follows them; comment-only chunks + are dropped. Unterminated constructs yield one trailing chunk, leaving + the syntax error to Postgres — same as executing the file whole. + """ + statements: list[str] = [] + buf: list[str] = [] + has_content = False # any non-comment, non-whitespace char in the chunk + i = 0 + n = len(sql) + state = "normal" + backslash_escapes = False # inside E'...' strings only + block_depth = 0 + dollar_tag = "" + + def flush() -> None: + nonlocal buf, has_content + chunk = "".join(buf).strip() + if has_content: + statements.append(chunk) + buf = [] + has_content = False + + while i < n: + ch = sql[i] + nxt = sql[i + 1] if i + 1 < n else "" + + if state == "normal": + if ch == "'": + # E'...' (E directly before the quote, not part of a longer + # identifier) uses backslash escapes; plain '...' does not + # (standard_conforming_strings=on). + backslash_escapes = buf[-1:] in (["e"], ["E"]) and ( + len(buf) < 2 or not (buf[-2].isalnum() or buf[-2] in "_$") + ) + state = "squote" + has_content = True + buf.append(ch) + i += 1 + elif ch == '"': + state = "dquote" + has_content = True + buf.append(ch) + i += 1 + elif ch == "-" and nxt == "-": + state = "line_comment" + buf.append(ch) + buf.append(nxt) + i += 2 + elif ch == "/" and nxt == "*": + state = "block_comment" + block_depth = 1 + buf.append(ch) + buf.append(nxt) + i += 2 + elif ( + ch == "$" + # A dollar-quote tag cannot immediately follow an identifier + # char — ``a$b$c`` is a legal identifier, not a quoted body + # (same rule as the E'...' detection above). buf holds + # everything since the last flush; empty buf (right after a + # ``;``) means no previous char, so the branch stays allowed. + # Why: isalnum() approximates PG's identifier-char rule (any + # char >= 0x80 counts as an identifier char there) — close + # enough for a splitter, and consistent with the E'...' gate. + and not (buf and (buf[-1].isalnum() or buf[-1] in "_$")) + and (m := _DOLLAR_TAG_RE.match(sql, i)) is not None + ): + state = "dollar" + dollar_tag = m.group(0) + has_content = True + buf.append(dollar_tag) + i = m.end() + elif ch == ";": + flush() + i += 1 + else: + if not ch.isspace(): + has_content = True + buf.append(ch) + i += 1 + elif state == "squote": + buf.append(ch) + if backslash_escapes and ch == "\\" and i + 1 < n: + buf.append(sql[i + 1]) + i += 2 + elif ch == "'": + if nxt == "'": # '' escape inside the string + buf.append(nxt) + i += 2 + else: + state = "normal" + i += 1 + else: + i += 1 + elif state == "dquote": + buf.append(ch) + if ch == '"': + if nxt == '"': # "" escape inside the identifier + buf.append(nxt) + i += 2 + else: + state = "normal" + i += 1 + else: + i += 1 + elif state == "line_comment": + buf.append(ch) + if ch == "\n" or ch == "\r": # Postgres ends -- comments at CR too (CR-only files) + state = "normal" + i += 1 + elif state == "block_comment": + if ch == "/" and nxt == "*": + block_depth += 1 + buf.append(ch) + buf.append(nxt) + i += 2 + elif ch == "*" and nxt == "/": + block_depth -= 1 + buf.append(ch) + buf.append(nxt) + i += 2 + if block_depth == 0: + state = "normal" + else: + buf.append(ch) + i += 1 + else: # dollar-quoted body: verbatim until the matching closing tag + if ch == "$" and sql.startswith(dollar_tag, i): + buf.append(dollar_tag) + i += len(dollar_tag) + state = "normal" + else: + buf.append(ch) + i += 1 + + flush() + return statements + + +def _skip_sql_trivia(sql: str, i: int) -> int: + """Advance past whitespace and comments (``--`` line, nested ``/* */``). + + Comments are valid trivia here because Postgres treats them as + whitespace when scanning keywords — ``SET /* x */ LOCAL`` is the same + statement to the server as ``SET LOCAL``, and block comments NEST, which + no single regex can skip. The guard must therefore skip exactly what the + server skips, or comment-wrapped forms slide past as silent no-ops. + """ + n = len(sql) + while i < n: + if sql[i].isspace(): + i += 1 + elif sql.startswith("--", i): + # Postgres ends -- comments at CR too (CR-only files). + end = i + 2 + while end < n and sql[end] not in "\r\n": + end += 1 + i = end + elif sql.startswith("/*", i): + depth = 1 + i += 2 + while i < n and depth > 0: + if sql.startswith("/*", i): + depth += 1 + i += 2 + elif sql.startswith("*/", i): + depth -= 1 + i += 2 + else: + i += 1 + # An unterminated block comment consumes the rest of the + # statement; the word read after it finds nothing and the guard + # defers to Postgres' own syntax error — same as split_statements. + else: + break + return i + + +def _read_sql_word(sql: str, i: int) -> tuple[str, int]: + """Read one keyword-shaped word (identifier chars) starting at ``i``, + lowercased, plus the index just past it.""" + start = i + n = len(sql) + while i < n and (sql[i].isalnum() or sql[i] in "_$"): + i += 1 + return sql[start:i].lower(), i + + +_TXN_CONTROL_WORDS = frozenset( + {"begin", "commit", "rollback", "end", "abort", "start", "savepoint", "release"} +) + + +def _transaction_control_word(statement: str) -> str | None: + """Uppercase transaction-control keyword the statement opens with + (``'SET LOCAL'`` / ``'SET TRANSACTION'`` for the two-word forms), or + ``None`` when the statement is allowed. + + Beyond BEGIN/COMMIT and friends this covers SAVEPOINT/RELEASE (only + valid inside a transaction) and SET LOCAL / SET TRANSACTION (silent + no-ops outside one — the author believes a setting applied when it did + not). Plain SET / SET SESSION stays allowed: it is session-scoped and + behaves identically either way. + """ + i = _skip_sql_trivia(statement, 0) + word, i = _read_sql_word(statement, i) + if word in _TXN_CONTROL_WORDS: + return word.upper() + if word == "set": + second, _ = _read_sql_word(statement, _skip_sql_trivia(statement, i)) + if second in ("local", "transaction"): + return f"SET {second.upper()}" + return None + + +def _reject_transaction_control(migration: Migration, statements: list[str]) -> None: + """Forbid transaction-control statements in a non-transactional migration. + + ``BEGIN``/``COMMIT`` and friends would silently re-open an explicit + transaction (defeating ``CREATE INDEX CONCURRENTLY``) and, on failure, + leave the caller's connection in an aborted transaction. ``SET LOCAL`` / + ``SET TRANSACTION`` are rejected for the opposite reason: outside a + transaction they are SILENT no-ops (server WARNING only), so an author + could believe e.g. ``statement_timeout`` was disabled for a long index + build when it was not. ``SAVEPOINT``/``RELEASE`` would fail loudly at + execution time, but the guard's value is rejecting the whole file before + any statement executes — a rejected file applies nothing. + """ + for statement in statements: + word = _transaction_control_word(statement) + if word is not None: + raise ValueError( + f"migration {migration.filename!r} is marked no-transaction but contains " + f"transaction-control statement {word!r}; remove it — " + "the runner manages transactions" + ) + + async def list_applied(conn: asyncpg.Connection, schema: str) -> set[str]: """Return ``{version}:{phase}`` keys recorded in ``schema_migrations``. @@ -133,6 +478,31 @@ async def list_applied(conn: asyncpg.Connection, schema: str) -> set[str]: return applied_keys +async def list_invalid_indexes(conn: asyncpg.Connection, schema: str) -> list[str]: + """Return the names of INVALID indexes in ``schema``, sorted by name. + + An interrupted ``CREATE INDEX CONCURRENTLY`` leaves an INVALID index + behind: the query planner ignores it, but writers still maintain it, so + it is pure overhead (and blocks the re-run's ``IF NOT EXISTS``). The CLI + surfaces these in its failure report so users never have to query the + catalogs by hand. + """ + if not _IDENT_RE.match(schema): + raise ValueError(f"invalid schema name {schema!r}") + rows = await conn.fetch( + """ + SELECT c.relname + FROM pg_class c + JOIN pg_index i ON i.indexrelid = c.oid + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = $1 AND NOT i.indisvalid + ORDER BY c.relname + """, + schema, + ) + return [r["relname"] for r in rows] + + async def apply_pending( conn: asyncpg.Connection, *, @@ -144,7 +514,13 @@ async def apply_pending( """Apply pending migrations. Each migration runs in its own transaction so a failure in one file does - not leave a half-applied schema. + not leave a half-applied schema — unless the file carries the + ``-- taskq:no-transaction`` header directive (:attr:`Migration.use_transaction`), + in which case it is executed statement by statement with no wrapping + transaction. That unlocks ``CREATE INDEX CONCURRENTLY`` and friends, but + nothing rolls back on failure: prior statements stay applied, the + migration is NOT recorded in the ledger, and it will be re-executed on + the next run — so non-transactional migrations must be idempotent. :param phase: restrict to ``pre`` or ``post`` migrations only. :param target: stop after applying this version (inclusive). @@ -224,22 +600,250 @@ async def apply_pending( ) eligible_keys.add(m.key) + # The ledger must be able to record how each migration runs. An existing + # ledger is upgraded once, up front, outside any migration transaction + # (pre-upgrade ledgers lack the column); a fresh install has no ledger + # until the initial migration creates one mid-loop, so the ensure runs + # lazily on the first record instead. + ledger_ready = False + if exists and pending: + await _ensure_ledger_use_transaction_column(conn, schema) + ledger_ready = True + applied_now: list[Migration] = [] for migration in effective: - async with conn.transaction(): - await conn.execute(migration.render(schema)) - await conn.execute( - f'INSERT INTO "{schema}".schema_migrations (version, checksum) VALUES ($1, $2)', - migration.key, - migration.checksum(schema), - ) + try: + if migration.use_transaction: + async with conn.transaction(): + await conn.execute(migration.render(schema)) + if not ledger_ready: + await _ensure_ledger_use_transaction_column(conn, schema) + ledger_ready = True + await _record_applied(conn, schema, migration) + else: + # No wrapping transaction: each statement commits on its own, so + # the ledger is written only after every statement succeeds. A + # failure leaves earlier statements in place and nothing recorded; + # idempotency makes the re-run safe (see module docstring). + logger.warning( + "migration-no-transaction", + key=migration.key, + filename=migration.filename, + ) + statements = split_statements(migration.render(schema)) + _reject_transaction_control(migration, statements) + for statement in statements: + await conn.execute(statement) + if not ledger_ready: + await _ensure_ledger_use_transaction_column(conn, schema) + ledger_ready = True + await _record_applied(conn, schema, migration) + except Exception as exc: + # Why tag and re-raise: diagnose_apply_failure's first-unrecorded + # heuristic misattributes the failure under --phase (an + # earlier-version pending migration of the OTHER phase sorts + # first in discover() order). The exception object itself is the + # reliable channel to the diagnosis — tagging, not wrapping, so + # the caller-visible exception type is unchanged. + exc.__dict__["taskq_failed_migration"] = migration + raise applied_now.append(migration) return applied_now +async def _ensure_ledger_use_transaction_column(conn: asyncpg.Connection, schema: str) -> None: + """Add the ledger's ``use_transaction`` column if it is missing. + + The ledger is runner bookkeeping (like Rails' ``schema_migrations`` or + Alembic's ``alembic_version``), so the runner owns its shape and upgrades + it in place — no dedicated migration file is needed for ledgers created + by older TaskQ versions. Pre-existing rows backfill to ``true``: every + migration applied before this column existed ran inside a transaction. + """ + await conn.execute( + f'ALTER TABLE "{schema}".schema_migrations ' + "ADD COLUMN IF NOT EXISTS use_transaction boolean NOT NULL DEFAULT true" + ) + + +async def _record_applied(conn: asyncpg.Connection, schema: str, migration: Migration) -> None: + """Record a successfully applied migration in ``schema_migrations``.""" + await conn.execute( + f'INSERT INTO "{schema}".schema_migrations (version, checksum, use_transaction) ' + "VALUES ($1, $2, $3)", + migration.key, + migration.checksum(schema), + migration.use_transaction, + ) + + _MIGRATION_LOCK_KEY: int = 1_234_567 +# ── apply-failure self-diagnosis ────────────────────────────────────────────── + +# Why this lives here and not in the CLI: both failure surfaces (``migrate +# up`` and worker/UI startup via ``apply_pending_locked``) need the same +# report, and the gathering logic reads ``discover``/``list_applied``/ +# ``list_invalid_indexes`` from this module — a helper module would +# circular-import them. + +# Why a separate action line for startup: a worker/startup apply failure is +# retried by restarting the process (migrations are idempotent and the +# runner self-heals on retry), not by re-running the CLI — the CLI guidance +# would send operators down the wrong path. Escalate to the CLI only when +# the restart loop does not heal. +_STARTUP_ACTION_LINE = ( + "Action: restart is safe — migrations are idempotent and self-heal on retry; " + "if the failure repeats, run `taskq migrate up` and report the output." +) + + +@dataclass(frozen=True, slots=True) +class ApplyFailureDiagnosis: + """Self-diagnosis of a failed migration apply, gathered on the still-open + connection by :func:`diagnose_apply_failure` and rendered by + :func:`render_apply_failure_lines`.""" + + headline: str + """First line of the original error (asyncpg messages can be multiline).""" + + failed_filename: str | None + """Filename of the migration that failed — taken from the exception's + ``taskq_failed_migration`` tag when :func:`apply_pending` attached one, + else the first unrecorded migration in ``discover()`` order (a heuristic + that misattributes under ``--phase``); ``None`` when it could not be + determined.""" + + use_transaction: bool | None + """``Migration.use_transaction`` of the failed migration; ``None`` when + the failed migration could not be determined.""" + + invalid_indexes: tuple[str, ...] + """INVALID indexes in the schema — debris of an interrupted ``CREATE + INDEX CONCURRENTLY``; only gathered for no-transaction failures.""" + + # Why carried here: the INVALID-index line names the schema and the + # renderer must stay pure (no re-querying), so the diagnosis owns every + # value the report needs. + schema: str + """Schema the report names in the INVALID-index line.""" + + def __post_init__(self) -> None: + # Why: the renderer branches on use_transaction whenever + # failed_filename is set, so the pair must stay consistent — a + # filename with use_transaction=None would silently render the + # no-transaction wording for a failure whose nature is unknown. + if self.failed_filename is not None and self.use_transaction is None: + raise ValueError( + "ApplyFailureDiagnosis: use_transaction is required when failed_filename is set" + ) + + +def _exception_headline(exc: Exception) -> str: + """First line of ``str(exc)`` — or the type name when that line is empty + or whitespace, so a report never opens with ``migration failed: `` and a + blank headline. asyncpg messages can be multiline (DETAIL/HINT lines); + the report keeps the first line only.""" + message = str(exc) + first_line = message.splitlines()[0] if message else "" + return first_line if first_line.strip() else type(exc).__name__ + + +async def diagnose_apply_failure( + conn: asyncpg.Connection, schema: str, exc: Exception +) -> ApplyFailureDiagnosis: + """Gather a self-diagnosis of a failed apply on the still-open conn. + + Both apply paths leave the connection reusable (a transactional failure + rolls back; the no-transaction path never opened one), so the same conn + can report what failed and what state the schema is in. Diagnosis must + NEVER mask the original error: every read is individually suppressed, + and whatever could not be gathered degrades to the generic report + (``failed_filename=None``). + """ + headline = _exception_headline(exc) + + failed: Migration | None = None + # apply_pending tags the exception with the migration that failed before + # re-raising — trust the tag over any heuristic. The isinstance guard + # keeps an exotic/mocked attribute from steering the report. + tagged = getattr(exc, "taskq_failed_migration", None) + if isinstance(tagged, Migration): + failed = tagged + else: + applied: set[str] | None = None + with contextlib.suppress(Exception): + applied = await list_applied(conn, schema) + if applied is not None: + # Fallback for exceptions from non-loop paths (e.g. the ledger + # ensure): apply_pending applies in discover() order and stops + # at the first failure, so the first unrecorded migration is the + # one that failed. Best-effort only — this misattributes under + # --phase, which is exactly why the loop tags. + with contextlib.suppress(Exception): + failed = next((m for m in discover() if m.key not in applied), None) + + invalid: list[str] = [] + if failed is not None and not failed.use_transaction: + with contextlib.suppress(Exception): + invalid = await list_invalid_indexes(conn, schema) + + return ApplyFailureDiagnosis( + headline=headline, + failed_filename=failed.filename if failed is not None else None, + use_transaction=failed.use_transaction if failed is not None else None, + invalid_indexes=tuple(invalid), + schema=schema, + ) + + +def render_apply_failure_lines(d: ApplyFailureDiagnosis, *, startup: bool = False) -> list[str]: + """Render a diagnosis as report lines (CLI stderr / startup SystemExit). + + ``startup=False`` reproduces the ``migrate up`` CLI report verbatim; + ``startup=True`` swaps only the action line for the restart-safe + variant (:data:`_STARTUP_ACTION_LINE`). + """ + if startup: + action = _STARTUP_ACTION_LINE + elif d.failed_filename is None: + action = ( + "Action: fix the error and re-run `taskq migrate up` — already-applied " + "migrations are skipped." + ) + elif d.use_transaction: + action = "Action: fix the error and re-run `taskq migrate up`." + else: + action = ( + "Action: re-run `taskq migrate up` — the migration is idempotent and " + "drops/rebuilds the debris itself." + ) + + if d.failed_filename is None: + return [f"migration failed: {d.headline}", action] + if d.use_transaction: + return [ + f"migration {d.failed_filename} failed: {d.headline}", + "It ran in a transaction and rolled back: nothing from the migration was applied.", + action, + ] + lines = [ + f"migration {d.failed_filename} failed: {d.headline}", + "It ran WITHOUT a transaction (-- taskq:no-transaction): statements " + "before the failure remain applied, and the migration was NOT recorded " + "in the ledger.", + ] + if d.invalid_indexes: + names = ", ".join(d.invalid_indexes) + lines.append( + f'INVALID index(es) in schema "{d.schema}": {names} — an interrupted ' + "CREATE INDEX CONCURRENTLY left them behind." + ) + lines.append(action) + return lines + + async def apply_pending_locked( dsn: str | None = None, *, @@ -293,7 +897,28 @@ async def apply_pending_locked( logger.info("no pending migrations") return applied except Exception as exc: - raise SystemExit(f"migration failed, aborting startup: {exc}") from exc + if c is None: + # The conn was never acquired (conn_factory/asyncpg.connect + # raised): there is nothing to diagnose on — keep the generic + # wrap. + raise SystemExit(f"migration failed, aborting startup: {exc}") from exc + diagnosis: ApplyFailureDiagnosis | None = None + with contextlib.suppress(Exception): + # Best-effort: the conn is still open (both apply paths leave it + # reusable), so report which migration failed and what state the + # schema is in instead of escaping a raw error. Diagnosis must + # never mask the original error — any surprise falls back to the + # generic wrap below. + diagnosis = await diagnose_apply_failure(c, schema, exc) + if diagnosis is None: + raise SystemExit(f"migration failed, aborting startup: {exc}") from exc + # Why one " — "-joined line: startup logs are grepped, not read as + # paragraphs, and the prefix stays stable for existing alerting + # rules pinned on it. + raise SystemExit( + "migration failed, aborting startup: " + + " — ".join(render_apply_failure_lines(diagnosis, startup=True)) + ) from exc finally: if c is not None: # Why the bounds: contextlib.suppress(Exception) catches errors diff --git a/tests/test_cli_migrate.py b/tests/test_cli_migrate.py index 56d2700..bffbc62 100644 --- a/tests/test_cli_migrate.py +++ b/tests/test_cli_migrate.py @@ -42,13 +42,16 @@ def terminate(self) -> None: self.close_wait.set() -def _make_migration(version: str, phase: str, filename: str) -> Migration: +def _make_migration( + version: str, phase: str, filename: str, *, use_transaction: bool = True +) -> Migration: return Migration( version=version, phase=phase, # type: ignore[arg-type] # Why: test fixture; Phase is Literal["pre", "post"]. description=f"{filename} description", filename=filename, sql_template="SELECT 1;", + use_transaction=use_transaction, ) @@ -164,6 +167,35 @@ def test_migrate_status_hung_close_does_not_mask_body_error_exit_code( assert fake_conn.terminated is True +def test_migrate_status_marks_no_transaction_migrations(monkeypatch: Any) -> None: + """migrate status annotates migrations that run outside a transaction so + operators can tell online-safe migrations from blocking ones.""" + _patch_connect(monkeypatch) + transactional = _make_migration("01.00.00_01", "pre", "01.00.00_01_pre_normal.sql") + no_transaction = Migration( + version="01.00.02_01", + phase="post", + description="concurrent index", + filename="01.00.02_01_post_concurrent_idx.sql", + sql_template="SELECT 1;", + use_transaction=False, + ) + + monkeypatch.setattr(cli_mod.migrate_mod, "list_applied", AsyncMock(return_value=set())) + monkeypatch.setattr( + cli_mod.migrate_mod, + "discover", + lambda: [transactional, no_transaction], + ) + + result = runner.invoke(app, ["migrate", "status"]) + assert result.exit_code == 0, f"stderr: {result.stderr}" + plain = plain_cli_output(result.output) + assert "01.00.02_01_post_concurrent_idx.sql (no transaction)" in plain + assert "01.00.00_01_pre_normal.sql (no transaction)" not in plain + assert "01.00.00_01_pre_normal.sql" in plain + + # ── migrate up ───────────────────────────────────────────────────────────── @@ -239,6 +271,229 @@ def test_migrate_up_closes_connection(monkeypatch: Any) -> None: assert fake_conn.terminated is False +# ── migrate up failure diagnosis ──────────────────────────────────────────── + + +def test_migrate_up_transactional_failure_reports_rollback_and_rerun( + monkeypatch: Any, +) -> None: + """A failed transactional migration rolls the whole file back: the CLI + names the migration, says nothing was applied, and prescribes + fix-and-re-run — never a traceback. The error line is truncated to its + first line (asyncpg messages can be multiline).""" + _patch_connect(monkeypatch) + failing = _make_migration("01.00.00_01", "pre", "01.00.00_01_pre_failing.sql") + monkeypatch.setattr( + cli_mod.migrate_mod, + "apply_pending", + AsyncMock(side_effect=RuntimeError("deadlock detected\nProcess 123 waits for ShareLock")), + ) + monkeypatch.setattr(cli_mod.migrate_mod, "discover", lambda: [failing]) + monkeypatch.setattr(cli_mod.migrate_mod, "list_applied", AsyncMock(return_value=set())) + + result = runner.invoke(app, ["migrate", "up"]) + assert result.exit_code == 1 + plain = plain_cli_output(result.output) + assert "migration 01.00.00_01_pre_failing.sql failed: deadlock detected" in plain + assert "Process 123 waits" not in plain, "only the exception's first line may be printed" + assert "transaction" in plain + assert "rolled back" in plain + assert "nothing" in plain and "applied" in plain + assert "taskq migrate up" in plain + assert "Traceback" not in plain + + +def test_migrate_up_no_transaction_failure_lists_invalid_indexes(monkeypatch: Any) -> None: + """A failed no-transaction migration keeps its partial effects: the CLI + says the migration was NOT recorded, names any INVALID indexes an + interrupted concurrent build left behind, and prescribes re-running the + idempotent migration.""" + _patch_connect(monkeypatch) + failing = _make_migration( + "01.00.02_01", "post", "01.00.02_01_post_concurrent_idx.sql", use_transaction=False + ) + monkeypatch.setattr( + cli_mod.migrate_mod, + "apply_pending", + AsyncMock(side_effect=RuntimeError("canceling statement due to statement timeout")), + ) + monkeypatch.setattr(cli_mod.migrate_mod, "discover", lambda: [failing]) + monkeypatch.setattr(cli_mod.migrate_mod, "list_applied", AsyncMock(return_value=set())) + monkeypatch.setattr( + cli_mod.migrate_mod, + "list_invalid_indexes", + AsyncMock(return_value=["jobs_queue_idx"]), + ) + + result = runner.invoke(app, ["migrate", "up"]) + assert result.exit_code == 1 + plain = plain_cli_output(result.output) + assert "migration 01.00.02_01_post_concurrent_idx.sql failed:" in plain + assert "WITHOUT a transaction" in plain + assert "NOT recorded" in plain + assert 'INVALID index(es) in schema "taskq": jobs_queue_idx' in plain + assert "idempotent" in plain + assert "taskq migrate up" in plain + assert "Traceback" not in plain + + +def test_migrate_up_no_transaction_failure_without_invalid_indexes(monkeypatch: Any) -> None: + """With no INVALID-index debris, the INVALID line is omitted while the + re-run guidance is still printed.""" + _patch_connect(monkeypatch) + failing = _make_migration( + "01.00.02_01", "post", "01.00.02_01_post_concurrent_idx.sql", use_transaction=False + ) + monkeypatch.setattr( + cli_mod.migrate_mod, + "apply_pending", + AsyncMock(side_effect=RuntimeError("connection was closed mid-build")), + ) + monkeypatch.setattr(cli_mod.migrate_mod, "discover", lambda: [failing]) + monkeypatch.setattr(cli_mod.migrate_mod, "list_applied", AsyncMock(return_value=set())) + monkeypatch.setattr(cli_mod.migrate_mod, "list_invalid_indexes", AsyncMock(return_value=[])) + + result = runner.invoke(app, ["migrate", "up"]) + assert result.exit_code == 1 + plain = plain_cli_output(result.output) + assert "migration 01.00.02_01_post_concurrent_idx.sql failed:" in plain + assert "INVALID" not in plain + assert "taskq migrate up" in plain + assert "Traceback" not in plain + + +def test_migrate_up_failure_diagnosis_never_masks_original_error(monkeypatch: Any) -> None: + """A diagnostic query failing (e.g. the connection died with the + migration) must not replace the original error: the base message and + re-run action are still printed, with no secondary exception and no + traceback.""" + _patch_connect(monkeypatch) + failing = _make_migration("01.00.00_01", "pre", "01.00.00_01_pre_failing.sql") + monkeypatch.setattr( + cli_mod.migrate_mod, + "apply_pending", + AsyncMock(side_effect=RuntimeError("boom")), + ) + monkeypatch.setattr(cli_mod.migrate_mod, "discover", lambda: [failing]) + monkeypatch.setattr( + cli_mod.migrate_mod, + "list_applied", + AsyncMock(side_effect=RuntimeError("conn dead")), + ) + + result = runner.invoke(app, ["migrate", "up"]) + assert result.exit_code == 1 + plain = plain_cli_output(result.output) + assert "boom" in plain + assert "conn dead" not in plain + assert "taskq migrate up" in plain + assert "Traceback" not in plain + + +def test_migrate_up_failure_report_survives_diagnosis_itself_raising( + monkeypatch: Any, +) -> None: + """Belt-and-braces: if diagnose_apply_failure itself blows up (a bug, or + a conn failure mode its suppressions don't cover), the CLI must still + print the ORIGINAL error and the re-run action — a diagnostic must never + mask the failure it diagnoses.""" + _patch_connect(monkeypatch) + monkeypatch.setattr( + cli_mod.migrate_mod, + "apply_pending", + AsyncMock(side_effect=RuntimeError("boom")), + ) + monkeypatch.setattr( + cli_mod.migrate_mod, + "diagnose_apply_failure", + AsyncMock(side_effect=RuntimeError("diagnosis exploded")), + ) + + result = runner.invoke(app, ["migrate", "up"]) + assert result.exit_code == 1 + plain = plain_cli_output(result.output) + assert "migration failed: boom" in plain + assert "diagnosis exploded" not in plain + assert "taskq migrate up" in plain + assert "Traceback" not in plain + + +def test_migrate_up_failure_report_exact_stderr_lines(monkeypatch: Any) -> None: + """Byte-exact pin of one full CLI failure report (content and order) — + complements the substring pins above; the renderer-level line lists are + pinned in tests/test_migrations_unit.py.""" + _patch_connect(monkeypatch) + failing = _make_migration("01.00.00_01", "pre", "01.00.00_01_pre_failing.sql") + monkeypatch.setattr( + cli_mod.migrate_mod, + "apply_pending", + AsyncMock(side_effect=RuntimeError("deadlock detected")), + ) + monkeypatch.setattr(cli_mod.migrate_mod, "discover", lambda: [failing]) + monkeypatch.setattr(cli_mod.migrate_mod, "list_applied", AsyncMock(return_value=set())) + + result = runner.invoke(app, ["migrate", "up"]) + assert result.exit_code == 1 + assert result.stderr.splitlines() == [ + "migration 01.00.00_01_pre_failing.sql failed: deadlock detected", + "It ran in a transaction and rolled back: nothing from the migration was applied.", + "Action: fix the error and re-run `taskq migrate up`.", + ] + + +def test_migrate_up_connect_failure_reports_generic_and_skips_close( + monkeypatch: Any, +) -> None: + """asyncpg.connect itself can fail (PG down, bad DSN) BEFORE the apply: + the CLI must still print the short report — original error plus re-run + action, never a traceback — and must not attempt a close on a + connection that was never acquired.""" + monkeypatch.setattr( + cli_mod.asyncpg, + "connect", + AsyncMock(side_effect=OSError("connection refused")), + ) + close_spy = AsyncMock() + monkeypatch.setattr(cli_mod, "close_conn_bounded", close_spy) + + result = runner.invoke(app, ["migrate", "up"]) + assert result.exit_code == 1 + plain = plain_cli_output(result.output) + assert "migration failed: connection refused" in plain + assert "taskq migrate up" in plain + assert "Traceback" not in plain + close_spy.assert_not_called() + + +def test_migrate_up_phase_failure_report_names_tagged_migration( + monkeypatch: Any, +) -> None: + """Under ``migrate up --phase pre`` an earlier-version :post migration can + still be pending and sorts FIRST in discover() order, so the naive + first-unrecorded heuristic would name the wrong file. apply_pending tags + the exception with the migration that actually failed; the report must + name the tagged (later) file.""" + _patch_connect(monkeypatch) + earlier_pending_post = _make_migration("01.00.00_01", "post", "01.00.00_01_post_pending.sql") + failing_pre = _make_migration("01.00.02_01", "pre", "01.00.02_01_pre_failing.sql") + exc = RuntimeError("deadlock detected") + exc.__dict__["taskq_failed_migration"] = failing_pre + monkeypatch.setattr(cli_mod.migrate_mod, "apply_pending", AsyncMock(side_effect=exc)) + monkeypatch.setattr( + cli_mod.migrate_mod, + "discover", + lambda: [earlier_pending_post, failing_pre], + ) + monkeypatch.setattr(cli_mod.migrate_mod, "list_applied", AsyncMock(return_value=set())) + + result = runner.invoke(app, ["migrate", "up", "--phase", "pre"]) + assert result.exit_code == 1 + plain = plain_cli_output(result.output) + assert "migration 01.00.02_01_pre_failing.sql failed: deadlock detected" in plain + assert "01.00.00_01_post_pending.sql" not in plain + assert "Traceback" not in plain + + async def test_migrate_up_terminates_hung_conn_close(monkeypatch: Any) -> None: """A hung conn close at migrate up exit (dead PG) is terminated after the bounded timeout and the command body completes.""" diff --git a/tests/test_migrate_no_transaction.py b/tests/test_migrate_no_transaction.py new file mode 100644 index 0000000..4b4fa75 --- /dev/null +++ b/tests/test_migrate_no_transaction.py @@ -0,0 +1,739 @@ +"""Integration tests (real Postgres) for non-transactional migrations. + +A migration carrying the ``-- taskq:no-transaction`` header directive is +discovered with ``use_transaction=False`` and applied WITHOUT the +per-migration transaction wrapper, making ``CREATE INDEX CONCURRENTLY`` / +``DROP INDEX CONCURRENTLY`` expressible. These tests pin the contract: + +- a non-transactional migration actually runs outside a transaction (a + concurrent index build succeeds — impossible inside a transaction block); +- the default path still wraps in a transaction (CONCURRENTLY is rejected, + and a failing mid-file statement rolls the whole file back); +- the ledger records a non-transactional migration only AFTER its statements + succeed — a failure leaves the key unrecorded while partial effects + persist, so such migrations must be idempotent and re-runnable; +- re-running after a failure is safe; +- the interrupted-``CREATE INDEX CONCURRENTLY`` failure mode (an INVALID + index left behind) is remedied by the documented drop-and-rebuild pattern; +- the ledger surfaces the distinction via ``schema_migrations.use_transaction``, + self-healed onto pre-upgrade ledgers. + +Synthetic migrations are layered on top of the bundled set by monkeypatching +``discover()`` (same pattern as ``test_migrate_coverage.py``) — except +``test_discover_directive_parsing_applies_end_to_end``, which patches +``importlib.resources.files`` instead so the REAL ``discover()`` directive +parsing is exercised against a real database. Each test uses its own +``new_base62()``-suffixed schema name. +""" + +from __future__ import annotations + +import asyncio +import contextlib +from importlib import resources +from pathlib import Path + +import asyncpg +import pytest +import structlog.testing +from typer.testing import CliRunner + +from taskq import migrate as migrate_mod +from taskq._ids import new_base62 +from taskq.cli import app +from taskq.migrate import Migration +from taskq.testing.assertions import plain_cli_output + +pytestmark = pytest.mark.integration + + +def _fake_migration( + version: str, phase: str, sql: str, *, use_transaction: bool = True +) -> Migration: + return Migration( + version=version, + phase=phase, # type: ignore[arg-type] # Why: test fixture; Phase is Literal["pre", "post"]. + description="synthetic", + filename=f"{version}_{phase}_synthetic.sql", + sql_template=sql, + use_transaction=use_transaction, + ) + + +async def _drop_schema(conn: asyncpg.Connection, schema: str) -> None: + await conn.execute(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE') + + +async def _bootstrap(conn: asyncpg.Connection, schema: str) -> list[Migration]: + """Apply the bundled migrations (schema + ledger now exist) and return + the real discovery list so tests can layer synthetics on top of it.""" + applied = await migrate_mod.apply_pending(conn, schema=schema) + assert applied, "expected bundled migrations to apply" + return migrate_mod.discover() + + +async def _index_validity(conn: asyncpg.Connection, schema: str, index: str) -> bool | None: + """``pg_index.indisvalid`` for ``index`` in ``schema``; None if absent.""" + return await conn.fetchval( + """ + SELECT i.indisvalid + FROM pg_class c + JOIN pg_index i ON i.indexrelid = c.oid + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = $1 AND c.relname = $2 + """, + schema, + index, + ) + + +async def _ledger_transactions(conn: asyncpg.Connection, schema: str) -> dict[str, bool]: + rows = await conn.fetch( + f'SELECT version, use_transaction FROM "{schema}".schema_migrations' # noqa: S608 # Why: schema is a test-generated identifier, not user input; asyncpg has no parameter binding for identifiers. + ) + return {r["version"]: r["use_transaction"] for r in rows} + + +# ── Non-transactional path: CONCURRENTLY works ────────────────────────────── + + +async def test_no_transaction_migration_runs_create_index_concurrently( + pg_dsn: str, monkeypatch: pytest.MonkeyPatch +) -> None: + schema = f"mig_nt_cc_{new_base62()}".lower() + conn = await asyncpg.connect(pg_dsn) + try: + await _drop_schema(conn, schema) + real = await _bootstrap(conn, schema) + m = _fake_migration( + "90.01.00_01", + "post", + "-- taskq:no-transaction\n" + 'DROP INDEX CONCURRENTLY IF EXISTS "{schema}".nt_jobs_queue_idx;\n' + "CREATE INDEX CONCURRENTLY IF NOT EXISTS nt_jobs_queue_idx " + 'ON "{schema}".jobs (queue);\n', + use_transaction=False, + ) + monkeypatch.setattr(migrate_mod, "discover", lambda: [*real, m]) + + with structlog.testing.capture_logs() as captured: + applied = await migrate_mod.apply_pending(conn, schema=schema) + + assert [x.key for x in applied] == [m.key] + assert any(e.get("event") == "migration-no-transaction" for e in captured), ( + "applying a migration outside a transaction must be logged" + ) + # A concurrent build cannot run inside a transaction block — success + # here proves the migration ran outside one, and left a VALID index. + assert await _index_validity(conn, schema, "nt_jobs_queue_idx") is True + ledger = await _ledger_transactions(conn, schema) + assert ledger[m.key] is False + # Bundled migrations applied through the default path record True. + assert all(ledger[x.key] is True for x in real) + finally: + await _drop_schema(conn, schema) + await conn.close() + + +async def test_apply_pending_locked_applies_no_transaction_migration( + pg_dsn: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """The startup path (--migrate / TASKQ_MIGRATE_ON_START) holds a session + advisory lock — not a transaction — so CONCURRENTLY still works.""" + schema = f"mig_nt_lock_{new_base62()}".lower() + conn = await asyncpg.connect(pg_dsn) + try: + await _drop_schema(conn, schema) + real = await _bootstrap(conn, schema) + m = _fake_migration( + "90.02.00_01", + "post", + "-- taskq:no-transaction\n" + "CREATE INDEX CONCURRENTLY IF NOT EXISTS nt_locked_idx " + 'ON "{schema}".jobs (status);\n', + use_transaction=False, + ) + monkeypatch.setattr(migrate_mod, "discover", lambda: [*real, m]) + finally: + await conn.close() + + async def _factory() -> asyncpg.Connection: + return await asyncpg.connect(pg_dsn) + + try: + applied = await migrate_mod.apply_pending_locked(schema=schema, conn_factory=_factory) + assert [x.key for x in applied] == [m.key] + + conn = await asyncpg.connect(pg_dsn) + try: + assert await _index_validity(conn, schema, "nt_locked_idx") is True + ledger = await _ledger_transactions(conn, schema) + assert ledger[m.key] is False + finally: + await conn.close() + finally: + conn = await asyncpg.connect(pg_dsn) + try: + await _drop_schema(conn, schema) + finally: + await conn.close() + + +# ── Directive parsing end-to-end (REAL discover()) ───────────────────────── + + +async def test_discover_directive_parsing_applies_end_to_end( + pg_dsn: str, monkeypatch: pytest.MonkeyPatch, tmp_path: Path +) -> None: + """Every other test here hand-sets ``use_transaction=False``, so the + directive-parsing path itself was never exercised against a real DB. This + test goes through the REAL ``discover()``: the synthetic file carries the + directive WITH a trailing note (the common real-world form), and the + migration must be applied non-transactionally based on that parse alone.""" + schema = f"mig_nt_disc_{new_base62()}".lower() + # Resolve and copy the bundled *.sql files BEFORE patching + # resources.files — apply_pending re-discovers on every call, so the + # patched dir must contain the full bundled set plus the synthetic file. + real_dir = resources.files("taskq.migrations") + for entry in real_dir.iterdir(): + if entry.is_file() and entry.name.endswith(".sql"): + (tmp_path / entry.name).write_bytes(entry.read_bytes()) + synth_name = "90.05.00_01_post_directive_file.sql" + (tmp_path / synth_name).write_text( + "-- taskq:no-transaction — CIC cannot run inside a transaction\n" + 'DROP INDEX CONCURRENTLY IF EXISTS "{schema}".nt_discovered_idx;\n' + "CREATE INDEX CONCURRENTLY IF NOT EXISTS nt_discovered_idx " + 'ON "{schema}".jobs (queue);\n', + encoding="utf-8", + ) + conn = await asyncpg.connect(pg_dsn) + try: + await _drop_schema(conn, schema) + # Bootstrap via the REAL package dir; only then patch, so the + # synthetic file is applied from the patched discovery below. + await _bootstrap(conn, schema) + monkeypatch.setattr(migrate_mod.resources, "files", lambda _pkg: tmp_path) + + applied = await migrate_mod.apply_pending(conn, schema=schema) + + # Parsed from the file — not hand-set. Asserted AFTER apply_pending + # so a parsing regression surfaces as its DB-level symptom + # (ActiveSQLTransactionError above) rather than a local assert. + m = next(x for x in migrate_mod.discover() if x.filename == synth_name) + assert m.use_transaction is False, "directive must be parsed from the file" + assert [x.key for x in applied] == [m.key] + assert await _index_validity(conn, schema, "nt_discovered_idx") is True + ledger = await _ledger_transactions(conn, schema) + assert ledger[m.key] is False + finally: + await _drop_schema(conn, schema) + await conn.close() + + +# ── Default path: still transactional ─────────────────────────────────────── + + +async def test_transactional_migration_rejects_concurrently( + pg_dsn: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """A migration WITHOUT the directive still runs inside a transaction, so + Postgres rejects CREATE INDEX CONCURRENTLY — pinning that the default + wrapper is intact and that the directive is what unlocks it.""" + schema = f"mig_nt_txcc_{new_base62()}".lower() + conn = await asyncpg.connect(pg_dsn) + try: + await _drop_schema(conn, schema) + real = await _bootstrap(conn, schema) + m = _fake_migration( + "90.03.00_01", + "post", + 'CREATE INDEX CONCURRENTLY IF NOT EXISTS nt_blocked_idx ON "{schema}".jobs (queue);\n', + ) + monkeypatch.setattr(migrate_mod, "discover", lambda: [*real, m]) + + with pytest.raises(asyncpg.ActiveSQLTransactionError): + await migrate_mod.apply_pending(conn, schema=schema) + + assert await _index_validity(conn, schema, "nt_blocked_idx") is None + assert m.key not in await migrate_mod.list_applied(conn, schema) + finally: + await _drop_schema(conn, schema) + await conn.close() + + +async def test_transactional_migration_rolls_back_on_failure( + pg_dsn: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """Regression guard for the default path: a failure mid-file rolls back + the whole migration — the table from the first statement must not persist.""" + schema = f"mig_nt_txrb_{new_base62()}".lower() + conn = await asyncpg.connect(pg_dsn) + try: + await _drop_schema(conn, schema) + real = await _bootstrap(conn, schema) + m = _fake_migration( + "90.04.00_01", + "post", + 'CREATE TABLE "{schema}".nt_rolled_back (id int);\nTHIS IS NOT VALID SQL;\n', + ) + monkeypatch.setattr(migrate_mod, "discover", lambda: [*real, m]) + + with pytest.raises(asyncpg.PostgresSyntaxError): + await migrate_mod.apply_pending(conn, schema=schema) + + table_exists = await conn.fetchval( + """ + SELECT EXISTS ( + SELECT 1 FROM information_schema.tables + WHERE table_schema = $1 AND table_name = 'nt_rolled_back' + ) + """, + schema, + ) + assert table_exists is False, "transactional migration must roll back fully" + assert m.key not in await migrate_mod.list_applied(conn, schema) + finally: + await _drop_schema(conn, schema) + await conn.close() + + +# ── Non-transactional failure modes ───────────────────────────────────────── + + +async def test_failed_no_transaction_migration_is_not_recorded_but_effects_persist( + pg_dsn: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """THE distinguishing failure mode: statements before the failure commit + independently (no rollback), yet the ledger records nothing — completion + is recorded only after every statement succeeds.""" + schema = f"mig_nt_fail_{new_base62()}".lower() + conn = await asyncpg.connect(pg_dsn) + try: + await _drop_schema(conn, schema) + real = await _bootstrap(conn, schema) + m = _fake_migration( + "90.05.00_01", + "post", + "-- taskq:no-transaction\n" + 'CREATE TABLE "{schema}".nt_partial (id int);\n' + "SELECT nonexistent_function_xyz();\n", + use_transaction=False, + ) + monkeypatch.setattr(migrate_mod, "discover", lambda: [*real, m]) + + with pytest.raises(asyncpg.UndefinedFunctionError): + await migrate_mod.apply_pending(conn, schema=schema) + + table_exists = await conn.fetchval( + """ + SELECT EXISTS ( + SELECT 1 FROM information_schema.tables + WHERE table_schema = $1 AND table_name = 'nt_partial' + ) + """, + schema, + ) + assert table_exists is True, ( + "no transaction wrapper: the first statement must NOT roll back" + ) + assert m.key not in await migrate_mod.list_applied(conn, schema), ( + "failed migration must not be marked applied" + ) + finally: + await _drop_schema(conn, schema) + await conn.close() + + +async def test_rerun_of_failed_no_transaction_migration_is_safe( + pg_dsn: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """After a failed non-transactional apply, re-running an idempotent + version of the same migration key succeeds and is recorded — without + duplicating effects from the first, partial run.""" + schema = f"mig_nt_rerun_{new_base62()}".lower() + conn = await asyncpg.connect(pg_dsn) + try: + await _drop_schema(conn, schema) + real = await _bootstrap(conn, schema) + broken = _fake_migration( + "90.06.00_01", + "post", + "-- taskq:no-transaction\n" + 'CREATE TABLE IF NOT EXISTS "{schema}".nt_rerun (id int);\n' + 'INSERT INTO "{schema}".nt_rerun VALUES (1);\n' + "SELECT nonexistent_function_xyz();\n", + use_transaction=False, + ) + monkeypatch.setattr(migrate_mod, "discover", lambda: [*real, broken]) + with pytest.raises(asyncpg.UndefinedFunctionError): + await migrate_mod.apply_pending(conn, schema=schema) + + fixed = _fake_migration( + "90.06.00_01", + "post", + "-- taskq:no-transaction\n" + 'CREATE TABLE IF NOT EXISTS "{schema}".nt_rerun (id int);\n' + 'INSERT INTO "{schema}".nt_rerun SELECT 1 ' + 'WHERE NOT EXISTS (SELECT 1 FROM "{schema}".nt_rerun);\n', + use_transaction=False, + ) + monkeypatch.setattr(migrate_mod, "discover", lambda: [*real, fixed]) + applied = await migrate_mod.apply_pending(conn, schema=schema) + + assert [m.key for m in applied] == [fixed.key] + assert fixed.key in await migrate_mod.list_applied(conn, schema) + rows = await conn.fetch(f'SELECT id FROM "{schema}".nt_rerun') # noqa: S608 # Why: schema is a test-generated identifier, not user input. + assert [r["id"] for r in rows] == [1], ( + "first run's INSERT persisted; the idempotent re-run must not duplicate it" + ) + finally: + await _drop_schema(conn, schema) + await conn.close() + + +async def test_interrupted_concurrent_build_remedy_drop_and_rebuild( + pg_dsn: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """An interrupted CREATE INDEX CONCURRENTLY leaves an INVALID index. The + documented remedy — DROP INDEX CONCURRENTLY IF EXISTS then CREATE INDEX + CONCURRENTLY IF NOT EXISTS in one non-transactional migration — replaces + it with a valid index and is then recorded.""" + schema = f"mig_nt_inv_{new_base62()}".lower() + conn = await asyncpg.connect(pg_dsn) + try: + await _drop_schema(conn, schema) + real = await _bootstrap(conn, schema) + + # Stage the debris of an interrupted build: a 2M-row table makes the + # concurrent build slow enough that a 100ms statement_timeout cancels + # it mid-build, leaving an INVALID index behind. On a heavily loaded + # runner the cancel can instead land during CIC's catalog-registration + # phase (no index row survives), so retry the staging a few times — + # each attempt is independent and well under a second. + await conn.execute( + f'CREATE TABLE "{schema}".nt_stage AS SELECT generate_series(1, 2000000) AS id' + ) + staged = False + for _attempt in range(3): + await conn.execute(f'DROP INDEX IF EXISTS "{schema}".nt_stage_idx') + await conn.execute("SET statement_timeout = '100ms'") + try: + with contextlib.suppress(asyncpg.QueryCanceledError): + await conn.execute( + f'CREATE INDEX CONCURRENTLY nt_stage_idx ON "{schema}".nt_stage (id)' + ) + finally: + await conn.execute("RESET statement_timeout") + validity = await _index_validity(conn, schema, "nt_stage_idx") + if validity is False: + staged = True + break + # None: cancelled before catalog registration (slow CI). True: + # build finished inside the timeout (absurdly fast machine). + assert staged, "could not stage an INVALID index via statement_timeout" + + m = _fake_migration( + "90.07.00_01", + "post", + "-- taskq:no-transaction\n" + 'DROP INDEX CONCURRENTLY IF EXISTS "{schema}".nt_stage_idx;\n' + "CREATE INDEX CONCURRENTLY IF NOT EXISTS nt_stage_idx " + 'ON "{schema}".nt_stage (id);\n', + use_transaction=False, + ) + monkeypatch.setattr(migrate_mod, "discover", lambda: [*real, m]) + applied = await migrate_mod.apply_pending(conn, schema=schema) + + assert [x.key for x in applied] == [m.key] + assert await _index_validity(conn, schema, "nt_stage_idx") is True, ( + "the remedy migration must replace the INVALID index with a valid one" + ) + ledger = await _ledger_transactions(conn, schema) + assert ledger[m.key] is False + finally: + await _drop_schema(conn, schema) + await conn.close() + + +async def test_list_invalid_indexes_reports_then_clears_staged_debris(pg_dsn: str) -> None: + """The CLI's failure report leans on ``list_invalid_indexes``: after an + interrupted CREATE INDEX CONCURRENTLY leaves an INVALID index behind, the + helper must name it; once the debris is dropped, it must report nothing.""" + schema = f"mig_nt_lii_{new_base62()}".lower() + conn = await asyncpg.connect(pg_dsn) + try: + await _drop_schema(conn, schema) + await _bootstrap(conn, schema) + + # Same staging pattern as + # test_interrupted_concurrent_build_remedy_drop_and_rebuild: a + # 2M-row table + 100ms statement_timeout cancels CIC mid-build, + # leaving an INVALID index (retry — the cancel can instead land + # before catalog registration on a loaded runner). + await conn.execute( + f'CREATE TABLE "{schema}".lii_stage AS SELECT generate_series(1, 2000000) AS id' + ) + staged = False + for _attempt in range(3): + await conn.execute(f'DROP INDEX IF EXISTS "{schema}".lii_stage_idx') + await conn.execute("SET statement_timeout = '100ms'") + try: + with contextlib.suppress(asyncpg.QueryCanceledError): + await conn.execute( + f'CREATE INDEX CONCURRENTLY lii_stage_idx ON "{schema}".lii_stage (id)' + ) + finally: + await conn.execute("RESET statement_timeout") + validity = await _index_validity(conn, schema, "lii_stage_idx") + if validity is False: + staged = True + break + assert staged, "could not stage an INVALID index via statement_timeout" + + assert await migrate_mod.list_invalid_indexes(conn, schema) == ["lii_stage_idx"] + + await conn.execute(f'DROP INDEX "{schema}".lii_stage_idx') + assert await migrate_mod.list_invalid_indexes(conn, schema) == [] + finally: + await _drop_schema(conn, schema) + await conn.close() + + +async def test_no_transaction_migration_rejects_transaction_control_statements( + pg_dsn: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """BEGIN/COMMIT inside a no-transaction file would re-open an explicit + transaction on the caller's connection — defeating CONCURRENTLY and, on + failure, leaving the connection in an aborted transaction. The runner + rejects the file BEFORE executing anything.""" + schema = f"mig_nt_txctl_{new_base62()}".lower() + conn = await asyncpg.connect(pg_dsn) + try: + await _drop_schema(conn, schema) + real = await _bootstrap(conn, schema) + m = _fake_migration( + "90.09.00_01", + "post", + "-- taskq:no-transaction\n" + "BEGIN;\n" + 'CREATE INDEX CONCURRENTLY IF NOT EXISTS nt_txctl_idx ON "{schema}".jobs (queue);\n' + "COMMIT;\n", + use_transaction=False, + ) + monkeypatch.setattr(migrate_mod, "discover", lambda: [*real, m]) + + with pytest.raises(ValueError, match="transaction-control"): + await migrate_mod.apply_pending(conn, schema=schema) + + assert await _index_validity(conn, schema, "nt_txctl_idx") is None, ( + "nothing must execute when the guard rejects the file" + ) + assert m.key not in await migrate_mod.list_applied(conn, schema) + # The caller's connection is untouched (no open/aborted transaction): + assert await conn.fetchval("SELECT 1") == 1 + finally: + await _drop_schema(conn, schema) + await conn.close() + + +# ── Ledger surfacing / upgrade path ───────────────────────────────────────── + + +async def test_runner_self_heals_ledger_column_and_backfills_default( + pg_dsn: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """Pre-upgrade deployments have a schema_migrations ledger without the + use_transaction column. The runner adds it when recording the next + migration; pre-existing rows backfill to true (they all ran inside a + transaction), and new rows record how they actually ran.""" + schema = f"mig_nt_heal_{new_base62()}".lower() + conn = await asyncpg.connect(pg_dsn) + try: + await _drop_schema(conn, schema) + await conn.execute(f'CREATE SCHEMA "{schema}"') + await conn.execute( + f""" + CREATE TABLE "{schema}".schema_migrations ( + version text PRIMARY KEY, + applied_at timestamptz NOT NULL DEFAULT now(), + checksum text NOT NULL + ) + """ + ) + await conn.execute( + f'INSERT INTO "{schema}".schema_migrations (version, checksum) VALUES ($1, $2)', # noqa: S608 + "00.00.00_01:pre", + "0" * 64, + ) + + tx = _fake_migration( + "90.08.00_01", "post", 'CREATE TABLE "{schema}".nt_heal_tx (id int);\n' + ) + nt = _fake_migration( + "90.08.00_02", + "post", + '-- taskq:no-transaction\nCREATE TABLE "{schema}".nt_heal_nt (id int);\n', + use_transaction=False, + ) + monkeypatch.setattr(migrate_mod, "discover", lambda: [tx, nt]) + applied = await migrate_mod.apply_pending(conn, schema=schema) + + assert [m.key for m in applied] == [tx.key, nt.key] + ledger = await _ledger_transactions(conn, schema) + assert ledger == { + "00.00.00_01:pre": True, + tx.key: True, + nt.key: False, + } + finally: + await _drop_schema(conn, schema) + await conn.close() + + +# ── CLI failure report (end-to-end) ───────────────────────────────────────── + + +def test_migrate_up_cli_reports_failed_no_transaction_migration( + pg_dsn: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """Owner requirement, end-to-end: a failed no-transaction ``migrate up`` + must itself report what failed, the state it left the schema in, and the + one action to take — never a traceback, never a manual-inspection + runbook. The partial table existing afterwards proves the report told + the truth.""" + schema = f"mig_nt_cli_{new_base62()}".lower() + + async def _setup() -> list[Migration]: + conn = await asyncpg.connect(pg_dsn) + try: + await _drop_schema(conn, schema) + return await _bootstrap(conn, schema) + finally: + await conn.close() + + real = asyncio.run(_setup()) + m = _fake_migration( + "90.10.00_01", + "post", + "-- taskq:no-transaction\n" + 'CREATE TABLE "{schema}".nt_cli_persist (id int);\n' + "THIS IS NOT VALID SQL;\n", + use_transaction=False, + ) + monkeypatch.setattr(migrate_mod, "discover", lambda: [*real, m]) + monkeypatch.setenv("TASKQ_PG_DSN", pg_dsn) + monkeypatch.setenv("TASKQ_SCHEMA_NAME", schema) + + # CliRunner runs in-process, so the monkeypatched discover and the env + # vars apply to the real CLI; it is synchronous, so this test drives + # async setup/verify/teardown through asyncio.run and must itself stay + # sync (asyncpg connections are bound to the loop that created them — + # one asyncio.run per phase, like conftest's _pg_admin). + result = CliRunner().invoke(app, ["migrate", "up"]) + + async def _table_exists() -> bool: + conn = await asyncpg.connect(pg_dsn) + try: + return bool( + await conn.fetchval( + """ + SELECT EXISTS ( + SELECT 1 FROM information_schema.tables + WHERE table_schema = $1 AND table_name = 'nt_cli_persist' + ) + """, + schema, + ) + ) + finally: + await conn.close() + + async def _cleanup() -> None: + conn = await asyncpg.connect(pg_dsn) + try: + await _drop_schema(conn, schema) + finally: + await conn.close() + + try: + assert result.exit_code == 1 + plain = plain_cli_output(result.output) + assert m.filename in plain + assert "WITHOUT a transaction" in plain + assert "NOT recorded" in plain + assert "taskq migrate up" in plain + assert "Traceback" not in plain + assert asyncio.run(_table_exists()) is True, ( + "the first statement must remain applied — the report said so" + ) + finally: + asyncio.run(_cleanup()) + + +# ── apply_pending_locked startup failure self-diagnosis ────────────────────── + + +async def test_apply_pending_locked_failure_self_diagnoses( + pg_dsn: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """Owner requirement, startup path: a migration failing under + ``apply_pending_locked`` (worker/UI startup) must abort with the SAME + self-diagnosis the CLI prints — which migration failed, the partial + state it left, the INVALID indexes it found, and the single action — + joined into ONE greppable SystemExit line, never a raw traceback.""" + schema = f"mig_nt_se_{new_base62()}".lower() + conn = await asyncpg.connect(pg_dsn) + try: + await _drop_schema(conn, schema) + real = await _bootstrap(conn, schema) + + # Stage INVALID-index debris (same statement_timeout pattern as + # test_list_invalid_indexes_reports_then_clears_staged_debris): an + # interrupted CIC leaves an INVALID index the diagnosis must name. + await conn.execute( + f'CREATE TABLE "{schema}".se_stage AS SELECT generate_series(1, 2000000) AS id' + ) + staged = False + for _attempt in range(3): + await conn.execute(f'DROP INDEX IF EXISTS "{schema}".se_stage_idx') + await conn.execute("SET statement_timeout = '100ms'") + try: + with contextlib.suppress(asyncpg.QueryCanceledError): + await conn.execute( + f'CREATE INDEX CONCURRENTLY se_stage_idx ON "{schema}".se_stage (id)' + ) + finally: + await conn.execute("RESET statement_timeout") + validity = await _index_validity(conn, schema, "se_stage_idx") + if validity is False: + staged = True + break + assert staged, "could not stage an INVALID index via statement_timeout" + + m = _fake_migration( + "90.11.00_01", + "post", + "-- taskq:no-transaction\n" + 'CREATE TABLE "{schema}".se_persist (id int);\n' + "THIS IS NOT VALID SQL;\n", + use_transaction=False, + ) + monkeypatch.setattr(migrate_mod, "discover", lambda: [*real, m]) + + async def _conn_factory() -> asyncpg.Connection: + return await asyncpg.connect(pg_dsn) + + with pytest.raises(SystemExit) as excinfo: + await migrate_mod.apply_pending_locked(schema=schema, conn_factory=_conn_factory) + message = str(excinfo.value) + assert "migration failed, aborting startup" in message + assert m.filename in message + assert "WITHOUT a transaction" in message + assert "NOT recorded" in message + assert f'INVALID index(es) in schema "{schema}": se_stage_idx' in message + assert "restart is safe" in message + assert "\n" not in message, "the startup report must be one greppable line" + assert "Traceback" not in message + finally: + await _drop_schema(conn, schema) + await conn.close() diff --git a/tests/test_migrate_runner.py b/tests/test_migrate_runner.py index c3954d1..5fe0282 100644 --- a/tests/test_migrate_runner.py +++ b/tests/test_migrate_runner.py @@ -171,6 +171,73 @@ async def test_guard_passes_when_pre_applied_earlier_in_same_run( assert conn.executed.count("SELECT 1;") == 5 +# ── apply_pending: failure tagging for self-diagnosis ─────────────────── + + +class _FailOnMarkerConn(_FakeMigrateConn): + """_FakeMigrateConn whose execute raises when the SQL contains a marker. + + Lets a test fail ONE specific migration mid-apply while every other + migration executes normally. + """ + + def __init__(self, applied: set[str], fail_marker: str) -> None: + super().__init__(applied) + self._fail_marker = fail_marker + + async def execute(self, sql: str, *args: object) -> str: + if self._fail_marker in sql: + raise RuntimeError(f"synthetic apply failure: {self._fail_marker}") + return await super().execute(sql, *args) + + +def _failing_migration( + version: str, phase: str, marker: str, *, use_transaction: bool +) -> Migration: + return Migration( + version=version, + phase=phase, # type: ignore[arg-type] # Why: test fixture; Phase is Literal["pre", "post"]. + description="fabricated failing", + filename=f"{version}_{phase}_fabricated.sql", + sql_template=f"SELECT '{marker}';", + use_transaction=use_transaction, + ) + + +async def test_apply_pending_tags_exception_with_failing_transactional_migration( + monkeypatch: Any, +) -> None: + """The self-diagnosis can only name the right file if the failure + carries WHICH migration failed — the first-unrecorded-in-discover-order + heuristic is wrong under ``--phase``. apply_pending must tag the raised + exception with the failing migration and re-raise the SAME object (the + type pins in test_migrate_no_transaction.py forbid wrapping).""" + failing = _failing_migration("02.00.00_01", "pre", "tx-fail-marker", use_transaction=True) + _patch_discover(monkeypatch, [_make_migration("01.00.00_01", "pre"), failing]) + conn = _FailOnMarkerConn(applied={"01.00.00_01:pre"}, fail_marker="tx-fail-marker") + + with pytest.raises(RuntimeError, match="synthetic apply failure") as excinfo: + await migrate_mod.apply_pending(conn, schema="taskq") # type: ignore[arg-type] + + assert getattr(excinfo.value, "taskq_failed_migration", None) is failing + + +async def test_apply_pending_tags_exception_from_no_transaction_statement( + monkeypatch: Any, +) -> None: + """The no-transaction statement loop gets the same tagging: a mid-file + statement failure is attributed to its own migration, not to whatever + sorts first in discover() order.""" + failing = _failing_migration("02.00.00_01", "pre", "nt-fail-marker", use_transaction=False) + _patch_discover(monkeypatch, [_make_migration("01.00.00_01", "pre"), failing]) + conn = _FailOnMarkerConn(applied={"01.00.00_01:pre"}, fail_marker="nt-fail-marker") + + with pytest.raises(RuntimeError, match="synthetic apply failure") as excinfo: + await migrate_mod.apply_pending(conn, schema="taskq") # type: ignore[arg-type] + + assert getattr(excinfo.value, "taskq_failed_migration", None) is failing + + # ── apply_pending_locked: bounded finally teardown (dead PG) ──────────── diff --git a/tests/test_migrations_populated.py b/tests/test_migrations_populated.py new file mode 100644 index 0000000..eff3eaf --- /dev/null +++ b/tests/test_migrations_populated.py @@ -0,0 +1,871 @@ +"""Bundled migrations applied STEPWISE onto a POPULATED database. + +Owner requirement: every bundled migration — present AND future — must run +cleanly as an end user against a database that already holds data. The +stepwise test below applies migrations one at a time (``max_steps=1``) and +seeds a full slice of realistic rows after each step, so migration N+1 +always applies onto the data volumes migration N left behind. A migration +that damages a populated database fails CI naming the offending migration +key in the assertion message. + +Two tiers live here: + +* ``test_seed_data_is_deterministic`` — pure (no PG) contract for the row + generator: identical output across calls, globally-unique idempotency + keys, exact status distribution, and no singleton-metadata on + active-status rows (which ``jobs_singleton_uniq`` would reject). +* ``test_bundled_migrations_apply_stepwise_onto_populated_database`` — the + integration harness: stepwise apply + per-step invariants + per-step + seeding, then a final integrity pass and a functional smoke through the + REAL backend API, ending with the end-user CLI contract (``migrate up`` + on a fully-migrated populated DB is a safe no-op). + +Runtime budget: the module must stay under ~45s (CI runs ``pytest -n 2``). +Tune ONLY via the ``SEED_*`` volume knobs below — halve ``SEED_JOBS`` +first, trim ``SEED_EVENTS_PER_JOB`` second. +""" + +from __future__ import annotations + +import asyncio +import uuid +from collections import Counter +from collections.abc import Awaitable, Callable +from contextlib import AsyncExitStack +from dataclasses import dataclass +from datetime import UTC, datetime, timedelta + +import asyncpg +import pytest +from typer.testing import CliRunner + +from taskq import migrate as migrate_mod +from taskq._ids import new_base62, new_uuid +from taskq._json import dumps, dumps_str, loads +from taskq.backend._sql_templates import COPY_FROM_COLUMNS +from taskq.backend.clock import SystemClock +from taskq.backend.postgres import PostgresBackend +from taskq.cli import app +from taskq.settings import WorkerSettings +from taskq.testing.assertions import plain_cli_output +from taskq.testing.jobs import make_enqueue_args +from taskq.testing.pg import DEFAULT_ACTORS, seed_actors +from taskq.testing.settings import make_integration_settings_dict +from taskq.worker.deps import open_worker_deps + +pytestmark = pytest.mark.integration + +# ── Volume knobs ────────────────────────────────────────────────────────── +# Tunable in ONE place. SEED_JOBS is per slice; across the full stepwise run +# ~10k ACTIVE-status rows land in the hot partial indexes (dispatch / +# scheduled-wake / running-lock), which is the volume that matters for +# index-building migrations. Keep SEED_JOBS a multiple of 100: the status +# wheel below then yields the declared distribution exactly. + +SEED_JOBS = 7_000 +SEED_EVENTS_PER_JOB = 2 +SEED_ATTEMPTS = 2_000 +SEED_ARCHIVE_ROWS = 1_000 + +# Declared status distribution (percent). The stepwise harness's job: prove +# migrations behave against terminal rows (result/error payloads), active +# rows (locks, future schedules), and everything between. +_STATUS_MIX: tuple[tuple[str, int], ...] = ( + ("succeeded", 55), + ("failed", 10), + ("cancelled", 5), + ("crashed", 3), + ("abandoned", 2), + ("pending", 15), + ("scheduled", 5), + ("running", 5), +) +_TERMINAL_ORDER: tuple[str, ...] = ("succeeded", "failed", "cancelled", "crashed", "abandoned") +_TERMINAL_STATUSES = frozenset(_TERMINAL_ORDER) +_ACTIVE_STATUSES = frozenset({"pending", "scheduled", "running"}) + +# Flat 100-slot wheel: ``i % 100`` indexes into it, so per-slice counts are +# EXACT whenever SEED_JOBS is a multiple of 100. +_STATUS_WHEEL: tuple[str, ...] = tuple( + status for status, weight in _STATUS_MIX for _ in range(weight) +) +assert len(_STATUS_WHEEL) == 100 + +# Fixed uuid5 namespace: seed ids are pure functions of (slice, counter), so +# two generator calls with the same ``now`` produce byte-identical rows and +# the final integrity pass can re-derive any row's id without the generator. +_SEED_NAMESPACE = uuid.uuid5(uuid.NAMESPACE_URL, "taskq/tests/test_migrations_populated") +_SPREAD_DAYS = 90 + + +def _uuid5(name: str) -> uuid.UUID: + return uuid.uuid5(_SEED_NAMESPACE, name) + + +def _payload(i: int) -> dict[str, object]: + if i % 10 == 0: + # ~2KB payload every 10th row: COPY, GIN jsonb indexing, and the + # idempotency index build must handle more than toy documents. + return {"job": i, "kind": "large", "blob": "x" * 2048} + return {"job": i, "kind": "small"} + + +def _tags(slice_id: int, i: int) -> list[str]: + return ["seed", f"slice-{slice_id}", f"bucket-{i % 7}"] + + +def _result(i: int) -> dict[str, object]: + return {"ok": True, "value": i} + + +# Attempt outcome for a terminal job status; active jobs are given an +# earlier FAILED attempt (they were retried and are still in flight). +_ATTEMPT_OUTCOME = { + "succeeded": "succeeded", + "failed": "failed", + "cancelled": "cancelled", + "crashed": "crashed", +} + + +def _job_row( + slice_id: int, + i: int, + *, + status: str, + row_id: uuid.UUID, + idempotency_key: str | None, + now: datetime, + created_at: datetime, + worker_id: uuid.UUID, +) -> dict[str, object]: + """One canonical ``jobs`` row (every column any schema version knows). + + ``idempotency_scope`` is included unconditionally — the seeder's runtime + column intersection drops it against pre-``01.00.03_01:pre`` schemas. + """ + terminal = status in _TERMINAL_STATUSES + started_at = ( + created_at + timedelta(minutes=1) if status not in ("pending", "scheduled") else None + ) + finished_at = started_at + timedelta(minutes=5) if terminal and started_at else None + result = _result(i) if status == "succeeded" else None + + error_class = error_message = error_traceback = None + if status == "failed": + error_class, error_message, error_traceback = ( + "SeedError", + f"boom {i}", + "simulated traceback", + ) + elif status == "crashed": + error_class, error_message = "WorkerCrashed", "worker died mid-flight" + elif status == "abandoned": + error_class, error_message = "MaxAttemptsExceeded", "attempts exhausted" + + metadata: dict[str, object] = {"seed": True} + if terminal and i % 97 == 0: + # Singleton metadata is legal ONLY on terminal rows (the partial + # unique index covers active statuses); the generator never puts it + # on active rows — pinned by test_seed_data_is_deterministic. + metadata["singleton"] = True + + # scheduled rows are genuinely future-due; pending rows are long since due + scheduled_at = now + timedelta(hours=i % 72) if status == "scheduled" else created_at + + return { + "id": row_id, + "actor": ("actor_a", "actor_b", "actor_c", "test_actor")[i % 4], + "queue": "default", + "identity_key": f"ident-{i % 50}" if i % 4 == 0 else None, + "fairness_key": f"fair-{i % 10}" if i % 7 == 0 else None, + "payload": _payload(i), + "payload_schema_ver": 1, + "status": status, + "priority": i % 3, + "attempt": 0 if status in ("pending", "scheduled") else 1, + "max_attempts": 3, + "retry_kind": "transient", + "schedule_to_close": now + timedelta(days=1) if status in _ACTIVE_STATUSES else None, + "start_to_close": timedelta(minutes=5), + "heartbeat_timeout": timedelta(seconds=30) if status == "running" else None, + "created_at": created_at, + "scheduled_at": scheduled_at, + "started_at": started_at, + "finished_at": finished_at, + "last_heartbeat_at": now - timedelta(seconds=5) if status == "running" else None, + "locked_by_worker": worker_id if status == "running" else None, + "lock_expires_at": now + timedelta(hours=1) if status == "running" else None, + "cancel_requested_at": created_at + timedelta(minutes=2) if status == "cancelled" else None, + "cancel_phase": 2 if status == "cancelled" else 0, + "error_class": error_class, + "error_message": error_message, + "error_traceback": error_traceback, + "progress_state": {}, + "progress_seq": 0, + "result": result, + "result_size_bytes": len(dumps(result)) if result is not None else None, + "result_expires_at": finished_at + timedelta(days=7) if finished_at and result else None, + "idempotency_scope": "", + "idempotency_key": idempotency_key, + "trace_id": f"{i:032x}" if i % 6 == 0 else None, + "span_id": f"{i:016x}" if i % 6 == 0 else None, + "metadata": metadata, + "tags": _tags(slice_id, i), + } + + +@dataclass(frozen=True, slots=True) +class _SliceSeed: + """All rows one slice contributes, keyed by table. Pure data — no PG.""" + + workers: list[dict[str, object]] + jobs: list[dict[str, object]] + job_events: list[dict[str, object]] + job_attempts: list[dict[str, object]] + jobs_archive: list[dict[str, object]] + job_attempts_archive: list[dict[str, object]] + queues: list[dict[str, object]] + cron_schedules: list[dict[str, object]] + rate_limit_buckets: list[dict[str, object]] + reservation_slots: list[dict[str, object]] + + +def _generate_slice(slice_id: int, *, now: datetime) -> _SliceSeed: + """Deterministically generate one slice of seed rows. + + Pure: same ``(slice_id, now)`` in → identical rows out. No unseeded RNG + anywhere — ids are uuid5 over a fixed namespace, keys are counter-based, + timestamps spread over the trailing ``_SPREAD_DAYS`` days relative to + ``now``. + """ + worker_id = _uuid5(f"worker:{slice_id}") + base = now - timedelta(days=_SPREAD_DAYS) + + jobs: list[dict[str, object]] = [] + events: list[dict[str, object]] = [] + attempts: list[dict[str, object]] = [] + job_step = timedelta(seconds=_SPREAD_DAYS * 86_400 / SEED_JOBS) + for i in range(SEED_JOBS): + status = _STATUS_WHEEL[i % 100] + created_at = base + job_step * i + # Counter-based keys are globally unique ACROSS slices: the legacy + # single-column jobs_idempotency_key_uniq index enforces global + # uniqueness until 01.00.03_01:post drops it. + key = f"idem-{slice_id * SEED_JOBS + i:05d}" if i % 5 == 0 else None + job = _job_row( + slice_id, + i, + status=status, + row_id=_uuid5(f"job:{slice_id}:{i}"), + idempotency_key=key, + now=now, + created_at=created_at, + worker_id=worker_id, + ) + jobs.append(job) + + # Kinds per the job_events contract comment: state_change (a subset + # with detail->>'reason'='lock_expired' so 01.00.02_01's partial + # reclaim index has real rows to index when it builds), progress, + # cancel_request, heartbeat_miss. + kind = ("state_change", "progress", "cancel_request", "heartbeat_miss")[i % 4] + if kind == "state_change": + detail: dict[str, object] = { + "reason": "lock_expired", + "from_state": "running", + "to_state": "pending", + } + elif kind == "progress": + detail = {"pct": i % 100} + elif kind == "cancel_request": + detail = {"requested_by": "seed"} + else: + detail = {"missed": 2} + events.append( + { + "job_id": job["id"], + "occurred_at": created_at + timedelta(minutes=1), + "kind": kind, + "detail": detail, + } + ) + events.append( + { + "job_id": job["id"], + "occurred_at": created_at + timedelta(minutes=2), + "kind": "state_change", + "detail": {"from_state": "pending", "to_state": status}, + } + ) + + if i < SEED_ATTEMPTS: + outcome = _ATTEMPT_OUTCOME.get(status, "failed") + attempts.append( + { + "job_id": job["id"], + "attempt": 1, + "started_at": created_at + timedelta(minutes=1), + "finished_at": created_at + timedelta(minutes=2), + "outcome": outcome, + "error_class": "SeedError" if outcome in ("failed", "crashed") else None, + "error_message": f"boom {i}" if outcome in ("failed", "crashed") else None, + "error_traceback": None, + "duration_ms": 60_000, + "worker_id": worker_id, + "metadata": {}, + } + ) + + archive: list[dict[str, object]] = [] + archive_attempts: list[dict[str, object]] = [] + archive_step = timedelta(seconds=_SPREAD_DAYS * 86_400 / SEED_ARCHIVE_ROWS) + for i in range(SEED_ARCHIVE_ROWS): + # jobs_archive mirrors jobs (and is ALTERed by 01.00.03_01:pre too), + # so it is seeded through the same canonical row builder. + status = _TERMINAL_ORDER[i % len(_TERMINAL_ORDER)] + key = f"idem-arc-{slice_id * SEED_ARCHIVE_ROWS + i:05d}" if i % 5 == 0 else None + created_at = base + archive_step * i + row = _job_row( + slice_id, + i, + status=status, + row_id=_uuid5(f"arc:{slice_id}:{i}"), + idempotency_key=key, + now=now, + created_at=created_at, + worker_id=worker_id, + ) + row["archived_at"] = now - timedelta(days=1) + row["expire_at"] = now + timedelta(days=300) + archive.append(row) + archive_attempts.append( + { + "job_id": row["id"], + "attempt": 1, + "started_at": created_at + timedelta(minutes=1), + "finished_at": created_at + timedelta(minutes=2), + "outcome": _ATTEMPT_OUTCOME.get(status, "failed"), + "error_class": None, + "error_message": None, + "error_traceback": None, + "duration_ms": 60_000, + "worker_id": worker_id, + "metadata": {}, + } + ) + + workers = [ + { + "id": worker_id, + "hostname": f"seed-host-{slice_id}", + "pid": 10_000 + slice_id, + "queues": ["default"], + "started_at": base, + "last_seen_at": now - timedelta(minutes=1), + "worker_label": f"seed-{slice_id}", + "workgroup_instance": None, + "metadata": {"seed": True}, + } + ] + queues = [ + { + "name": "default" if slice_id == 0 and q == 0 else f"q_{slice_id}_{q}", + "mode": ("strict_fifo", "round_robin")[q % 2], + "created_at": base, + "updated_at": now - timedelta(days=1), + "max_concurrent": None if q == 0 else 10, + } + for q in range(2) + ] + cron = [ + { + "id": _uuid5(f"cron:{slice_id}:{c}"), + # DISTINCT actors: cron_schedules has UNIQUE(actor) until + # 01.00.01_01 replaces it with UNIQUE(actor, name). + "actor": f"cron_actor_{slice_id}_{c}", + "cron_expr": "*/5 * * * *", + "timezone": "UTC", + "dst_strategy": "skip", + "payload_factory": None, + "enabled": True, + "last_fired_at": None, + "last_fire_error": None, + "consecutive_failures": 0, + "next_fire_at": now + timedelta(hours=1), + "metadata": {}, + "name": f"sched_{slice_id}_{c}", + "identity_key": None, + } + for c in range(2) + ] + buckets = [ + { + "bucket_name": f"bucket_{slice_id}_{b}", + "kind": "token_bucket", + "state": {"tokens": 5.0, "capacity": 10.0}, + "updated_at": now - timedelta(minutes=5), + } + for b in range(2) + ] + # Two free + two held slots per bucket; held slots reference this + # slice's running jobs (wheel positions 95/96). + slots = [ + { + "bucket_name": f"rsv_{slice_id}", + "slot_index": s, + "job_id": jobs[95 + s - 2]["id"] if s >= 2 else None, + "held_by_worker_id": worker_id if s >= 2 else None, + "acquired_at": now - timedelta(minutes=5) if s >= 2 else None, + "lease_expires_at": now + timedelta(hours=1) if s >= 2 else None, + } + for s in range(4) + ] + + return _SliceSeed( + workers=workers, + jobs=jobs, + job_events=events, + job_attempts=attempts, + jobs_archive=archive, + job_attempts_archive=archive_attempts, + queues=queues, + cron_schedules=cron, + rate_limit_buckets=buckets, + reservation_slots=slots, + ) + + +# ── Column supersets ───────────────────────────────────────────────────── +# Canonical ORDERED supersets of every column any schema version knows, in +# COPY order. For `jobs` that superset IS the library's own COPY_FROM_COLUMNS +# (single source of truth for the COPY path; pinned to the live table by +# tests/test_migrations.py::test_copy_from_columns_match_jobs_table_exactly). + +_WORKERS_COLUMNS: tuple[str, ...] = ( + "id", + "hostname", + "pid", + "queues", + "started_at", + "last_seen_at", + "worker_label", + "workgroup_instance", + "metadata", +) +_QUEUES_COLUMNS: tuple[str, ...] = ("name", "mode", "created_at", "updated_at", "max_concurrent") +_CRON_COLUMNS: tuple[str, ...] = ( + "id", + "actor", + "cron_expr", + "timezone", + "dst_strategy", + "payload_factory", + "enabled", + "last_fired_at", + "last_fire_error", + "consecutive_failures", + "next_fire_at", + "metadata", + "name", + "identity_key", +) +# `id` omitted deliberately: bigserial assigns it. +_JOB_EVENTS_COLUMNS: tuple[str, ...] = ("job_id", "occurred_at", "kind", "detail") +_JOB_ATTEMPTS_COLUMNS: tuple[str, ...] = ( + "job_id", + "attempt", + "started_at", + "finished_at", + "outcome", + "error_class", + "error_message", + "error_traceback", + "duration_ms", + "worker_id", + "metadata", +) +_JOBS_ARCHIVE_COLUMNS: tuple[str, ...] = (*COPY_FROM_COLUMNS, "archived_at", "expire_at") +_RATE_LIMIT_COLUMNS: tuple[str, ...] = ("bucket_name", "kind", "state", "updated_at") +_RESERVATION_COLUMNS: tuple[str, ...] = ( + "bucket_name", + "slot_index", + "job_id", + "held_by_worker_id", + "acquired_at", + "lease_expires_at", +) + +# asyncpg's binary COPY encodes jsonb from str — dict values must be +# serialized (dumps_str) before they go into a record. +_JSONB_COLUMNS: dict[str, frozenset[str]] = { + "workers": frozenset({"metadata"}), + "jobs": frozenset({"payload", "progress_state", "result", "metadata"}), + "job_events": frozenset({"detail"}), + "job_attempts": frozenset({"metadata"}), + "jobs_archive": frozenset({"payload", "progress_state", "result", "metadata"}), + "job_attempts_archive": frozenset({"metadata"}), + "queues": frozenset(), + "cron_schedules": frozenset({"metadata"}), + "rate_limit_buckets": frozenset({"state"}), + "reservation_slots": frozenset(), +} + + +async def _live_columns(conn: asyncpg.Connection, schema: str, table: str) -> set[str]: + rows = await conn.fetch( + """ + SELECT column_name FROM information_schema.columns + WHERE table_schema = $1 AND table_name = $2 + """, + schema, + table, + ) + return {str(r["column_name"]) for r in rows} + + +async def _seed_slice( + conn: asyncpg.Connection, schema: str, *, slice_id: int, counts: dict[str, int] +) -> None: + """Bulk-load one generated slice into the CURRENT schema shape. + + The runtime column intersection (live ``information_schema.columns`` ∩ + canonical superset) is deliberate: when a future migration adds a + column, the seeder adapts and keeps loading; when a future migration + adds a NOT NULL-without-default column the seeder FAILS — and that + failure is the CI alarm, because the harness's job is precisely to + surface migrations that cannot run against a populated database. + + ``counts`` accumulates expected per-table row totals across slices so + the caller can assert nothing was silently lost or duplicated. + """ + seed = _generate_slice(slice_id, now=datetime.now(UTC)) + # actor_config goes through the library's own seeder (ON CONFLICT DO + # NOTHING makes it idempotent across slices). + await seed_actors(conn, schema) + counts["actor_config"] = len(DEFAULT_ACTORS) + + plan: list[tuple[str, tuple[str, ...], list[dict[str, object]]]] = [ + ("workers", _WORKERS_COLUMNS, seed.workers), + ("jobs", COPY_FROM_COLUMNS, seed.jobs), + ("job_events", _JOB_EVENTS_COLUMNS, seed.job_events), + ("job_attempts", _JOB_ATTEMPTS_COLUMNS, seed.job_attempts), + ("jobs_archive", _JOBS_ARCHIVE_COLUMNS, seed.jobs_archive), + ("job_attempts_archive", _JOB_ATTEMPTS_COLUMNS, seed.job_attempts_archive), + ("queues", _QUEUES_COLUMNS, seed.queues), + ("cron_schedules", _CRON_COLUMNS, seed.cron_schedules), + ("rate_limit_buckets", _RATE_LIMIT_COLUMNS, seed.rate_limit_buckets), + ("reservation_slots", _RESERVATION_COLUMNS, seed.reservation_slots), + ] + for table, superset, rows in plan: + live = await _live_columns(conn, schema, table) + columns = [c for c in superset if c in live] + jsonb = _JSONB_COLUMNS[table] + records = [ + tuple( + dumps_str(row[c]) if c in jsonb and row[c] is not None else row[c] for c in columns + ) + for row in rows + ] + # The tr_notify_job_insert trigger fires pg_notify per COPY'd + # pending row; with no listener Postgres discards them — harmless. + await conn.copy_records_to_table( + table, records=records, columns=columns, schema_name=schema + ) + counts[table] = counts.get(table, 0) + len(records) + + +# ── Catalog helpers (same shapes as tests/test_migrate_no_transaction.py) ── + + +async def _drop_schema(conn: asyncpg.Connection, schema: str) -> None: + await conn.execute(f'DROP SCHEMA IF EXISTS "{schema}" CASCADE') + + +async def _index_validity(conn: asyncpg.Connection, schema: str, index: str) -> bool | None: + """``pg_index.indisvalid`` for ``index`` in ``schema``; None if absent.""" + return await conn.fetchval( + """ + SELECT i.indisvalid + FROM pg_class c + JOIN pg_index i ON i.indexrelid = c.oid + JOIN pg_namespace n ON n.oid = c.relnamespace + WHERE n.nspname = $1 AND c.relname = $2 + """, + schema, + index, + ) + + +async def _ledger_transactions(conn: asyncpg.Connection, schema: str) -> dict[str, bool]: + rows = await conn.fetch(f'SELECT version, use_transaction FROM "{schema}".schema_migrations') + return {r["version"]: r["use_transaction"] for r in rows} + + +async def _assert_table_counts( + conn: asyncpg.Connection, schema: str, counts: dict[str, int] +) -> None: + for table, expected in sorted(counts.items()): + actual = await conn.fetchval(f'SELECT count(*) FROM "{schema}"."{table}"') + assert actual == expected, ( + f"{table}: expected {expected} cumulative rows after seeding, found {actual}" + ) + + +# ── Migration-specific checks ───────────────────────────────────────────── +# Generic per-step invariants (runner order, no INVALID indexes, ledger +# use_transaction) apply to EVERY discovered key — unknown keys get ONLY +# those, so future migrations (incl. PRs #25/#27's CIC index rebuilds) +# automatically join this harness the day they land. Entries below are for +# migrations whose populated-DB effect deserves a sharper assertion. + +_MigrationCheck = Callable[[asyncpg.Connection, str], Awaitable[None]] + + +async def _noop_check(conn: asyncpg.Connection, schema: str) -> None: + pass + + +async def _check_idempotency_scope_post(conn: asyncpg.Connection, schema: str) -> None: + """After ``01.00.03_01:post``: the legacy global-unique index is GONE + and the composite scope-key index is VALID — with seeded idempotency + keys present in the table while the swap happened.""" + assert await _index_validity(conn, schema, "jobs_idempotency_key_uniq") is None, ( + "01.00.03_01:post must drop the legacy single-column idempotency index" + ) + assert await _index_validity(conn, schema, "jobs_idempotency_scope_key_uniq") is True, ( + "the composite (idempotency_scope, idempotency_key) index must be VALID" + ) + + +_MIGRATION_SPECIFIC_CHECKS: dict[str, _MigrationCheck] = { + "01.00.03_01:post": _check_idempotency_scope_post, +} + + +# ── Tier 1: generator contract (no PG) ──────────────────────────────────── + + +def test_seed_data_is_deterministic() -> None: + """Pure generator contract (no PG). + + The seeder must be fully deterministic (fixed ``now`` in, identical + rows out) so failures reproduce byte-for-byte, and its output must + respect the invariants of the EARLIEST schema version it seeds: the + legacy single-column ``jobs_idempotency_key_uniq`` index enforces + GLOBAL key uniqueness until ``01.00.03_01:post`` drops it, and + ``jobs_singleton_uniq`` forbids singleton metadata on active rows. + """ + fixed_now = datetime(2026, 3, 1, tzinfo=UTC) + + first = _generate_slice(0, now=fixed_now) + second = _generate_slice(0, now=fixed_now) + assert [r["id"] for r in first.jobs] == [r["id"] for r in second.jobs] + assert [r["idempotency_key"] for r in first.jobs] == [r["idempotency_key"] for r in second.jobs] + assert first == second + + # Idempotency keys must be globally unique ACROSS slices — the old + # single-column unique index spans every seeded generation until the + # :post migration drops it. + keys = [ + row["idempotency_key"] + for slice_id in range(8) + for row in _generate_slice(slice_id, now=fixed_now).jobs + if row["idempotency_key"] is not None + ] + assert len(keys) > 0 + assert len(keys) == len(set(keys)), "idempotency keys must be globally unique" + + # Status counts match the declared distribution exactly. + counts = Counter(str(row["status"]) for row in first.jobs) + assert counts == {status: SEED_JOBS * weight // 100 for status, weight in _STATUS_MIX} + + # jobs_singleton_uniq forbids active-status rows carrying singleton + # metadata — the seeder must never produce one. + for row in first.jobs: + if str(row["status"]) in _ACTIVE_STATUSES: + metadata = row["metadata"] + assert isinstance(metadata, dict) + assert not metadata.get("singleton"), ( + f"active-status job {row['id']} must not carry singleton metadata" + ) + + +# ── Tier 2: stepwise apply onto populated data ──────────────────────────── + + +def test_bundled_migrations_apply_stepwise_onto_populated_database( + pg_dsn: str, monkeypatch: pytest.MonkeyPatch +) -> None: + """Apply every bundled migration one step at a time onto populated data. + + Each iteration: apply exactly ONE migration → generic invariants → + migration-specific checks → seed the next slice → assert cumulative + counts. After the loop: data round-trips, index spot-checks, a + functional smoke through the real backend API, and the end-user CLI + contract. CliRunner runs in-process, so the monkeypatched env vars + apply to the real CLI; it is synchronous, so this test drives async + phases through asyncio.run and must itself stay sync (asyncpg + connections are bound to the loop that created them — one asyncio.run + per phase, like conftest's _pg_admin and tests/test_migrate_no_transaction.py). + """ + schema = f"mig_pop_{new_base62()}".lower() + + async def _stepwise() -> None: + conn = await asyncpg.connect(pg_dsn) + try: + await _drop_schema(conn, schema) + counts: dict[str, int] = {} + for slice_id, migration in enumerate(migrate_mod.discover()): + applied = await migrate_mod.apply_pending(conn, schema=schema, max_steps=1) + assert [m.key for m in applied] == [migration.key], ( + f"step {slice_id}: expected exactly {migration.key!r} to apply, got " + f"{[m.key for m in applied]} — runner order drifted or the " + "migration failed against a populated database" + ) + assert await migrate_mod.list_invalid_indexes(conn, schema) == [], ( + f"{migration.key} left INVALID indexes behind on populated data" + ) + ledger = await _ledger_transactions(conn, schema) + assert ledger[migration.key] is migration.use_transaction, ( + f"ledger use_transaction mismatch for {migration.key}" + ) + await _MIGRATION_SPECIFIC_CHECKS.get(migration.key, _noop_check)(conn, schema) + await _seed_slice(conn, schema, slice_id=slice_id, counts=counts) + await _assert_table_counts(conn, schema, counts) + finally: + await conn.close() + + async def _round_trips_and_index_spot_checks() -> None: + conn = await asyncpg.connect(pg_dsn) + try: + # Payload/tags/result round-trip on deterministic sample ids: + # i=1 (small payload, no key), i=10 (~2KB payload, keyed). + small = await conn.fetchrow( + f'SELECT payload, tags, result, status, idempotency_key FROM "{schema}".jobs WHERE id = $1', + _uuid5("job:0:1"), + ) + assert small is not None + assert loads(small["payload"]) == _payload(1) + assert list(small["tags"]) == _tags(0, 1) + assert loads(small["result"]) == _result(1) + assert small["status"] == "succeeded" + assert small["idempotency_key"] is None + + large = await conn.fetchrow( + f'SELECT payload, idempotency_key FROM "{schema}".jobs WHERE id = $1', + _uuid5("job:0:10"), + ) + assert large is not None + assert loads(large["payload"]) == _payload(10), "the ~2KB payload must survive intact" + assert large["idempotency_key"] == "idem-00010" + + running = await conn.fetchrow( + f'SELECT status, locked_by_worker, lock_expires_at FROM "{schema}".jobs WHERE id = $1', + _uuid5("job:0:95"), + ) + assert running is not None + assert running["status"] == "running" + assert running["locked_by_worker"] == _uuid5("worker:0") + assert running["lock_expires_at"] > datetime.now(UTC) + + archived = await conn.fetchrow( + f'SELECT status, result FROM "{schema}".jobs_archive WHERE id = $1', + _uuid5("arc:0:0"), + ) + assert archived is not None + assert archived["status"] == "succeeded" + assert loads(archived["result"]) == _result(0) + + # Index spot-checks: the hot dispatch index is VALID, and the + # lock_expired reclaim partial index (built by 01.00.02_01 over + # seeded events) exists. + assert await _index_validity(conn, schema, "jobs_dispatch_idx") is True + assert await _index_validity(conn, schema, "job_events_reclaim_idx") is True + finally: + await conn.close() + + async def _real_api_smoke() -> None: + """End-state functional smoke through the REAL public API. + + The real API is used ONLY here, against the final schema: the + current enqueue SQL inserts ``idempotency_scope``, which raises + UndefinedColumnError against pre-``01.00.03_01:pre`` schemas, and + ``job_events`` has no public append API at all — intermediate + schemas are therefore seeded raw (schema-shaped), never via enqueue. + """ + settings = WorkerSettings.load_from_dict(make_integration_settings_dict(pg_dsn)) + settings.schema_name = schema + stack = AsyncExitStack() + deps = await stack.enter_async_context(open_worker_deps(settings)) + try: + backend = PostgresBackend( + deps, + clock=SystemClock(), + cancellation_grace_period=timedelta( + seconds=deps.settings.cancellation_grace_period + ), + cleanup_grace_period=timedelta(seconds=deps.settings.cleanup_grace_period), + ) + key = f"smoke-{schema}" + row = await backend.enqueue( + make_enqueue_args(actor="test_actor", idempotency_key=key, priority=10) + ) + assert row.status == "pending" # type: ignore[comparison-overlap] # Why: JobStatus is Literal[...]; pyright narrows too conservatively across frozen dataclass fields (same as tests/test_dispatch_pg.py). + + duplicate = await backend.enqueue( + make_enqueue_args( + actor="test_actor", idempotency_key=key, payload={"v": 2}, priority=10 + ) + ) + assert duplicate.id == row.id, "same key + same scope must dedupe" + + scoped = await backend.enqueue( + make_enqueue_args( + actor="test_actor", idempotency_key=key, idempotency_scope="smoke-b" + ) + ) + assert scoped.id != row.id, ( + "the SAME key under a SECOND scope must insert a second row — only legal " + "after 01.00.03_01:post dropped the legacy global-unique index; the " + "sharpest end-to-end proof the deploy sequence landed on populated data" + ) + + # priority=10 outranks every seeded row (priorities 0..2), so + # limit=1 leases exactly our row. + worker_id = new_uuid() + dispatched = await backend.dispatch_batch( + worker_id=worker_id, + queues=["default"], + limit=1, + lock_lease=timedelta(seconds=30), + ) + assert len(dispatched) == 1 + assert dispatched[0].id == row.id + assert dispatched[0].status == "running" # type: ignore[comparison-overlap] # Why: see enqueue assertion above. + assert dispatched[0].locked_by_worker == worker_id + finally: + await stack.aclose() + + async def _cleanup() -> None: + conn = await asyncpg.connect(pg_dsn) + try: + await _drop_schema(conn, schema) + finally: + await conn.close() + + try: + asyncio.run(_stepwise()) + asyncio.run(_round_trips_and_index_spot_checks()) + asyncio.run(_real_api_smoke()) + + # End-user contract: re-running `migrate up` on a fully-migrated, + # populated database is a safe no-op (same CliRunner/env pattern as + # tests/test_migrate_no_transaction.py). + monkeypatch.setenv("TASKQ_PG_DSN", pg_dsn) + monkeypatch.setenv("TASKQ_SCHEMA_NAME", schema) + result = CliRunner().invoke(app, ["migrate", "up"]) + assert result.exit_code == 0 + assert "no pending migrations" in plain_cli_output(result.output) + finally: + asyncio.run(_cleanup()) diff --git a/tests/test_migrations_unit.py b/tests/test_migrations_unit.py index b35f41e..a7a9bfb 100644 --- a/tests/test_migrations_unit.py +++ b/tests/test_migrations_unit.py @@ -1,16 +1,49 @@ """Unit tests for the migration runner's pure functions (no PG required). -Covers ``discover``, ``render``, ``Migration`` dataclass behavior, and -``apply_pending``/``list_applied`` error paths that don't need a live -database connection. +Covers ``discover``, ``render``, ``Migration`` dataclass behavior, +``split_statements``, and ``apply_pending``/``list_applied`` error paths +that don't need a live database connection. """ from __future__ import annotations +from collections.abc import Callable +from unittest.mock import AsyncMock + import pytest +import structlog.testing from taskq import migrate -from taskq.migrate import Migration +from taskq.migrate import ( # pyright: ignore[reportPrivateUsage] # Why: unit-testing the guard directly; the end-to-end path is pinned separately in test_migrate_no_transaction.py. + Migration, + _reject_transaction_control, +) + + +def _fake_package(files: dict[str, str]) -> type: + """Build a ``importlib.resources.files()`` stand-in for monkeypatching. + + Keys are filenames, values are file contents. Contents are stored as + UTF-8 bytes and decoded with the requested encoding so tests exercise + ``discover()``'s real decode path (e.g. ``utf-8-sig`` BOM handling). + """ + + class FakeEntry: + def __init__(self, name: str, content: str) -> None: + self.name = name + self._content_bytes = content.encode("utf-8") + + def is_file(self) -> bool: + return True + + def read_text(self, encoding: str = "utf-8") -> str: + return self._content_bytes.decode(encoding) + + class FakePackage: + def iterdir(self): + return [FakeEntry(name, content) for name, content in files.items()] + + return FakePackage def test_discover_returns_sorted_migrations() -> None: @@ -93,6 +126,14 @@ def test_list_applied_rejects_invalid_schema() -> None: asyncio.run(migrate.list_applied(object(), "invalid;schema")) # type: ignore[arg-type] +def test_list_invalid_indexes_rejects_invalid_schema() -> None: + """list_invalid_indexes validates schema before touching the DB.""" + with pytest.raises(ValueError, match="invalid schema name"): + import asyncio + + asyncio.run(migrate.list_invalid_indexes(object(), "invalid;schema")) # type: ignore[arg-type] + + def test_apply_pending_rejects_invalid_schema() -> None: """apply_pending validates schema before touching the DB.""" with pytest.raises(ValueError, match="invalid schema name"): @@ -167,3 +208,707 @@ def iterdir(self): monkeypatch.setattr(migrate.resources, "files", lambda _pkg: FakePackage()) assert migrate.discover() == [] + + +# ── use_transaction / no-transaction directive ───────────────────────────── + + +def _discover_with_content(monkeypatch: pytest.MonkeyPatch, content: str) -> Migration: + files = {"90.00.00_01_post_directive.sql": content} + monkeypatch.setattr(migrate.resources, "files", lambda _pkg: _fake_package(files)()) + (m,) = migrate.discover() + return m + + +def test_migration_use_transaction_defaults_to_true() -> None: + """Existing call sites construct Migration without the field: default True.""" + m = Migration( + version="01.00.00_01", + phase="pre", + description="test", + filename="test.sql", + sql_template="SELECT 1;", + ) + assert m.use_transaction is True + + +def test_bundled_migrations_are_all_transactional() -> None: + """No bundled migration uses the no-transaction directive yet. + + Guards against accidentally retrofitting the mechanism onto existing + migrations, which the framework requires to stay transactional. + """ + assert migrate.discover(), "expected bundled migrations" + assert all(m.use_transaction for m in migrate.discover()) + + +def test_discover_marks_directive_in_leading_comment_block( + monkeypatch: pytest.MonkeyPatch, +) -> None: + m = _discover_with_content( + monkeypatch, + "-- Rebuild an index without locking writes.\n" + "-- taskq:no-transaction\n" + "-- The literal {schema} token is substituted at apply time.\n" + 'CREATE INDEX CONCURRENTLY IF NOT EXISTS t_idx ON "{schema}".t (id);\n', + ) + assert m.use_transaction is False + + +def test_discover_defaults_transactional_without_directive( + monkeypatch: pytest.MonkeyPatch, +) -> None: + m = _discover_with_content( + monkeypatch, + '-- Just a normal migration.\nCREATE TABLE "{schema}".t (id int);\n', + ) + assert m.use_transaction is True + + +def test_discover_accepts_directive_whitespace_variants( + monkeypatch: pytest.MonkeyPatch, +) -> None: + m = _discover_with_content( + monkeypatch, + " --taskq:no-transaction \nSELECT 1;\n", + ) + assert m.use_transaction is False + + +def test_discover_ignores_directive_after_first_sql_statement( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The directive only counts in the leading comment block — a stray + ``-- taskq:no-transaction`` later in the file must not flip semantics.""" + m = _discover_with_content( + monkeypatch, + 'CREATE TABLE "{schema}".t (id int);\n-- taskq:no-transaction\n', + ) + assert m.use_transaction is True + + +def test_discover_ignores_directive_in_later_comment_block( + monkeypatch: pytest.MonkeyPatch, +) -> None: + m = _discover_with_content( + monkeypatch, + "SELECT 1;\n-- a mid-file comment block\n-- taskq:no-transaction\nSELECT 2;\n", + ) + assert m.use_transaction is True + + +def test_discover_honors_directive_in_bom_prefixed_file( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """A UTF-8 BOM (Windows editors) must not silently disable the directive.""" + m = _discover_with_content( + monkeypatch, + "\ufeff-- taskq:no-transaction\nSELECT 1;\n", + ) + assert m.use_transaction is False + + +def test_discover_accepts_directive_with_trailing_note( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Real-world directives carry a reason after the token + (``-- taskq:no-transaction needed for CIC``); silently ignoring that + form defeats the opt-out — the migration runs transactional anyway.""" + m = _discover_with_content( + monkeypatch, + "-- taskq:no-transaction needed for CIC\nSELECT 1;\n", + ) + assert m.use_transaction is False + + +def test_discover_accepts_directive_mixed_case( + monkeypatch: pytest.MonkeyPatch, +) -> None: + m = _discover_with_content( + monkeypatch, + "-- TaskQ:No-Transaction\nSELECT 1;\n", + ) + assert m.use_transaction is False + + +def test_discover_rejects_directive_lookalike_and_warns( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """``-- taskq:no-transactional`` must NOT match (the token boundary stops + prefix drift), but it looks like a directive attempt — warn so the author + notices instead of silently running transactional.""" + with structlog.testing.capture_logs() as captured: + m = _discover_with_content( + monkeypatch, + "-- taskq:no-transactional\nSELECT 1;\n", + ) + assert m.use_transaction is True + warnings = [e for e in captured if e.get("event") == "migration-directive-unrecognized"] + assert warnings, "expected an unrecognized-directive warning" + assert warnings[0].get("filename") == "90.00.00_01_post_directive.sql" + assert "taskq:no-transactional" in str(warnings[0].get("line")) + + +def test_discover_warns_on_directive_typo( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with structlog.testing.capture_logs() as captured: + m = _discover_with_content( + monkeypatch, + "-- taskq:no-transction\nSELECT 1;\n", + ) + assert m.use_transaction is True + assert any(e.get("event") == "migration-directive-unrecognized" for e in captured) + + +def test_discover_exact_directive_logs_no_unrecognized_warning( + monkeypatch: pytest.MonkeyPatch, +) -> None: + with structlog.testing.capture_logs() as captured: + m = _discover_with_content( + monkeypatch, + "-- taskq:no-transaction\nSELECT 1;\n", + ) + assert m.use_transaction is False + assert not any(e.get("event") == "migration-directive-unrecognized" for e in captured) + + +def test_directive_changes_checksum(monkeypatch: pytest.MonkeyPatch) -> None: + """The directive lives inside the SQL template, so toggling it changes + the checksum — drift detection catches an edited-in-place directive.""" + sql = "-- taskq:no-transaction\nSELECT 1;\n" + with_directive = Migration( + version="90.00.00_01", + phase="post", + description="d", + filename="f.sql", + sql_template=sql, + ) + without_directive = Migration( + version="90.00.00_01", + phase="post", + description="d", + filename="f.sql", + sql_template="SELECT 1;\n", + ) + assert with_directive.checksum("taskq") != without_directive.checksum("taskq") + + +# ── split_statements ──────────────────────────────────────────────────────── + + +def test_split_statements_empty_input() -> None: + assert migrate.split_statements("") == [] + assert migrate.split_statements(" \n ") == [] + + +def test_split_statements_single_statement_without_semicolon() -> None: + assert migrate.split_statements("SELECT 1") == ["SELECT 1"] + + +def test_split_statements_two_statements() -> None: + assert migrate.split_statements("SELECT 1; SELECT 2;") == ["SELECT 1", "SELECT 2"] + + +def test_split_statements_ignores_semicolon_in_string_literal() -> None: + assert migrate.split_statements("INSERT INTO t VALUES ('a;b'); SELECT 1;") == [ + "INSERT INTO t VALUES ('a;b')", + "SELECT 1", + ] + + +def test_split_statements_handles_doubled_quote_escape() -> None: + assert migrate.split_statements("SELECT 'it''s;';") == ["SELECT 'it''s;'"] + + +def test_split_statements_handles_e_string_backslash_escape() -> None: + # In E'...' strings a backslash escapes the next char, so the first ' + # here does NOT terminate the string. + assert migrate.split_statements("SELECT E'a\\';b';") == ["SELECT E'a\\';b'"] + + +def test_split_statements_ignores_semicolon_in_quoted_identifier() -> None: + assert migrate.split_statements('SELECT 1 AS "we;ird";') == ['SELECT 1 AS "we;ird"'] + + +def test_split_statements_keeps_leading_comment_attached() -> None: + assert migrate.split_statements("-- a;b\nSELECT 1;") == ["-- a;b\nSELECT 1"] + + +def test_split_statements_ignores_semicolons_in_block_comment() -> None: + assert migrate.split_statements("/* one ; two ; */ SELECT 1; SELECT 2;") == [ + "/* one ; two ; */ SELECT 1", + "SELECT 2", + ] + + +def test_split_statements_handles_nested_block_comments() -> None: + sql = "/* outer /* inner ; */ still outer ; */ SELECT 1;" + assert migrate.split_statements(sql) == ["/* outer /* inner ; */ still outer ; */ SELECT 1"] + + +def test_split_statements_ignores_semicolons_in_dollar_quoted_body() -> None: + sql = ( + "CREATE FUNCTION f() RETURNS void AS $$\n" + "BEGIN\n" + " RAISE NOTICE 'x;y';\n" + "END;\n" + "$$ LANGUAGE plpgsql;\n" + "SELECT 1;" + ) + assert migrate.split_statements(sql) == [ + "CREATE FUNCTION f() RETURNS void AS $$\n" + "BEGIN\n" + " RAISE NOTICE 'x;y';\n" + "END;\n" + "$$ LANGUAGE plpgsql", + "SELECT 1", + ] + + +def test_split_statements_handles_tagged_dollar_quotes() -> None: + sql = "DO $body$\nBEGIN\n PERFORM 1;\nEND\n$body$;\nSELECT 1;" + assert migrate.split_statements(sql) == [ + "DO $body$\nBEGIN\n PERFORM 1;\nEND\n$body$", + "SELECT 1", + ] + + +def test_split_statements_handles_dollar_quote_inside_other_tag() -> None: + # A different tag inside a dollar-quoted body must not close it. + sql = "SELECT $outer$ a $inner$ b $inner$ c $outer$; SELECT 1;" + assert migrate.split_statements(sql) == [ + "SELECT $outer$ a $inner$ b $inner$ c $outer$", + "SELECT 1", + ] + + +def test_split_statements_drops_comment_only_chunks() -> None: + assert migrate.split_statements("-- only a comment;\n") == [] + assert migrate.split_statements("SELECT 1; -- trailing\n; SELECT 2;") == [ + "SELECT 1", + "SELECT 2", + ] + + +def test_split_statements_strips_trailing_semicolon_and_whitespace() -> None: + assert migrate.split_statements("\n SELECT 1 ;\n") == ["SELECT 1"] + + +def test_split_statements_handles_empty_dollar_quoted_string() -> None: + assert migrate.split_statements("SELECT $$$$;") == ["SELECT $$$$"] + + +def test_split_statements_handles_unicode_dollar_tag() -> None: + # Postgres dollar tags follow identifier rules and may be non-ASCII. + assert migrate.split_statements("SELECT $é$a;b$é$; SELECT 2;") == [ + "SELECT $é$a;b$é$", + "SELECT 2", + ] + + +def test_split_statements_cr_only_line_ending_ends_line_comment() -> None: + """Postgres ends ``--`` comments at ``\\r`` too; otherwise the comment + swallows the rest of a CR-only file and statements are silently dropped.""" + assert migrate.split_statements("SELECT 1; -- c\rSELECT 2;") == [ + "SELECT 1", + "-- c\rSELECT 2", + ] + + +def test_split_statements_e_string_right_after_statement_boundary() -> None: + assert migrate.split_statements("SELECT 1;E'\\'';") == ["SELECT 1", "E'\\''"] + + +def test_split_statements_dollar_sign_inside_identifier() -> None: + """``a$b$c`` is a legal Postgres identifier; ``$b$`` must not be read as + a dollar-quote opener here — a tag cannot immediately follow an + identifier character (same rule as the E'...' detection).""" + assert migrate.split_statements("SELECT a$b$c; SELECT 2;") == [ + "SELECT a$b$c", + "SELECT 2", + ] + + +def test_split_statements_dollar_quote_after_punctuation_still_parses() -> None: + """The identifier-char gate must not block real dollar quotes: after + whitespace, a comma, or an open paren there is no identifier char before + the ``$``, so the tag still opens a quoted body.""" + assert migrate.split_statements("SELECT 1, $tag$x$tag$ FROM t; SELECT 2;") == [ + "SELECT 1, $tag$x$tag$ FROM t", + "SELECT 2", + ] + assert migrate.split_statements("SELECT f($tag$x$tag$); SELECT 2;") == [ + "SELECT f($tag$x$tag$)", + "SELECT 2", + ] + + +def test_split_statements_unicode_escape_string_regression() -> None: + """Regression pins for U&'...' literals: backslash is NOT an escape + outside E'...', and '' still doubles.""" + assert migrate.split_statements("SELECT U&'d\\0061t\\+000061'; SELECT 2;") == [ + "SELECT U&'d\\0061t\\+000061'", + "SELECT 2", + ] + assert migrate.split_statements("SELECT U&'x''y'; SELECT 2;") == [ + "SELECT U&'x''y'", + "SELECT 2", + ] + + +# ── transaction-control guard for non-transactional migrations ────────────── + + +def _nt_migration() -> Migration: + return Migration( + version="90.00.00_01", + phase="post", + description="d", + filename="f.sql", + sql_template="-- taskq:no-transaction\nSELECT 1;\n", + use_transaction=False, + ) + + +def test_reject_transaction_control_accepts_plain_statements() -> None: + _reject_transaction_control(_nt_migration(), ["SELECT 1", "CREATE TABLE t (id int)"]) + + +@pytest.mark.parametrize( + "keyword", ["BEGIN", "begin", "COMMIT", "ROLLBACK", "END", "ABORT", "START"] +) +def test_reject_transaction_control_rejects_keywords(keyword: str) -> None: + with pytest.raises(ValueError, match="transaction-control"): + _reject_transaction_control(_nt_migration(), [f"{keyword} WORK"]) + + +def test_reject_transaction_control_rejects_after_leading_comments() -> None: + """split_statements keeps leading comments attached; the guard must see + past them or a ``-- comment\\nBEGIN`` would slip through.""" + with pytest.raises(ValueError, match="transaction-control"): + _reject_transaction_control(_nt_migration(), ["-- explain\n/* plus */\nBEGIN"]) + + +def test_reject_transaction_control_message_names_migration_and_keyword() -> None: + with pytest.raises(ValueError, match=r"f\.sql.*'COMMIT'"): + _reject_transaction_control(_nt_migration(), ["COMMIT"]) + + +@pytest.mark.parametrize( + "statement", + [ + "SAVEPOINT sp1", + "RELEASE SAVEPOINT sp1", + "SET LOCAL statement_timeout = 0", + "SET TRANSACTION ISOLATION LEVEL SERIALIZABLE", + "set local statement_timeout = 0", # lowercase variant + ], +) +def test_reject_transaction_control_rejects_transaction_scoped_statements( + statement: str, +) -> None: + """``SAVEPOINT``/``RELEASE`` would fail loudly at execution outside a + transaction, but the guard's value is rejecting the file BEFORE anything + runs. ``SET LOCAL``/``SET TRANSACTION`` are worse: outside a transaction + they are SILENT no-ops (server WARNING only), so the author believes e.g. + the statement timeout was disabled for a long build when it was not.""" + with pytest.raises(ValueError, match="transaction-control"): + _reject_transaction_control(_nt_migration(), [statement]) + + +def test_reject_transaction_control_rejects_set_local_after_leading_comments() -> None: + with pytest.raises(ValueError, match="transaction-control"): + _reject_transaction_control( + _nt_migration(), ["-- tune for the long build\nSET LOCAL statement_timeout = 0"] + ) + + +def test_reject_transaction_control_message_names_set_local_keyword() -> None: + with pytest.raises(ValueError, match=r"f\.sql.*'SET LOCAL'"): + _reject_transaction_control(_nt_migration(), ["SET LOCAL statement_timeout = 0"]) + + +def test_reject_transaction_control_accepts_session_set_and_checkpoint() -> None: + # Deliberate allowlist: plain SET / SET SESSION is session-scoped — it + # behaves identically inside and outside a transaction, so opting out + # changes nothing about it and it is not deceptive. CHECKPOINT is a + # cluster-level maintenance statement with no transaction semantics at + # all. Neither pretends the runner is managing a transaction, so the + # guard stays out of the way. + _reject_transaction_control( + _nt_migration(), + ["SET work_mem = '1GB'", "SET SESSION work_mem = '1GB'", "CHECKPOINT"], + ) + + +@pytest.mark.parametrize( + "statement", + [ + "SET /* x */ LOCAL statement_timeout = 0", + "SET -- x\nLOCAL statement_timeout = 0", + "SET/**/LOCAL statement_timeout = 0", + "/* a /* b */ c */ SET LOCAL statement_timeout = 0", + ], +) +def test_reject_transaction_control_rejects_comment_trivia_bypasses(statement: str) -> None: + """Comments are valid trivia here because Postgres treats them as + whitespace between keywords — ``SET /* x */ LOCAL`` is the same statement + to the server as ``SET LOCAL`` (and /* */ comments NEST), so the guard + must skip them too or these forms slip past as silent no-ops.""" + with pytest.raises(ValueError, match="transaction-control"): + _reject_transaction_control(_nt_migration(), [statement]) + + +def test_reject_transaction_control_rejects_begin_after_nested_comment() -> None: + """A nested /* ... /* ... */ ... */ block comment before BEGIN must not + hide it — Postgres sees straight through to the keyword.""" + with pytest.raises(ValueError, match="transaction-control"): + _reject_transaction_control(_nt_migration(), ["/* a /* b */ c */ BEGIN"]) + + +def test_reject_transaction_control_message_names_set_local_through_comments() -> None: + """The rejection message names the keyword even when comments intervene.""" + with pytest.raises(ValueError, match=r"f\.sql.*'SET LOCAL'"): + _reject_transaction_control(_nt_migration(), ["SET /* tune */ LOCAL statement_timeout = 0"]) + + +def test_reject_transaction_control_allows_session_set_through_comments() -> None: + """Comments between SET and a session-scoped keyword keep it ALLOWED: + the keyword is what decides, not the trivia around it.""" + _reject_transaction_control( + _nt_migration(), + ["SET /* tune */ work_mem = '1GB'", "SET -- tune\nSESSION work_mem = '1GB'"], + ) + + +# ── apply-failure diagnosis rendering ──────────────────────────────────────── + +# The startup action line differs from the CLI's: a worker/startup apply +# failure is retried by restarting the process, not by re-running the CLI. +_STARTUP_ACTION_LINE = ( + "Action: restart is safe — migrations are idempotent and self-heal on retry; " + "if the failure repeats, run `taskq migrate up` and report the output." +) + + +def _tx_diagnosis() -> migrate.ApplyFailureDiagnosis: + return migrate.ApplyFailureDiagnosis( + headline="deadlock detected", + failed_filename="01.00.00_01_pre_failing.sql", + use_transaction=True, + invalid_indexes=(), + schema="taskq", + ) + + +def _nt_diagnosis() -> migrate.ApplyFailureDiagnosis: + return migrate.ApplyFailureDiagnosis( + headline="canceling statement due to statement timeout", + failed_filename="01.00.02_01_post_concurrent_idx.sql", + use_transaction=False, + invalid_indexes=("jobs_queue_idx",), + schema="taskq", + ) + + +def _generic_diagnosis() -> migrate.ApplyFailureDiagnosis: + return migrate.ApplyFailureDiagnosis( + headline="boom", + failed_filename=None, + use_transaction=None, + invalid_indexes=(), + schema="taskq", + ) + + +# Expected line lists mirror tests/test_cli_migrate.py verbatim so CLI +# byte-identity is pinned at the renderer level. +_TX_DEFAULT_LINES = [ + "migration 01.00.00_01_pre_failing.sql failed: deadlock detected", + "It ran in a transaction and rolled back: nothing from the migration was applied.", + "Action: fix the error and re-run `taskq migrate up`.", +] +_NT_DEFAULT_LINES = [ + "migration 01.00.02_01_post_concurrent_idx.sql failed: canceling statement due to statement timeout", + "It ran WITHOUT a transaction (-- taskq:no-transaction): statements " + "before the failure remain applied, and the migration was NOT recorded " + "in the ledger.", + 'INVALID index(es) in schema "taskq": jobs_queue_idx — an interrupted ' + "CREATE INDEX CONCURRENTLY left them behind.", + "Action: re-run `taskq migrate up` — the migration is idempotent and " + "drops/rebuilds the debris itself.", +] +_GENERIC_DEFAULT_LINES = [ + "migration failed: boom", + "Action: fix the error and re-run `taskq migrate up` — already-applied migrations are skipped.", +] + + +@pytest.mark.parametrize( + ("make", "startup", "expected"), + [ + pytest.param(_tx_diagnosis, False, _TX_DEFAULT_LINES, id="transactional-default"), + pytest.param( + _tx_diagnosis, + True, + [*_TX_DEFAULT_LINES[:-1], _STARTUP_ACTION_LINE], + id="transactional-startup", + ), + pytest.param(_nt_diagnosis, False, _NT_DEFAULT_LINES, id="no-transaction-default"), + pytest.param( + _nt_diagnosis, + True, + [*_NT_DEFAULT_LINES[:-1], _STARTUP_ACTION_LINE], + id="no-transaction-startup", + ), + pytest.param(_generic_diagnosis, False, _GENERIC_DEFAULT_LINES, id="generic-default"), + pytest.param( + _generic_diagnosis, + True, + [*_GENERIC_DEFAULT_LINES[:-1], _STARTUP_ACTION_LINE], + id="generic-startup", + ), + ], +) +def test_render_apply_failure_lines( + make: Callable[[], migrate.ApplyFailureDiagnosis], + startup: bool, + expected: list[str], +) -> None: + """The renderer reproduces the CLI's three report variants verbatim and + swaps ONLY the action line for startup (worker/UI) failures.""" + assert migrate.render_apply_failure_lines(make(), startup=startup) == expected + + +def test_render_apply_failure_lines_omits_invalid_line_without_debris() -> None: + """A no-transaction failure with no INVALID-index debris omits the + INVALID line (mirrors the CLI's conditional).""" + d = migrate.ApplyFailureDiagnosis( + headline="connection was closed mid-build", + failed_filename="01.00.02_01_post_concurrent_idx.sql", + use_transaction=False, + invalid_indexes=(), + schema="taskq", + ) + lines = migrate.render_apply_failure_lines(d) + assert len(lines) == 3 + assert not any("INVALID" in line for line in lines) + + +def test_apply_failure_diagnosis_dataclass_is_frozen() -> None: + from dataclasses import FrozenInstanceError + + d = _generic_diagnosis() + with pytest.raises(FrozenInstanceError): + d.headline = "changed" # type: ignore[misc] + + +def test_apply_failure_diagnosis_rejects_filename_without_use_transaction() -> None: + """The renderer branches on use_transaction whenever failed_filename is + set, so the pair must be consistent: a filename with use_transaction=None + would silently render the no-transaction wording. Fail fast instead.""" + with pytest.raises(ValueError, match="use_transaction"): + migrate.ApplyFailureDiagnosis( + headline="boom", + failed_filename="01.00.00_01_pre_failing.sql", + use_transaction=None, + invalid_indexes=(), + schema="taskq", + ) + + +# ── apply-failure diagnosis attribution ───────────────────────────────────── + + +def _phase_scenario_migrations() -> tuple[Migration, Migration]: + """discover() order for the --phase misattribution scenario: an + earlier-version pending :post sorts BEFORE a later-version :pre (the + sort key is version first), so under ``migrate up --phase pre`` — where + only the :pre applies and fails — the first-unrecorded heuristic would + blame the wrong file.""" + earlier_pending_post = Migration( + version="01.00.00_01", + phase="post", + description="d", + filename="01.00.00_01_post_pending.sql", + sql_template="SELECT 1;", + ) + later_failing_pre = Migration( + version="01.00.02_01", + phase="pre", + description="d", + filename="01.00.02_01_pre_failing.sql", + sql_template="SELECT 1;", + ) + return earlier_pending_post, later_failing_pre + + +async def test_diagnose_apply_failure_prefers_tagged_migration( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """Under --phase the first unrecorded migration in discover() order is + NOT necessarily the one that failed. apply_pending tags the exception + with the failing migration; the diagnosis must trust the tag over the + heuristic.""" + earlier_pending_post, later_failing_pre = _phase_scenario_migrations() + monkeypatch.setattr(migrate, "discover", lambda: [earlier_pending_post, later_failing_pre]) + monkeypatch.setattr(migrate, "list_applied", AsyncMock(return_value=set())) + exc = RuntimeError("boom") + exc.__dict__["taskq_failed_migration"] = later_failing_pre + + d = await migrate.diagnose_apply_failure(object(), "taskq", exc) # type: ignore[arg-type] # Why: the tagged path never touches the conn (transactional tag), so a stand-in suffices. + + assert d.failed_filename == "01.00.02_01_pre_failing.sql" + assert d.use_transaction is True + + +async def test_diagnose_apply_failure_untagged_falls_back_to_first_unrecorded( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """An untagged exception (raised outside apply_pending's per-migration + loop — e.g. the ledger ensure) keeps the original heuristic: first + unrecorded in discover() order.""" + earlier_pending_post, later_failing_pre = _phase_scenario_migrations() + monkeypatch.setattr(migrate, "discover", lambda: [earlier_pending_post, later_failing_pre]) + monkeypatch.setattr(migrate, "list_applied", AsyncMock(return_value=set())) + + d = await migrate.diagnose_apply_failure(object(), "taskq", RuntimeError("boom")) # type: ignore[arg-type] # Why: list_applied is monkeypatched, so the conn stand-in is never used. + + assert d.failed_filename == "01.00.00_01_post_pending.sql" + assert d.use_transaction is True + + +async def test_diagnose_apply_failure_tagged_no_transaction_gathers_invalid_indexes( + monkeypatch: pytest.MonkeyPatch, +) -> None: + """The tagged path must still gather INVALID-index debris when the + tagged migration opted out of the transaction wrapper — an interrupted + CREATE INDEX CONCURRENTLY is exactly the failure the tag exists for.""" + tagged = Migration( + version="01.00.02_01", + phase="post", + description="d", + filename="01.00.02_01_post_concurrent_idx.sql", + sql_template="SELECT 1;", + use_transaction=False, + ) + monkeypatch.setattr(migrate, "list_invalid_indexes", AsyncMock(return_value=["jobs_queue_idx"])) + exc = RuntimeError("boom") + exc.__dict__["taskq_failed_migration"] = tagged + + d = await migrate.diagnose_apply_failure(object(), "taskq", exc) # type: ignore[arg-type] # Why: list_invalid_indexes is monkeypatched, so the conn stand-in is never used. + + assert d.failed_filename == "01.00.02_01_post_concurrent_idx.sql" + assert d.use_transaction is False + assert d.invalid_indexes == ("jobs_queue_idx",) + + +async def test_diagnose_apply_failure_headline_falls_back_to_type_name() -> None: + """An exception whose first str() line is empty/whitespace must not + render ``migration failed: `` with an empty headline — use the + exception's type name instead.""" + d = await migrate.diagnose_apply_failure( # type: ignore[arg-type] # Why: the generic path suppresses every read, so a conn stand-in is fine. + object(), "taskq", RuntimeError(" \nDETAIL: something") + ) + assert d.headline == "RuntimeError"