py-app-runner 0.5.49.dev0__py3-none-any.whl

py_app_runner/migrations/commands.py

@@ -0,0 +1,386 @@
+import pathlib
+import time
+from collections.abc import Callable
+from datetime import datetime
+
+import psycopg
+
+from py_app_runner.migrations.discovery import (
+    MigrationError,
+    count_statements,
+    discover,
+    find_meta_commands,
+    load_migration,
+    new_filename,
+)
+from py_app_runner.migrations.states import MigrationState, State, blocking, compute_states, pending
+from py_app_runner.migrations.tracker import Tracker
+
+Out = Callable[[str], None]
+
+
+async def _load_states(
+    conn: psycopg.AsyncConnection, directory: pathlib.Path, table: str
+) -> tuple[Tracker, list[MigrationState]]:
+    tracker = Tracker(conn, table)
+    await tracker.ensure_table()
+    files = discover(directory)
+    rows = await tracker.applied_rows()
+
+    return tracker, compute_states(files, rows)
+
+
+def _report_blocking(states: list[MigrationState], table: str, out: Out) -> None:
+    """The two blocking states need different advice. `repair` re-reads the file to re-stamp
+    its checksum, so it is the DRIFT remedy and is useless for MISSING - pointing an operator
+    at it there just earns them "no such migration file"."""
+
+    blocked = blocking(states)
+    for state in blocked:
+        if state.state is State.DRIFT:
+            out(f"  {state.state:<8} {state.name}  (file changed after it was applied)")
+        else:
+            out(f"  {state.state:<8} {state.name}  (applied, but the file is gone)")
+
+    out("")
+    if any(state.state is State.DRIFT for state in blocked):
+        out("DRIFT: revert the edit, or run `migrations repair <name>` if the edit was deliberate.")
+
+    for state in blocked:
+        if state.state is not State.MISSING:
+            continue
+
+        out(f"MISSING: {state.name} is gone, so `migrations repair` cannot help. Restore the file")
+        out("from version control, or - if it is gone for good and its schema change is known to")
+        out(f"be in place - drop the tracking row:  DELETE FROM {table} WHERE name = '{state.name}';")
+
+
+async def cmd_status(
+    conn: psycopg.AsyncConnection,
+    directory: pathlib.Path,
+    table: str,
+    check: bool,
+    out: Out,
+) -> int:
+    try:
+        _tracker, states = await _load_states(conn, directory, table)
+    except MigrationError as e:
+        out(f"error: {e}")
+        return 1
+
+    if not states:
+        out("No migrations found.")
+        return 0
+
+    for state in states:
+        applied_at = state.row.applied_at.strftime("%Y-%m-%d %H:%M") if state.row else ""
+        out(f"  {state.state:<8} {state.name:<52} {applied_at}")
+
+    blocked = blocking(states)
+    waiting = pending(states)
+
+    out("")
+    out(f"{len(states) - len(waiting) - len(blocked)} applied, {len(waiting)} pending, {len(blocked)} blocked")
+
+    if check and (waiting or blocked):
+        return 1
+
+    return 0
+
+
+async def cmd_apply(
+    conn: psycopg.AsyncConnection,
+    directory: pathlib.Path,
+    table: str,
+    dry_run: bool,
+    to: str | None,
+    applied_by: str,
+    out: Out,
+) -> int:
+    """Apply pending migrations to `conn`.
+
+    `conn` must be in autocommit mode. `tracker.lock()`'s advisory-lock statement would
+    otherwise implicitly open a transaction, so each migration's `conn.transaction()`
+    degrades to a SAVEPOINT instead of a real BEGIN/COMMIT - a mid-run failure would then
+    roll the whole outer transaction back on unlock, undoing migrations this function
+    already reported as applied. `CREATE INDEX CONCURRENTLY` in the no-transaction branch
+    also requires autocommit outright.
+    """
+
+    if not conn.autocommit:
+        out("error: cmd_apply requires an autocommit connection")
+        return 1
+
+    tracker = Tracker(conn, table)
+
+    async with tracker.lock():
+        try:
+            _tracker, states = await _load_states(conn, directory, table)
+        except MigrationError as e:
+            out(f"error: {e}")
+            return 1
+
+        if blocking(states):
+            _report_blocking(states, table, out)
+            return 1
+
+        queue = pending(states)
+        if to is not None:
+            # blocking(states) was empty, so no MISSING state remains here and every
+            # state has a file - the `is not None` filter is structural narrowing for
+            # pyrefly, not a real behavioural filter.
+            known = {state.file.prefix for state in states if state.file is not None}
+            if to not in known:
+                out(f"error: no migration with prefix {to!r}")
+                return 1
+
+            queue = [state for state in queue if state.file is not None and state.file.prefix <= to]
+
+        if not queue:
+            out("Database is up to date; nothing to apply.")
+            return 0
+
+        for state in queue:
+            assert state.file is not None
+            meta = find_meta_commands(state.file.sql)
+            if meta:
+                out(f"error: {state.name} contains psql meta-command(s) psycopg cannot execute:")
+                for line_number, line in meta:
+                    out(f"    line {line_number}: {line}")
+
+                out("Strip them (pg_dump emits \\restrict / \\unrestrict) and try again.")
+                return 1
+
+            # Postgres wraps a multi-statement simple-Query message in an *implicit*
+            # transaction, so a no-transaction file only genuinely runs outside one when it
+            # holds exactly one statement. Since the whole file goes to `cur.execute()` in a
+            # single call with no statement splitter, the only honest answer is to refuse.
+            if state.file.no_transaction and count_statements(state.file.sql) > 1:
+                out(f"error: {state.name} is marked `-- migrations:no-transaction` but contains")
+                out("more than one statement. Postgres runs a multi-statement send inside an")
+                out("implicit transaction, which would defeat the directive. Split the file so")
+                out("each no-transaction migration contains exactly one statement.")
+                return 1
+
+        for state in queue:
+            assert state.file is not None
+            if dry_run:
+                out(f"would apply {state.name}")
+                continue
+
+            started = time.monotonic()
+            duration_ms = 0
+            try:
+                if state.file.no_transaction:
+                    async with conn.cursor() as cur:
+                        await cur.execute(state.file.sql.encode(conn.info.encoding))
+
+                    duration_ms = int((time.monotonic() - started) * 1000)
+                    await tracker.record(state.name, state.file.checksum, duration_ms, applied_by)
+                else:
+                    # The tracking row is written inside the same transaction as the
+                    # migration, so a file either fully lands and is recorded, or neither.
+                    async with conn.transaction():
+                        async with conn.cursor() as cur:
+                            await cur.execute(state.file.sql.encode(conn.info.encoding))
+
+                        duration_ms = int((time.monotonic() - started) * 1000)
+                        await tracker.record(state.name, state.file.checksum, duration_ms, applied_by)
+
+            except Exception as e:
+                out(f"FAILED {state.name}: {e}")
+                if state.file.no_transaction:
+                    # No transaction means no rollback: a failed CREATE INDEX CONCURRENTLY
+                    # leaves an invalid index behind, and a re-run then dies on "relation
+                    # already exists". Claiming a clean stop here is what strands a deploy.
+                    out("This file ran OUTSIDE a transaction, so it may have PARTIALLY applied.")
+                    out("It was not recorded as applied. Inspect the database and undo whatever")
+                    out("landed (a failed CREATE INDEX CONCURRENTLY leaves an invalid index)")
+                    out("before re-running.")
+                else:
+                    out("Stopped. Nothing after this migration was applied.")
+
+                return 1
+
+            out(f"applied {state.name} ({duration_ms} ms)")
+
+        if dry_run:
+            out(f"{len(queue)} migration(s) would be applied; nothing was changed.")
+        else:
+            out(f"Applied {len(queue)} migration(s).")
+
+        return 0
+
+
+_NEW_FILE_TEMPLATE = """-- {name}
+--
+-- Runs in a transaction. Add `-- migrations:no-transaction` as the very first line if this
+-- file needs CREATE INDEX CONCURRENTLY or anything else Postgres refuses inside one.
+--
+-- A no-transaction file must contain exactly ONE statement. Postgres wraps a
+-- multi-statement send in an implicit transaction, which would defeat the directive, so
+-- `apply` refuses such a file. Split it into one file per statement.
+"""
+
+
+async def cmd_baseline(
+    conn: psycopg.AsyncConnection,
+    directory: pathlib.Path,
+    table: str,
+    to: str | None,
+    assume_yes: bool,
+    applied_by: str,
+    prompt: Callable[[str], str],
+    out: Out,
+) -> int:
+    """Writes tracking rows without executing anything - the adoption path for a database
+    that already has the schema. Nothing is written until every answer is in, so `q` really
+    does leave the table untouched.
+
+    `conn` must be in autocommit mode, for the same reason as `cmd_apply`: `tracker.record()`
+    issues its own INSERT per row with no wrapping transaction and no explicit commit, so on a
+    non-autocommit connection those rows would sit uncommitted - silently discarded if the
+    caller never commits, rather than genuinely "written".
+    """
+
+    if not conn.autocommit:
+        out("error: cmd_baseline requires an autocommit connection")
+        return 1
+
+    tracker = Tracker(conn, table)
+
+    async with tracker.lock():
+        try:
+            _tracker, states = await _load_states(conn, directory, table)
+        except MigrationError as e:
+            out(f"error: {e}")
+            return 1
+
+        candidates = pending(states)
+        if to is not None:
+            # Unlike cmd_apply, this runs without a blocking() guard in front of it, so a
+            # MISSING state (file is None) can still be in `states` here.
+            known = {state.file.prefix for state in states if state.file is not None}
+            if to not in known:
+                out(f"error: no migration with prefix {to!r}")
+                return 1
+
+            candidates = [state for state in candidates if state.file is not None and state.file.prefix <= to]
+
+        if not candidates:
+            out("Nothing to baseline; every migration on disk is already recorded.")
+            return 0
+
+        chosen = []
+        stamp_rest = assume_yes
+        for state in candidates:
+            if stamp_rest:
+                chosen.append(state)
+                continue
+
+            answer = prompt(f"  {state.name:<52} mark as already applied? [y/N/a/q] ").strip().lower()
+            if answer == "q":
+                out("Aborted; nothing was written.")
+                return 1
+
+            if answer == "a":
+                stamp_rest = True
+                chosen.append(state)
+            elif answer == "y":
+                chosen.append(state)
+
+        try:
+            # One transaction for the whole write step: `chosen` was decided in memory
+            # above, so partway-through failure here must not leave some of those decisions
+            # stamped and others not - that would put `apply` in exactly the position
+            # baseline exists to avoid, half-recognising a schema it should treat as known.
+            async with conn.transaction():
+                for state in chosen:
+                    assert state.file is not None
+                    await tracker.record(state.name, state.file.checksum, 0, applied_by)
+        except Exception as e:
+            out(f"error: failed to write tracking rows: {e}")
+            out("Nothing was written.")
+            return 1
+
+        out("")
+        out(
+            f"stamped {len(chosen)} migration(s) as applied (not executed); "
+            f"{len(candidates) - len(chosen)} left pending"
+        )
+
+        return 0
+
+
+def cmd_new(directory: pathlib.Path, name: str, now: datetime, out: Out) -> int:
+    try:
+        filename = new_filename(name, now)
+    except MigrationError as e:
+        out(f"error: {e}")
+        return 1
+
+    if not directory.is_dir():
+        out(f"error: migrations directory does not exist: {directory}")
+        return 1
+
+    path = directory / filename
+    if path.exists():
+        out(f"error: {path} already exists")
+        return 1
+
+    path.write_text(_NEW_FILE_TEMPLATE.format(name=name))
+    out(f"created {path}")
+
+    return 0
+
+
+async def cmd_repair(
+    conn: psycopg.AsyncConnection,
+    directory: pathlib.Path,
+    table: str,
+    name: str,
+    out: Out,
+) -> int:
+    """Re-stamps one migration's checksum after a deliberate edit. The only way out of DRIFT
+    short of reverting the file.
+
+    `conn` must be in autocommit mode, for the same reason as `cmd_baseline`: `update_checksum`
+    is a single UPDATE with no explicit commit of its own.
+    """
+
+    if not conn.autocommit:
+        out("error: cmd_repair requires an autocommit connection")
+        return 1
+
+    # A bare basename only. `directory / "../elsewhere/x.sql"` resolves outside the migrations
+    # directory, and stamping that file's checksum under its basename would leave the real
+    # migration permanently in DRIFT while reporting a successful repair.
+    if name != pathlib.PurePath(name).name:
+        out(f"error: {name!r} must be a plain migration filename, not a path")
+        return 1
+
+    path = directory / name
+    if not path.is_file():
+        out(f"error: no such migration file: {path}")
+        return 1
+
+    try:
+        migration = load_migration(path)
+    except MigrationError as e:
+        out(f"error: {e}")
+        return 1
+
+    tracker = Tracker(conn, table)
+    try:
+        await tracker.ensure_table()
+    except MigrationError as e:
+        out(f"error: {e}")
+        return 1
+
+    if not await tracker.update_checksum(migration.name, migration.checksum):
+        out(f"error: {migration.name} has no tracking row - it was never applied, so there is nothing to repair")
+        return 1
+
+    out(f"repaired {migration.name}: checksum re-stamped to {migration.checksum[:12]}...")
+
+    return 0

