Migrations

Forward-only SQL migration runner.

migrate

Migration runner.

Forward-only by design. The runner:

1. Discovers *.sql files under :mod:taskq.migrations in lexicographic order, the naming convention is {ver}_{nn}_{pre|post}_{description}.sql.
2. Substitutes the {schema} placeholder with the configured schema name.
3. Applies migrations not already recorded in {schema}.schema_migrations, 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.

Published variants

An applied migration's ledger checksum must match the bundled file or the runner refuses (checksum drift, fail-closed). The one exception: a file whose bytes were amended after first shipping may exist in more than one published version, and databases that applied the earlier bytes keep that checksum forever. The _variants/ directory inside :mod:taskq.migrations carries the published history (one file per vintage, named for the commit that published it); the drift check accepts a ledger checksum matching the current file OR any bundled variant, each rendered with the checking database's own schema. A checksum nothing bundled produced is still drift. Variants are code: reviewed, bundled, versioned — adding one is a PR, never a runtime knob.

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.

__all__ module-attribute

__all__ = [
    "ApplyFailureDiagnosis",
    "ChecksumDrift",
    "ChecksumDriftError",
    "Migration",
    "MigrationLockTimeoutError",
    "Phase",
    "apply_pending",
    "apply_pending_locked",
    "checksum_drifts",
    "diagnose_apply_failure",
    "discover",
    "list_applied",
    "list_invalid_indexes",
    "migration_advisory_lock",
    "migration_lock_name",
    "render",
    "render_apply_failure_lines",
    "split_statements",
]

logger module-attribute

logger = structlog.get_logger('taskq.migrate')

Phase module-attribute

Phase: TypeAlias = Literal['pre', 'post']

DEFAULT_MIGRATION_DDL_LOCK_TIMEOUT module-attribute

DEFAULT_MIGRATION_DDL_LOCK_TIMEOUT: float = 30.0

DEFAULT_MIGRATION_LOCK_TIMEOUT module-attribute

DEFAULT_MIGRATION_LOCK_TIMEOUT: float = 120.0

Migration dataclass

Migration(
    version: str,
    phase: Phase,
    description: str,
    filename: str,
    sql_template: str,
    use_transaction: bool = True,
)

A single SQL migration file.

version instance-attribute

version: str

{ver}_{nn}, e.g. 01.00.00_01.

phase instance-attribute

phase: Phase

description instance-attribute

description: str

filename instance-attribute

filename: str

sql_template instance-attribute

sql_template: str

use_transaction class-attribute instance-attribute

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.

key property

key: str

Identity stored in schema_migrations.version: {version}:{phase}.

render

render(schema: str) -> str

Source code in src/taskq/migrate.py

def render(self, schema: str) -> str:
    return render(self.sql_template, schema)

checksum

checksum(schema: str) -> str

Source code in src/taskq/migrate.py

def checksum(self, schema: str) -> str:
    return hashlib.sha256(self.render(schema).encode("utf-8")).hexdigest()

MigrationLockTimeoutError

MigrationLockTimeoutError(
    migration: Migration | None, timeout: float
)

Bases: RuntimeError

A migration run's DDL waited ddl_lock_timeout for a table lock and gave up; its transaction rolled back and nothing was applied.

migration is the transactional migration whose DDL timed out, or None when the wait was the runner's own upgrade of the schema_migrations ledger (an ALTER TABLE that runs before the first migration of a run, or after a -- taskq:no-transaction file's statements). Raised in place of the driver's LockNotAvailableError (kept as the cause) so the report names what waited, the bound that expired, and the remedy, the holder must be found and ended, or the bound raised, rather than the bare canceling statement due to lock timeout.

Source code in src/taskq/migrate.py

def __init__(self, migration: Migration | None, timeout: float) -> None:
    self.migration = migration
    self.timeout = timeout
    if migration is not None:
        subject = f"migration {migration.filename}"
        what = "a lock on a table it alters"
        outcome = "The migration rolled back and nothing was applied."
    else:
        subject = "the schema_migrations ledger upgrade"
        what = "a lock on the ledger table"
        outcome = "It rolled back; no migration was applied or recorded."
    super().__init__(
        f"{subject} waited {timeout}s for {what} and gave up (ddl_lock_timeout): "
        "another session holds a conflicting lock on it, a long-running or "
        "idle-in-transaction connection, pg_dump, or an admin query. "
        f"{outcome} Find the holder in pg_stat_activity / pg_locks and end it or "
        "wait for it, then re-run; to wait longer, raise ddl_lock_timeout "
        "(apply_pending / apply_pending_locked), 0 waits indefinitely, at the cost "
        "of parking every statement on the table behind the queued DDL."
    )

migration instance-attribute

migration = migration

timeout instance-attribute

timeout = timeout

ChecksumDrift dataclass

ChecksumDrift(key: str, stored: str, current: str)

One applied migration whose ledger checksum differs from the file.

key is the ledger identity ({version}:{phase}), stored is the checksum recorded when the migration ran, current is the checksum the bundled file renders to now.

key instance-attribute

key: str

stored instance-attribute

stored: str

current instance-attribute

current: str

ChecksumDriftError

ChecksumDriftError(drifts: list[ChecksumDrift])

Bases: RuntimeError

An APPLIED migration's ledger checksum differs from the bundled file.

