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.

__all__ module-attribute

__all__ = [
    "Migration",
    "Phase",
    "apply_pending",
    "apply_pending_locked",
    "discover",
    "list_applied",
    "render",
]

logger module-attribute

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

Phase module-attribute

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

Migration dataclass

Migration(
    version: str,
    phase: Phase,
    description: str,
    filename: str,
    sql_template: str,
)

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

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

discover

discover() -> list[Migration]

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

Source code in src/taskq/migrate.py

def discover() -> list[Migration]:
    """Return all bundled migrations sorted by version, then ``pre`` before ``post``."""
    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.
        found.append(
            Migration(
                version=version,
                phase=phase,
                description=match.group("desc"),
                filename=entry.name,
                sql_template=entry.read_text(encoding="utf-8"),
            )
        )
    found.sort(key=lambda m: (m.version, 0 if m.phase == "pre" else 1))
    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)

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

apply_pending async

apply_pending(
    conn: Connection,
    *,
    schema: str,
    phase: Phase | None = None,
    target: str | None = None,
    max_steps: int | None = None,
) -> 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.

:param phase: restrict to pre or post migrations only. :param target: stop after applying this version (inclusive). :param max_steps: stop after this many applies. :returns: migrations that were applied (in order).

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,
) -> 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.

    :param phase: restrict to ``pre`` or ``post`` migrations only.
    :param target: stop after applying this version (inclusive).
    :param max_steps: stop after this many applies.
    :returns: migrations that were applied (in order).
    """
    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()

    for m in all_migrations:
        if m.key in applied_checksums:
            stored = applied_checksums[m.key]
            current = m.checksum(schema)
            if stored != current:
                logger.warning(
                    "migration-checksum-drift",
                    key=m.key,
                    stored_checksum=stored,
                    current_checksum=current,
                )

    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]

    applied_now: list[Migration] = []
    for migration in pending:
        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),
            )
        applied_now.append(migration)
        if target is not None and migration.version == target:
            break
        if max_steps is not None and len(applied_now) >= max_steps:
            break
    return applied_now

apply_pending_locked async

apply_pending_locked(
    dsn: str,
    *,
    schema: str,
    phase: Phase | None = None,
    target: str | None = None,
    max_steps: int | None = None,
) -> list[Migration]

Apply pending migrations under a session-level advisory lock.

Opens its own connection, acquires pg_advisory_lock to prevent concurrent startup races, applies pending migrations, and closes the connection. 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.

Source code in src/taskq/migrate.py

async def apply_pending_locked(
    dsn: str,
    *,
    schema: str,
    phase: Phase | None = None,
    target: str | None = None,
    max_steps: int | None = None,
) -> list[Migration]:
    """Apply pending migrations under a session-level advisory lock.

    Opens its own connection, acquires ``pg_advisory_lock`` to prevent
    concurrent startup races, applies pending migrations, and closes the
    connection.  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.
    """
    conn: asyncpg.Connection | None = None
    try:
        conn = await asyncpg.connect(dsn)
        await conn.execute("SELECT pg_advisory_lock($1)", _MIGRATION_LOCK_KEY)
        applied = await apply_pending(
            conn, schema=schema, phase=phase, target=target, max_steps=max_steps
        )
        if applied:
            logger.info("applied migrations before startup", count=len(applied))
        else:
            logger.info("no pending migrations")
        return applied
    except Exception as exc:
        raise SystemExit(f"migration failed, aborting startup: {exc}") from exc
    finally:
        if conn is not None:
            with contextlib.suppress(Exception):
                await conn.execute("SELECT pg_advisory_unlock($1)", _MIGRATION_LOCK_KEY)
            with contextlib.suppress(Exception):
                await conn.close()