py_app_runner/migrations/discovery.py

@@ -0,0 +1,108 @@
+import hashlib
+import pathlib
+import re
+from dataclasses import dataclass
+from datetime import datetime
+
+MIGRATION_FILENAME_RE = re.compile(r"^(\d{4}-\d{2}-\d{2}-\d{6})-([a-z0-9-]+)\.sql$")
+NO_TRANSACTION_DIRECTIVE = "-- migrations:no-transaction"
+
+# psql meta-commands. pg_dump emits \restrict / \unrestrict, and psycopg sends SQL to the
+# server directly - there is no psql to interpret these, so they arrive as syntax errors.
+_META_COMMAND_RE = re.compile(r"^\s*\\")
+
+_LINE_COMMENT_RE = re.compile(r"--[^\n]*")
+
+
+class MigrationError(Exception):
+    """Anything wrong with the migration files themselves, as opposed to the SQL failing."""
+
+
+@dataclass(frozen=True)
+class MigrationFile:
+    name: str
+    prefix: str
+    path: pathlib.Path
+    sql: str
+    checksum: str
+    no_transaction: bool
+
+
+def checksum_bytes(data: bytes) -> str:
+    return hashlib.sha256(data).hexdigest()
+
+
+def load_migration(path: pathlib.Path) -> MigrationFile:
+    match = MIGRATION_FILENAME_RE.match(path.name)
+    if not match:
+        raise MigrationError(f"Bad migration filename {path.name!r}: expected YYYY-MM-DD-HHMMSS-kebab-name.sql")
+
+    raw = path.read_bytes()
+    try:
+        sql = raw.decode("utf-8")
+    except UnicodeDecodeError as exc:
+        raise MigrationError(f"Migration {path.name!r} is not valid UTF-8: {exc}") from exc
+
+    first_line = sql.split("\n", 1)[0].strip()
+
+    return MigrationFile(
+        name=path.name,
+        prefix=match.group(1),
+        path=path,
+        sql=sql,
+        checksum=checksum_bytes(raw),
+        no_transaction=first_line == NO_TRANSACTION_DIRECTIVE,
+    )
+
+
+def discover(directory: pathlib.Path) -> list[MigrationFile]:
+    """Every *.sql in `directory`, chronological. Filenames sort chronologically as plain
+    text, which is the whole point of the timestamp prefix."""
+
+    if not directory.is_dir():
+        raise MigrationError(f"Migrations directory does not exist: {directory}")
+
+    migrations = [load_migration(path) for path in sorted(directory.glob("*.sql"))]
+
+    seen: dict[str, str] = {}
+    for migration in migrations:
+        if migration.prefix in seen:
+            raise MigrationError(
+                f"Duplicate migration timestamp {migration.prefix}: {seen[migration.prefix]} and {migration.name}"
+            )
+
+        seen[migration.prefix] = migration.name
+
+    return migrations
+
+
+def find_meta_commands(sql: str) -> list[tuple[int, str]]:
+    """1-indexed line numbers of psql meta-commands. A backslash inside a string literal is
+    not one, and only a line that *starts* with a backslash can be."""
+
+    return [
+        (number, line.rstrip())
+        for number, line in enumerate(sql.split("\n"), start=1)
+        if _META_COMMAND_RE.match(line)
+    ]
+
+
+def count_statements(sql: str) -> int:
+    """Crude statement count: strip `--` comments, split on `;`, count non-empty remainders.
+
+    A semicolon inside a dollar-quoted block or a string literal is counted as a separator
+    here, so such a file is over-counted. That is acceptable because the only caller uses this
+    to *refuse* a multi-statement no-transaction file before anything executes - the
+    false positive is a loud refusal at scan time, and the operations the no-transaction
+    directive exists for (CREATE INDEX CONCURRENTLY and friends) are never dollar-quoted.
+    """
+
+    return len([part for part in _LINE_COMMENT_RE.sub("", sql).split(";") if part.strip()])
+
+
+def new_filename(name: str, now: datetime) -> str:
+    slug = re.sub(r"[^a-z0-9]+", "-", name.lower()).strip("-")
+    if not slug:
+        raise MigrationError(f"Migration name {name!r} contains no usable characters")
+
+    return f"{now:%Y-%m-%d-%H%M%S}-{slug}.sql"