The ledger checksum is the SHA-256 of the SQL as it ran; a mismatch means the file no longer matches the schema the database actually has. Applying further migrations on top would build on an unverified premise (the recorded schema may not be what the current files produce, so the next file's assumptions can be silently wrong), which is why the runner refuses by default instead of warning and continuing. NOT-YET-APPLIED files cannot drift: nothing is recorded for them yet.

Remedies: restore the migration file to the version that actually ran (editing an already-applied migration is never picked up, forward-only by design), or, after verifying the edited file is exactly what the database already has (or that the edit is safe to build on), pass allow_checksum_drift=True (the CLI's --allow-checksum-drift) to proceed past the refusal. The drift is still logged as a warning either way, the ledger keeps the stored checksum.

Source code in src/taskq/migrate.py

def __init__(self, drifts: list[ChecksumDrift]) -> None:
    self.drifts = drifts
    detail = "; ".join(
        f"{d.key} stored {d.stored[:12]} vs file {d.current[:12]}" for d in drifts
    )
    super().__init__(
        f"{len(drifts)} applied migration(s) no longer match the bundled "
        f"files (checksum drift): {detail}. The database schema was built "
        "from different SQL than the files now contain, so refusing to "
        "apply anything on top of it. Restore the migration file(s) to "
        "what actually ran, or verify the edit is safe and re-run with "
        "--allow-checksum-drift (apply_pending / apply_pending_locked's "
        "allow_checksum_drift) to proceed past the refusal."
    )

drifts instance-attribute

drifts = drifts

ApplyFailureDiagnosis dataclass

ApplyFailureDiagnosis(
    headline: str,
    failed_filename: str | None,
    use_transaction: bool | None,
    invalid_indexes: tuple[str, ...],
    schema: str,
    guard_rejected: bool = False,
)

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 instance-attribute

headline: str

First line of the original error (asyncpg messages can be multiline).

failed_filename instance-attribute

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 instance-attribute

use_transaction: bool | None

Migration.use_transaction of the failed migration; None when the failed migration could not be determined.

invalid_indexes instance-attribute

invalid_indexes: tuple[str, ...]

INVALID indexes in the schema, debris of an interrupted CREATE INDEX CONCURRENTLY; only gathered for no-transaction failures.

schema instance-attribute

schema: str

Schema the report names in the INVALID-index line.

guard_rejected class-attribute instance-attribute

guard_rejected: bool = False

True when the failure was the transaction-control guard refusing the file (:func:_reject_transaction_control) rather than a statement failing mid-file, read from the exception's taskq_guard_rejection tag. The report then says nothing was executed and that the fix is removing the statement, instead of the generic no-transaction wording ("statements before the failure remain applied", "re-run, the migration is idempotent"), both false for a guard rejection.

__post_init__

__post_init__() -> None

Source code in src/taskq/migrate.py

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"
        )

discover

discover() -> list[Migration]

Return all bundled migrations sorted by version, then pre before post.

Raises :class:ValueError when two bundled files map to one ledger identity (see :func:_refuse_duplicate_ledger_keys).

Source code in src/taskq/migrate.py

def discover() -> list[Migration]:
    """Return all bundled migrations sorted by version, then ``pre`` before ``post``.

    Raises :class:`ValueError` when two bundled files map to one ledger
    identity (see :func:`_refuse_duplicate_ledger_keys`).
    """
    found: list[Migration] = []
    package = resources.files("taskq.migrations")
    for entry in package.iterdir():
        if not entry.is_file() or not entry.name.endswith(".sql"):
            continue
        match = _NAME_RE.match(entry.name)
        if match is None:
            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=sql_template,
                use_transaction=_uses_transaction(sql_template, entry.name),
            )
        )
    found.sort(key=lambda m: (m.version, 0 if m.phase == "pre" else 1))
    _refuse_duplicate_ledger_keys(found)
    return found

render

render(template: str, schema: str) -> str

Substitute {schema} in a SQL template.