py_app_runner/migrations/states.py

@@ -0,0 +1,63 @@
+from dataclasses import dataclass
+from datetime import datetime
+from enum import StrEnum
+
+from py_app_runner.migrations.discovery import MigrationFile
+
+
+class State(StrEnum):
+    APPLIED = "applied"
+    PENDING = "PENDING"
+    DRIFT = "DRIFT"
+    MISSING = "MISSING"
+
+
+@dataclass(frozen=True)
+class AppliedRow:
+    name: str
+    checksum: str
+    applied_at: datetime
+
+
+@dataclass(frozen=True)
+class MigrationState:
+    name: str
+    state: State
+    file: MigrationFile | None
+    row: AppliedRow | None
+
+
+def compute_states(files: list[MigrationFile], rows: list[AppliedRow]) -> list[MigrationState]:
+    """Union of disk and table, ordered by name. A row with no file (MISSING) sorts in
+    among the files rather than being appended, so `status` reads chronologically."""
+
+    by_name = {file.name: file for file in files}
+    rows_by_name = {row.name: row for row in rows}
+
+    states = []
+    for name in sorted(by_name.keys() | rows_by_name.keys()):
+        file = by_name.get(name)
+        row = rows_by_name.get(name)
+
+        if file is None:
+            state = State.MISSING
+        elif row is None:
+            state = State.PENDING
+        elif row.checksum != file.checksum:
+            state = State.DRIFT
+        else:
+            state = State.APPLIED
+
+        states.append(MigrationState(name=name, state=state, file=file, row=row))
+
+    return states
+
+
+def pending(states: list[MigrationState]) -> list[MigrationState]:
+    return [state for state in states if state.state is State.PENDING]
+
+
+def blocking(states: list[MigrationState]) -> list[MigrationState]:
+    """States that must be resolved before anything may be applied."""
+
+    return [state for state in states if state.state in (State.DRIFT, State.MISSING)]

py_app_runner/migrations/tracker.py

@@ -0,0 +1,141 @@
+import contextlib
+from collections.abc import AsyncIterator
+
+import psycopg
+from psycopg import sql
+from psycopg.pq import TransactionStatus
+
+from py_app_runner.migrations.discovery import MigrationError
+from py_app_runner.migrations.states import AppliedRow
+
+# One fixed key for the whole system. Two processes applying migrations at once would
+# interleave DDL and duplicate tracking rows, so `apply` serialises on this.
+ADVISORY_LOCK_KEY = 4829173465872001
+
+EXPECTED_COLUMNS = frozenset({"id", "name", "checksum", "applied_at", "duration_ms", "applied_by"})
+
+# The schema is resolved from the oid rather than matched on `table_name` alone, so a
+# same-named table in another schema cannot answer for the one search_path actually picks.
+_COLUMNS_SQL = """
+    SELECT column_name
+    FROM information_schema.columns
+    WHERE (table_schema, table_name) = (
+        SELECT n.nspname, c.relname
+        FROM pg_class c
+        JOIN pg_namespace n ON n.oid = c.relnamespace
+        WHERE c.oid = to_regclass(%s)
+    )
+"""
+
+
+class Tracker:
+    """Reads and writes the migration tracking table.
+
+    The table bootstraps itself on every invocation rather than being migration zero -
+    otherwise there is nothing to record the fact that it was created.
+    """
+
+    def __init__(self, conn: "psycopg.AsyncConnection", table: str = "migrations") -> None:
+        self.conn = conn
+        self.table_name = table
+        self.table = sql.Identifier(table)
+
+    async def ensure_table(self) -> None:
+        await self._create_table()
+        await self._verify_columns()
+
+    async def _create_table(self) -> None:
+        statement = sql.SQL(
+            """
+            CREATE TABLE IF NOT EXISTS {table} (
+                id          INT8 GENERATED ALWAYS AS IDENTITY PRIMARY KEY,
+                name        TEXT NOT NULL UNIQUE,
+                checksum    TEXT NOT NULL,
+                applied_at  TIMESTAMPTZ NOT NULL DEFAULT CURRENT_TIMESTAMP,
+                duration_ms INT4 NOT NULL,
+                applied_by  TEXT NOT NULL
+            )
+            """
+        ).format(table=self.table)
+
+        async with self.conn.cursor() as cur:
+            await cur.execute(statement)
+
+    async def _verify_columns(self) -> None:
+        """`CREATE TABLE IF NOT EXISTS` matches on name alone, and `migrations` is about as
+        generic a table name as exists. Without this check an unrelated pre-existing table is
+        silently adopted and the next query fails with a bare `column "name" does not exist`,
+        with nothing pointing at this tool - exactly when an operator is adopting a database.
+        """
+
+        async with self.conn.cursor() as cur:
+            await cur.execute(_COLUMNS_SQL, (self.table.as_string(self.conn),))
+            found = {row[0] for row in await cur.fetchall()}
+
+        missing = EXPECTED_COLUMNS - found
+        if not missing:
+            return
+
+        raise MigrationError(
+            f"Table {self.table_name!r} exists but is not a migration tracking table: "
+            f"missing column(s) {', '.join(sorted(missing))}; it has {', '.join(sorted(found)) or 'no columns'}. "
+            f'Point the tracker at a free name via config["migrations"]["table"], '
+            f"or drop/rename the existing table if it is not in use."
+        )
+
+    async def applied_rows(self) -> list[AppliedRow]:
+        statement = sql.SQL("SELECT name, checksum, applied_at FROM {table} ORDER BY name").format(
+            table=self.table
+        )
+
+        async with self.conn.cursor() as cur:
+            await cur.execute(statement)
+            return [AppliedRow(name=row[0], checksum=row[1], applied_at=row[2]) for row in await cur.fetchall()]
+
+    async def record(self, name: str, checksum: str, duration_ms: int, applied_by: str) -> None:
+        statement = sql.SQL(
+            "INSERT INTO {table} (name, checksum, duration_ms, applied_by) VALUES (%s, %s, %s, %s)"
+        ).format(table=self.table)
+
+        async with self.conn.cursor() as cur:
+            await cur.execute(statement, (name, checksum, duration_ms, applied_by))
+
+    async def update_checksum(self, name: str, checksum: str) -> bool:
+        """False when no such row exists - the caller reports that rather than claiming a
+        repair that never happened."""
+
+        statement = sql.SQL("UPDATE {table} SET checksum = %s WHERE name = %s").format(table=self.table)
+
+        async with self.conn.cursor() as cur:
+            await cur.execute(statement, (checksum, name))
+            return cur.rowcount == 1
+
+    @contextlib.asynccontextmanager
+    async def lock(self) -> AsyncIterator[None]:
+        """Session-level advisory lock held for the whole run. A crashed process drops its
+        connection, and Postgres releases the lock with it.
+
+        The lock will wrap real migration execution, so a failing body may leave a
+        non-autocommit connection in an aborted transaction. Releasing the lock then needs a
+        rollback first, and its own failure must never replace the body's real exception as
+        what the caller sees.
+        """
+
+        async with self.conn.cursor() as cur:
+            await cur.execute("SELECT pg_advisory_lock(%s)", (ADVISORY_LOCK_KEY,))
+
+        try:
+            yield
+        except BaseException:
+            with contextlib.suppress(Exception):
+                await self._unlock()
+            raise
+        else:
+            await self._unlock()
+
+    async def _unlock(self) -> None:
+        if self.conn.info.transaction_status == TransactionStatus.INERROR:
+            await self.conn.rollback()
+
+        async with self.conn.cursor() as cur:
+            await cur.execute("SELECT pg_advisory_unlock(%s)", (ADVISORY_LOCK_KEY,))