SQL files escape literal curly braces by doubling them ({{ → {) because :func:str.format is the substitution engine.

Source code in src/taskq/migrate.py

def render(template: str, schema: str) -> str:
    """Substitute ``{schema}`` in a SQL template.

    SQL files escape literal curly braces by doubling them (``{{`` → ``{``)
    because :func:`str.format` is the substitution engine.
    """
    if not _IDENT_RE.match(schema):
        raise ValueError(f"invalid schema name {schema!r}")
    return template.format(schema=schema)

split_statements

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.

Source code in src/taskq/migrate.py

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

list_applied async

list_applied(conn: Connection, schema: str) -> set[str]

Return {version}:{phase} keys recorded in schema_migrations.

Returns an empty set if the schema or table does not yet exist, a fresh database is the common case on first migrate up.

Source code in src/taskq/migrate.py

async def list_applied(conn: asyncpg.Connection, schema: str) -> set[str]:
    """Return ``{version}:{phase}`` keys recorded in ``schema_migrations``.

    Returns an empty set if the schema or table does not yet exist, a
    fresh database is the common case on first ``migrate up``.
    """
    if not _IDENT_RE.match(schema):
        raise ValueError(f"invalid schema name {schema!r}")
    exists = await conn.fetchval(
        """
        SELECT EXISTS (
            SELECT 1 FROM information_schema.tables
            WHERE table_schema = $1 AND table_name = 'schema_migrations'
        )
        """,
        schema,
    )
    if not exists:
        return set()
    rows = await conn.fetch(f'SELECT version, checksum FROM "{schema}".schema_migrations')
    applied_keys: set[str] = set()
    for r in rows:
        applied_keys.add(r["version"])
    return applied_keys

checksum_drifts async

checksum_drifts(
    conn: Connection, *, schema: str
) -> dict[str, ChecksumDrift]

Compare the ledger's recorded checksums against the bundled files.

Returns one :class:ChecksumDrift per APPLIED migration whose stored checksum differs from what its file renders to now, keyed by the migration key, empty when the ledger and the files agree (or when no ledger exists yet). Shared by migrate up's fail-closed refusal (:func:apply_pending) and migrate status's drift report, so both surfaces answer with the same comparison.

A drifted checksum means the file no longer matches the SQL that built the schema: the file was edited after it ran, or the ledger was tampered with. Neither is recoverable by applying more migrations, which is why callers should surface the result prominently rather than bury it among pending entries.

Source code in src/taskq/migrate.py

async def checksum_drifts(conn: asyncpg.Connection, *, schema: str) -> dict[str, ChecksumDrift]:
    """Compare the ledger's recorded checksums against the bundled files.

    Returns one :class:`ChecksumDrift` per APPLIED migration whose stored
    checksum differs from what its file renders to now, keyed by the
    migration key, empty when the ledger and the files agree (or when no
    ledger exists yet). Shared by ``migrate up``'s fail-closed refusal
    (:func:`apply_pending`) and ``migrate status``'s drift report, so both
    surfaces answer with the same comparison.

    A drifted checksum means the file no longer matches the SQL that
    built the schema: the file was edited after it ran, or the ledger was
    tampered with. Neither is recoverable by applying more migrations,
    which is why callers should surface the result prominently rather
    than bury it among pending entries.
    """
    if not _IDENT_RE.match(schema):
        raise ValueError(f"invalid schema name {schema!r}")
    exists = await conn.fetchval(
        """
        SELECT EXISTS (
            SELECT 1 FROM information_schema.tables
            WHERE table_schema = $1 AND table_name = 'schema_migrations'
        )
        """,
        schema,
    )
    if not exists:
        return {}
    rows = await conn.fetch(f'SELECT version, checksum FROM "{schema}".schema_migrations')
    applied_checksums: dict[str, str] = {r["version"]: r["checksum"] for r in rows}
    drift_list = _detect_checksum_drifts(discover(), applied_checksums, schema)
    return {d.key: d for d in drift_list}

list_invalid_indexes async

list_invalid_indexes(
    conn: 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.

Source code in src/taskq/migrate.py

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]

apply_pending async

apply_pending(
    conn: Connection,
    *,
    schema: str,
    phase: Phase | None = None,
    target: str | None = None,
    max_steps: int | None = None,
    ddl_lock_timeout: float = DEFAULT_MIGRATION_DDL_LOCK_TIMEOUT,
    allow_checksum_drift: bool = False,
) -> list[Migration]

Apply pending migrations.

Each migration runs in its own transaction so a failure in one file does 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). Must name a version in the discovered set: an unknown target, a deleted or renumbered migration, raises ValueError rather than silently applying every pending migration. :param max_steps: stop after this many applies. :param ddl_lock_timeout: seconds a transactional migration may wait for a table lock (:data:DEFAULT_MIGRATION_DDL_LOCK_TIMEOUT); applied as SET LOCAL lock_timeout inside each migration's own transaction, so it never outlives the migration, and never applied to -- taskq:no-transaction files, whose CONCURRENTLY waits are heavyweight-lock waits by design. A wait that outlives it raises :class:MigrationLockTimeoutError; 0 waits indefinitely (stated explicitly, so a session-level lock_timeout default cannot bound the wait either); negative, non-finite, and sub-millisecond values are refused with ValueError. :param allow_checksum_drift: proceed past :class:ChecksumDriftError when an APPLIED migration's ledger checksum differs from the bundled file. The drift is still logged as a warning. NOT-YET- APPLIED files cannot drift (nothing is recorded for them yet), so this flag never changes which migrations run, only whether a drifted ledger refuses the run. :returns: migrations that were applied (in order). :raises ChecksumDriftError: when an applied migration's stored checksum differs from the file and allow_checksum_drift is False. Raised BEFORE anything is applied: a ledger that no longer matches the files means the schema's provenance is unverified, and further migrations must not build on it.

Source code in src/taskq/migrate.py

async def apply_pending(
    conn: asyncpg.Connection,
    *,
    schema: str,
    phase: Phase | None = None,
    target: str | None = None,
    max_steps: int | None = None,
    ddl_lock_timeout: float = DEFAULT_MIGRATION_DDL_LOCK_TIMEOUT,
    allow_checksum_drift: bool = False,
) -> list[Migration]:
    """Apply pending migrations.

    Each migration runs in its own transaction so a failure in one file does
    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). Must
        name a version in the discovered set: an unknown target, a
        deleted or renumbered migration, raises ``ValueError`` rather
        than silently applying every pending migration.
    :param max_steps: stop after this many applies.
    :param ddl_lock_timeout: seconds a transactional migration may wait for
        a table lock (:data:`DEFAULT_MIGRATION_DDL_LOCK_TIMEOUT`); applied
        as ``SET LOCAL lock_timeout`` inside each migration's own
        transaction, so it never outlives the migration, and never applied
        to ``-- taskq:no-transaction`` files, whose ``CONCURRENTLY`` waits
        are heavyweight-lock waits by design. A wait that outlives it
        raises :class:`MigrationLockTimeoutError`; ``0`` waits indefinitely
        (stated explicitly, so a session-level ``lock_timeout`` default
        cannot bound the wait either); negative, non-finite, and
        sub-millisecond values are refused with ``ValueError``.
    :param allow_checksum_drift: proceed past :class:`ChecksumDriftError`
        when an APPLIED migration's ledger checksum differs from the
        bundled file. The drift is still logged as a warning. NOT-YET-
        APPLIED files cannot drift (nothing is recorded for them yet), so
        this flag never changes which migrations run, only whether a
        drifted ledger refuses the run.
    :returns: migrations that were applied (in order).
    :raises ChecksumDriftError: when an applied migration's stored checksum
        differs from the file and ``allow_checksum_drift`` is False.
        Raised BEFORE anything is applied: a ledger that no longer matches
        the files means the schema's provenance is unverified, and further
        migrations must not build on it.
    """
    _validate_timeout_bound(ddl_lock_timeout, "ddl_lock_timeout")
    if not _IDENT_RE.match(schema):
        raise ValueError(f"invalid schema name {schema!r}")

    exists = await conn.fetchval(
        """
        SELECT EXISTS (
            SELECT 1 FROM information_schema.tables
            WHERE table_schema = $1 AND table_name = 'schema_migrations'
        )
        """,
        schema,
    )
    if exists:
        applied_rows = await conn.fetch(
            f'SELECT version, checksum FROM "{schema}".schema_migrations'
        )
        applied_keys: set[str] = {r["version"] for r in applied_rows}
        applied_checksums: dict[str, str] = {r["version"]: r["checksum"] for r in applied_rows}
    else:
        applied_keys = set()
        applied_checksums: dict[str, str] = {}

    all_migrations = discover()

    # A target the runner no longer bundles must fail loudly. The
    # truncation loop below only stops when it SEES the target, so an
    # unknown one falls through and silently applies EVERY pending
    # migration: the exact inversion of what a caller passing a target
    # expects, and reachable the moment any release deletes or renumbers
    # files (a deleted 01.00.11_01 or renumbered 01.00.12_05:post is a
    # target no longer bundled). Refuse before anything is applied.
    if target is not None and target not in {m.version for m in all_migrations}:
        range_hint = (
            f"discovered versions run {all_migrations[0].version} through "
            f"{all_migrations[-1].version}"
            if all_migrations
            else "no migrations are bundled"
        )
        raise ValueError(
            f"unknown migration target {target!r}: not in the bundled migration "
            f"set ({range_hint}). A target the runner no longer bundles is a "
            "deleted or renumbered migration; refusing instead of silently "
            "applying every pending migration. List the discovered versions "
            "with `taskq migrate status`."
        )

    drifts = _detect_checksum_drifts(all_migrations, applied_checksums, schema)
    for d in drifts:
        logger.warning(
            "migration-checksum-drift",
            key=d.key,
            stored_checksum=d.stored,
            current_checksum=d.current,
        )
    # Fail-closed: a drifted ledger means the database was built from SQL
    # the files no longer contain, so this run's premise (the files
    # describe the schema) is unverified. The refusal happens BEFORE the
    # apply loop, so a refused run applies nothing. Allowed through, the
    # warnings above still land: the drift stays visible in the logs.
    if drifts and not allow_checksum_drift:
        raise ChecksumDriftError(drifts)

    pending = [m for m in all_migrations if m.key not in applied_keys]
    if phase is not None:
        pending = [m for m in pending if m.phase == phase]

    # Truncate to the prefix this run will actually apply (target is
    # inclusive; max_steps caps the count) BEFORE validating phase ordering,
    # so the guard below reasons about exactly this run's apply list and
    # cannot refuse over a migration the truncation would have skipped.
    effective: list[Migration] = []
    for migration in pending:
        effective.append(migration)
        if target is not None and migration.version == target:
            break
        if max_steps is not None and len(effective) >= max_steps:
            break

    # A post-phase migration is only eligible once its same-version pre-phase
    # counterpart is applied (or will be applied earlier in this same run).
    # Post-phase migrations remove the structures pre-phase migrations add ,
    # e.g. 01.00.03_01:post drops the old idempotency index the rolling-deploy
    # overlap window depends on. Applying one first (e.g. `migrate up --phase
    # post` against a schema whose pre phase hasn't run) both strands the pre
    # phase's protections permanently (the post is already recorded) and can
    # break not-yet-upgraded workers immediately. Refuse loudly instead.
    pre_versions: set[str] = {m.version for m in all_migrations if m.phase == "pre"}
    eligible_keys = set(applied_keys)
    for m in effective:
        if m.phase == "post" and m.version in pre_versions:
            pre_key = f"{m.version}:pre"
            if pre_key not in eligible_keys:
                exc = ValueError(
                    f"migration {m.key} cannot be applied before its pre-phase "
                    f"counterpart {pre_key}. Run `taskq migrate up --phase pre` "
                    f"(or a plain `taskq migrate up`) first."
                )
                # Why tag: this raises BEFORE the per-migration loop, so the
                # loop's except-clause tagging never runs and the diagnosis
                # would fall back to the first-unrecorded heuristic, which
                # blames an unrelated, earlier-version file under --phase.
                # Same mechanism the loop uses (tag, not wrap) so the
                # caller-visible exception type is unchanged.
                exc.__dict__["taskq_failed_migration"] = m
                raise exc
        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 _upgrade_ledger_bounded(conn, schema, ddl_lock_timeout)
        ledger_ready = True

    applied_now: list[Migration] = []
    for migration in effective:
        try:
            if migration.use_transaction:
                async with _lock_bounded_transaction(conn, ddl_lock_timeout):
                    try:
                        await conn.execute(migration.render(schema))
                    except asyncpg.LockNotAvailableError as exc:
                        raise MigrationLockTimeoutError(migration, ddl_lock_timeout) from exc
                    if not ledger_ready:
                        # Inside the migration's own transaction: the bound
                        # above already covers the ledger ALTER.
                        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 _upgrade_ledger_bounded(conn, schema, ddl_lock_timeout)
                    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

migration_lock_name

migration_lock_name(schema: str) -> str

Schema-qualified migration advisory-lock name: taskq:migrate:{schema}.

Advisory locks live in a per-database namespace, so the pre-qualification fixed bigint key was shared by every schema in the database: two deployments sharing one database serialized their migrations on one lock, and the loser of a cross-deployment race, a schema with nothing to migrate, burned its bounded startup wait and SystemExited while an unrelated schema's long DDL ran. Qualifying with the schema gives each schema its own lock; same-schema replicas still serialize (the reason the lock exists, a virgin schema's bare CREATE TABLE races two concurrent appliers).

Deliberately the same taskq:{purpose}:{schema} shape as :func:taskq.constants.schema_lock_name (built from it). That helper is how every purpose lock in the system derives its key, one lock per schema and purpose (maintenance_leader, prune, archive_expiry, cron, migrate). Keyed locks, which must distinguish a resource within a schema, embed the schema before their key instead: the sliding-window bucket lock taskq:{schema}:sw:{name} (mirroring the bucket's Redis key) and the enqueue unique-for lock taskq:unique_for:{schema}:{actor}:{identity_key}. Both conventions qualify with the schema; they differ only in where the distinguishing key segments sit.

Upgrade discipline: the qualified key replaces the unqualified bigint outright, a mixed old/new fleet holds different keys, so an old and a new replica can apply migrations to the SAME schema concurrently during the deploy window. Adopt by restarting the fleet onto the new release rather than rolling it; the window is the deploy, not the steady state. (Same caveat as the schema-qualified leader locks.)

Source code in src/taskq/migrate.py

def migration_lock_name(schema: str) -> str:
    """Schema-qualified migration advisory-lock name: ``taskq:migrate:{schema}``.

    Advisory locks live in a per-database namespace, so the pre-qualification
    fixed bigint key was shared by every schema in the database: two
    deployments sharing one database serialized their migrations on one
    lock, and the loser of a *cross-deployment* race, a schema with
    nothing to migrate, burned its bounded startup wait and SystemExited
    while an unrelated schema's long DDL ran. Qualifying with the schema
    gives each schema its own lock; same-schema replicas still serialize
    (the reason the lock exists, a virgin schema's bare ``CREATE TABLE``
    races two concurrent appliers).

    Deliberately the same ``taskq:{purpose}:{schema}`` shape as
    :func:`taskq.constants.schema_lock_name` (built from it). That helper
    is how every *purpose* lock in the system derives its key, one lock
    per schema and purpose (maintenance_leader, prune, archive_expiry,
    cron, migrate). *Keyed* locks, which must distinguish a resource
    within a schema, embed the schema before their key instead: the
    sliding-window bucket lock ``taskq:{schema}:sw:{name}`` (mirroring
    the bucket's Redis key) and the enqueue unique-for lock
    ``taskq:unique_for:{schema}:{actor}:{identity_key}``. Both
    conventions qualify with the schema; they differ only in where the
    distinguishing key segments sit.

    Upgrade discipline: the qualified key replaces the unqualified bigint
    outright, a mixed old/new fleet holds different keys, so an old and a
    new replica can apply migrations to the SAME schema concurrently
    during the deploy window. Adopt by restarting the fleet onto the new
    release rather than rolling it; the window is the deploy, not the
    steady state. (Same caveat as the schema-qualified leader locks.)
    """
    require_schema(schema)
    return schema_lock_name("migrate", schema)

diagnose_apply_failure async

diagnose_apply_failure(
    conn: 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).

Source code in src/taskq/migrate.py

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)

    # Guard rejections carry their own tag (set by
    # _reject_transaction_control, preserved by the loop's re-raise of the
    # same object); `is True` keeps an exotic/mocked attribute from
    # steering the report, mirroring the Migration tag guard above.
    guard_rejected = getattr(exc, "taskq_guard_rejection", False) is True

    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,
        guard_rejected=guard_rejected,
    )

render_apply_failure_lines

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), except for guard rejections, where neither action line is truthful (nothing executed, and re-run or restart fails identically until the statement is removed), so the guard branch renders the same remove-and-rerun guidance either way.

Source code in src/taskq/migrate.py

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`), except for guard rejections,
    where neither action line is truthful (nothing executed, and re-run or
    restart fails identically until the statement is removed), so the
    guard branch renders the same remove-and-rerun guidance either way.
    """
    if d.failed_filename is not None and d.guard_rejected:
        return [
            f"migration {d.failed_filename} failed: {d.headline}",
            "Nothing was executed, the file was rejected before any "
            "statement ran (transaction-control statements are forbidden "
            "in a -- taskq:no-transaction migration).",
            "Action: remove the transaction-control statement and re-run `taskq migrate up`.",
        ]

    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

migration_advisory_lock async

migration_advisory_lock(
    conn: Connection,
    lock_timeout: float = DEFAULT_MIGRATION_LOCK_TIMEOUT,
    *,
    schema: str,
) -> AsyncGenerator[None]

Hold schema's migration advisory lock on conn, with a bounded wait.

Extracted so the CLI and :func:apply_pending_locked serialize on the SAME lock without duplicating the acquire/reset/release protocol. The CLI cannot delegate to apply_pending_locked: it owns the connection so it can run _report_up_failure diagnostics on it after a failure, and apply_pending_locked converts failures to SystemExit before that could run.

The lock key is :func:migration_lock_name, schema-qualified, because advisory locks are database-scoped and an unqualified key serializes every schema in the database (see that function's docstring for the cross-deployment failure mode and the mixed-fleet adopt note). schema is required, not defaulted: a silent default would point at one fixed schema while the caller's settings name another, the same defect class the qualification itself fixes.

lock_timeout bounds only the WAIT for the advisory lock, via Postgres' lock_timeout GUC. It is reset to unlimited before the body runs: the session-wide value must not bound the -- taskq:no-transaction files, whose CONCURRENTLY phases wait on heavyweight locks by design, while each transactional migration bounds its own table-lock wait with SET LOCAL (see :func:apply_pending's ddl_lock_timeout). Neither bound can interrupt a statement that already holds its lock, a long index build is governed by statement_timeout, widened below. 0 waits indefinitely (stated explicitly, so a session-level lock_timeout default cannot bound the acquire either); negative, non-finite, and sub-millisecond values are refused with ValueError, the same rule ddl_lock_timeout follows. If the reset fails (a wedged caller-owned connection), the failure is logged as a warning naming the connection and the lock flow proceeds, that connection keeps the wait bound as its session-wide lock_timeout until it is closed.

statement_timeout is widened to unlimited immediately after the lock is acquired, once, for the whole apply phase, and NEVER on the contention path: a caller-owned connection that merely lost the race gets SystemExit above and returns to its owner, which must not silently gain a session-wide unlimited statement_timeout from losing a lock race. A session statement_timeout (server_settings on the connection, a role/database default, a DSN options clause) would abort a long DDL statement, an index build over a large table , mid-statement, and identically on every retry, since the runner re-executes the same migration after a failure. Like the lock_timeout doctrine, the widened state is deliberately NOT restored on exit: the unlock must not trade a GUC restore for a failed migration run, and a failed widening is logged as a warning naming the connection while the apply proceeds, that connection keeps its session statement_timeout.

Raises :class:SystemExit on contention rather than blocking until the container platform kills the process.

Source code in src/taskq/migrate.py

@contextlib.asynccontextmanager
async def migration_advisory_lock(
    conn: asyncpg.Connection,
    lock_timeout: float = DEFAULT_MIGRATION_LOCK_TIMEOUT,
    *,
    schema: str,
) -> AsyncGenerator[None]:
    """Hold *schema*'s migration advisory lock on *conn*, with a bounded wait.

    Extracted so the CLI and :func:`apply_pending_locked` serialize on the SAME
    lock without duplicating the acquire/reset/release protocol. The CLI cannot
    delegate to ``apply_pending_locked``: it owns the connection so it
    can run ``_report_up_failure`` diagnostics on it after a failure, and
    ``apply_pending_locked`` converts failures to ``SystemExit`` before that
    could run.

    The lock key is :func:`migration_lock_name`, schema-qualified, because
    advisory locks are database-scoped and an unqualified key serializes
    every schema in the database (see that function's docstring for the
    cross-deployment failure mode and the mixed-fleet adopt note).
    ``schema`` is required, not defaulted: a silent default would point at
    one fixed schema while the caller's settings name another, the same
    defect class the qualification itself fixes.

    ``lock_timeout`` bounds only the WAIT for the advisory lock, via
    Postgres' ``lock_timeout`` GUC. It is reset to unlimited before the body
    runs: the session-wide value must not bound the ``-- taskq:no-transaction``
    files, whose ``CONCURRENTLY`` phases wait on heavyweight locks by design,
    while each transactional migration bounds its own table-lock wait with
    ``SET LOCAL`` (see :func:`apply_pending`'s ``ddl_lock_timeout``). Neither
    bound can interrupt a statement that already holds its lock, a long
    index build is governed by ``statement_timeout``, widened below. ``0``
    waits indefinitely (stated explicitly, so a session-level
    ``lock_timeout`` default cannot bound the acquire either); negative,
    non-finite, and sub-millisecond values are refused with ``ValueError``,
    the same rule ``ddl_lock_timeout`` follows. If the reset fails (a
    wedged caller-owned connection), the failure is logged as a warning
    naming the connection and the lock flow proceeds, that connection keeps
    the wait bound as its session-wide ``lock_timeout`` until it is closed.

    ``statement_timeout`` is widened to unlimited immediately after the
    lock is acquired, once, for the whole apply phase, and NEVER on the
    contention path: a caller-owned connection that merely lost the race
    gets ``SystemExit`` above and returns to its owner, which must not
    silently gain a session-wide unlimited statement_timeout from losing
    a lock race. A session statement_timeout (server_settings on the
    connection, a role/database default, a DSN options clause) would
    abort a long DDL statement, an index build over a large table ,
    mid-statement, and identically on every retry, since the runner
    re-executes the same migration after a failure. Like the
    ``lock_timeout`` doctrine, the widened state is deliberately NOT
    restored on exit: the unlock must not trade a GUC restore for a
    failed migration run, and a failed widening is logged as a warning
    naming the connection while the apply proceeds, that connection
    keeps its session ``statement_timeout``.

    Raises :class:`SystemExit` on contention rather than blocking until the
    container platform kills the process.
    """
    lock_name = migration_lock_name(schema)
    _validate_timeout_bound(lock_timeout, "lock_timeout")
    if lock_timeout > 0:
        # Milliseconds; applies to the advisory-lock acquire below.
        await conn.execute(f"SET lock_timeout = {int(lock_timeout * 1000)}")
    else:
        # 0 promises an indefinite wait; a session whose role default, DSN
        # options clause or server_settings bounds lock_timeout would bound
        # the acquire after all unless the unlimited value is stated
        # explicitly. The bounded path's reset below leaves the session at
        # 0 too, so both paths end with the same session state.
        await conn.execute("SET lock_timeout = 0")
    try:
        await conn.execute(_MIGRATION_LOCK_SQL, lock_name)
    except asyncpg.LockNotAvailableError as exc:
        msg = (
            f"could not acquire the migration advisory lock within {lock_timeout}s: "
            "another process is applying migrations (or is wedged holding the lock). "
            "Migrations must run once, from a single pre-deploy job or init "
            "container -- not from every replica."
        )
        raise SystemExit(msg) from exc
    finally:
        # Reset before the DDL: the transactional files bound their own
        # table-lock waits with SET LOCAL, and the no-transaction files'
        # CONCURRENTLY waits must stay unbounded.
        if lock_timeout > 0:
            try:
                await conn.execute("SET lock_timeout = 0")
            except Exception as exc:
                # Why warn, not raise: the lock is held and the migrations
                # still need to run, so aborting here would trade a stale
                # GUC for a failed migration run. But the swallowed failure
                # must be visible, a caller-owned connection keeps the
                # wait bound as its session-wide lock_timeout, so later
                # deliberate long lock waits abort at the wait bound.
                # repr(conn) is the only name an anonymous connection has.
                logger.warning(
                    "migration-lock-timeout-reset-failed",
                    conn=repr(conn),
                    lock_timeout_ms=int(lock_timeout * 1000),
                    error=repr(exc),
                )
    # Only a session that WON the lock reaches here. Widen
    # statement_timeout for the whole apply phase, AFTER the acquire, so
    # the contention exit above cannot widen a caller-owned connection it
    # is about to hand back. Unconditional because this function never
    # sets a statement_timeout itself, the bound being widened was
    # supplied by the session (server_settings, a role/database default,
    # a DSN options clause), and it would abort a multi-minute index
    # build identically on every deploy and every retry.
    try:
        await conn.execute("SET statement_timeout = 0")
    except Exception as exc:
        # Why warn, not raise: the lock is held and the migrations
        # still need to run, so aborting here would trade a widened
        # timeout for a failed migration run. But the swallowed failure
        # must be visible, a caller-owned connection keeps its session
        # statement_timeout, so a long DDL step on it can still be
        # aborted mid-statement. repr(conn) is the only name an
        # anonymous connection has.
        logger.warning(
            "migration-statement-timeout-widen-failed",
            conn=repr(conn),
            error=repr(exc),
        )
    try:
        yield
    finally:
        # Why bounded: contextlib.suppress catches errors but cannot stop a
        # call that never returns; a dead PG wedges the unlock indefinitely.
        with contextlib.suppress(Exception):
            await asyncio.wait_for(
                conn.execute(_MIGRATION_UNLOCK_SQL, lock_name),
                timeout=CLOSE_TIMEOUT_SECS,
            )

apply_pending_locked async

apply_pending_locked(
    dsn: str | None = None,
    *,
    schema: str,
    phase: Phase | None = "pre",
    target: str | None = None,
    max_steps: int | None = None,
    conn: Connection | None = None,
    conn_factory: Callable[[], Awaitable[Connection]]
    | None = None,
    lock_timeout: float = DEFAULT_MIGRATION_LOCK_TIMEOUT,
    ddl_lock_timeout: float = DEFAULT_MIGRATION_DDL_LOCK_TIMEOUT,
    allow_checksum_drift: bool = False,
) -> list[Migration]

Apply pending migrations under a session-level advisory lock.

Acquires pg_advisory_lock to prevent concurrent startup races, applies pending migrations, and releases the lock.

lock_timeout bounds only the wait for the advisory lock, via Postgres' lock_timeout GUC. It is reset to unlimited before the migrations run; each transactional migration then bounds its own table-lock waits with ddl_lock_timeout (see :func:apply_pending), which cannot interrupt a statement that already holds its lock. Pass 0 to wait indefinitely (the old behaviour). The session's statement_timeout is widened to unlimited for the apply phase for the same reason, a caller-supplied session default must not abort a long DDL statement mid-build (it would fail identically on every retry). Losing the race raises :class:SystemExit naming the contention, rather than hanging until the platform kills the container.

allow_checksum_drift passes through to :func:apply_pending: a ledger checksum that no longer matches the bundled files refuses the run (as :class:ChecksumDriftError, converted to SystemExit by this wrapper's failure handling) unless the flag is set. Startup callers (--migrate, TASKQ_MIGRATE_ON_START) should leave it False: a drifted ledger is exactly the state a fleet must not start on.

Connection sources (mutually exclusive): * conn, pre-constructed, caller-owned; NOT closed here. * conn_factory, zero-arg async factory; closed in finally. * dsn, asyncpg.connect(dsn); closed in finally.

Raises :class:SystemExit on failure so the calling process aborts cleanly. This is the recommended entry point for CLI --migrate and admin sidecar TASKQ_MIGRATE_ON_START paths.

phase defaults to "pre" because this entry point fires on process lifecycle events nobody sequences, a pod restart, a rollout, an autoscale event, not on an operator's decision. A post-phase migration exists precisely to be withheld until the whole fleet is confirmed upgraded, and no single process can know that, so applying one here would let a restart close a rolling-deploy overlap window mid-rollout. Dropping the old single-column idempotency index while the not-yet-upgraded half of the fleet is still issuing ON CONFLICT (idempotency_key) takes that half's ENTIRE enqueue path down: the arbiter index is resolved at plan time, so every enqueue fails regardless of row values. The post phase stays behind the operator's explicit taskq migrate up --phase post; pass phase=None to apply every phase from a context that has that knowledge.

Source code in src/taskq/migrate.py

async def apply_pending_locked(
    dsn: str | None = None,
    *,
    schema: str,
    phase: Phase | None = "pre",
    target: str | None = None,
    max_steps: int | None = None,
    conn: asyncpg.Connection | None = None,
    conn_factory: Callable[[], Awaitable[asyncpg.Connection]] | None = None,
    lock_timeout: float = DEFAULT_MIGRATION_LOCK_TIMEOUT,
    ddl_lock_timeout: float = DEFAULT_MIGRATION_DDL_LOCK_TIMEOUT,
    allow_checksum_drift: bool = False,
) -> list[Migration]:
    """Apply pending migrations under a session-level advisory lock.

    Acquires ``pg_advisory_lock`` to prevent concurrent startup races,
    applies pending migrations, and releases the lock.

    ``lock_timeout`` bounds only the **wait** for the advisory lock, via
    Postgres' ``lock_timeout`` GUC. It is reset to unlimited before the
    migrations run; each transactional migration then bounds its own
    table-lock waits with ``ddl_lock_timeout`` (see :func:`apply_pending`),
    which cannot interrupt a statement that already holds its lock. Pass
    ``0`` to wait indefinitely (the old behaviour).
    The session's ``statement_timeout`` is widened to unlimited for the
    apply phase for the same reason, a caller-supplied session default
    must not abort a long DDL statement mid-build (it would fail
    identically on every retry).
    Losing the race raises :class:`SystemExit` naming the contention, rather
    than hanging until the platform kills the container.

    ``allow_checksum_drift`` passes through to :func:`apply_pending`: a
    ledger checksum that no longer matches the bundled files refuses the
    run (as :class:`ChecksumDriftError`, converted to ``SystemExit`` by
    this wrapper's failure handling) unless the flag is set. Startup
    callers (``--migrate``, ``TASKQ_MIGRATE_ON_START``) should leave it
    False: a drifted ledger is exactly the state a fleet must not start
    on.

    Connection sources (mutually exclusive):
    * ``conn``, pre-constructed, caller-owned; NOT closed here.
    * ``conn_factory``, zero-arg async factory; closed in ``finally``.
    * ``dsn``, ``asyncpg.connect(dsn)``; closed in ``finally``.

    Raises :class:`SystemExit` on failure so the calling process aborts
    cleanly.  This is the recommended entry point for CLI ``--migrate``
    and admin sidecar ``TASKQ_MIGRATE_ON_START`` paths.

    ``phase`` defaults to ``"pre"`` because this entry point fires on
    process lifecycle events nobody sequences, a pod restart, a rollout,
    an autoscale event, not on an operator's decision. A post-phase
    migration exists precisely to be withheld until the whole fleet is
    confirmed upgraded, and no single process can know that, so applying
    one here would let a restart close a rolling-deploy overlap window
    mid-rollout. Dropping the old single-column idempotency index while
    the not-yet-upgraded half of the fleet is still issuing
    ``ON CONFLICT (idempotency_key)`` takes that half's ENTIRE enqueue
    path down: the arbiter index is resolved at plan time, so every
    enqueue fails regardless of row values. The post phase stays behind
    the operator's explicit ``taskq migrate up --phase post``; pass
    ``phase=None`` to apply every phase from a context that has that
    knowledge.
    """
    if conn is not None and conn_factory is not None:
        raise ValueError("apply_pending_locked: provide 'conn' or 'conn_factory', not both")
    if dsn is not None and (conn is not None or conn_factory is not None):
        raise ValueError(
            "apply_pending_locked: 'dsn' is mutually exclusive with 'conn' and 'conn_factory'"
        )
    if conn is None and conn_factory is None and dsn is None:
        raise ValueError("apply_pending_locked: provide 'dsn', 'conn', or 'conn_factory'")

    owns_conn = conn is None  # factory/DSN → we close; caller-owned → we don't
    c: asyncpg.Connection | None = None
    try:
        if conn is not None:
            c = conn
        elif conn_factory is not None:
            c = await conn_factory()
        else:
            assert dsn is not None  # guarded by validation above
            c = await asyncpg.connect(dsn)
        async with migration_advisory_lock(c, lock_timeout, schema=schema):
            applied = await apply_pending(
                c,
                schema=schema,
                phase=phase,
                target=target,
                max_steps=max_steps,
                ddl_lock_timeout=ddl_lock_timeout,
                allow_checksum_drift=allow_checksum_drift,
            )
        if applied:
            logger.info("migrations-applied-before-startup", count=len(applied))
        else:
            logger.info("no-pending-migrations")
        return applied
    except SystemExit:
        # Lock-contention exit above is already precise; do not re-wrap it as
        # "migration failed", which would misreport a queueing problem as a
        # broken migration.
        raise
    except Exception as 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 and owns_conn:
            # The advisory unlock now belongs to migration_advisory_lock's own
            # finally (bounded there for the same reason), so only the owned
            # connection close remains here. Unlocking again would fire
            # pg_advisory_unlock on a lock this session no longer holds, which
            # Postgres answers with a WARNING and a false return.
            #
            # Why bounded: contextlib.suppress catches errors but cannot stop a
            # call that never returns, and this finally runs before any lifespan
            # exit stack exists, so an unbounded close would wedge CLI/UI
            # startup forever. close_conn_bounded terminates on timeout.
            await close_conn_bounded(c, "migrate", CLOSE_TIMEOUT_SECS)