PacketStation/server/packetstation/migrate.py

"""Plain SQL migrations: migrations/NNNN_name.sql applied in order, each recorded with its sha256."""

from __future__ import annotations

import hashlib
import logging
import re
from contextlib import AbstractContextManager
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Protocol

log = logging.getLogger(__name__)

NAME_RE = re.compile(r"^(\d{4})_[a-z0-9_]+\.sql$")
MIGRATE_LOCK = 7310

CREATE_TABLE = """
CREATE TABLE IF NOT EXISTS schema_migrations (
    name       text PRIMARY KEY,
    applied_at timestamptz NOT NULL DEFAULT now(),
    sha256     text NOT NULL
)
"""


class MigrationError(Exception):
    pass


class Conn(Protocol):
    """The subset of psycopg.Connection the runner uses (a fake works in tests)."""

    def execute(self, query: Any, params: Any = None) -> Any: ...
    def transaction(self) -> AbstractContextManager[Any]: ...


@dataclass(frozen=True)
class Migration:
    name: str
    path: Path
    sha256: str

    def sql(self) -> str:
        return self.path.read_text(encoding="utf-8")


def sha256_file(path: Path) -> str:
    return hashlib.sha256(path.read_bytes()).hexdigest()


def discover(directory: Path) -> list[Migration]:
    """Every NNNN_name.sql in name order; other .sql names and duplicate numbers are errors."""
    out: list[Migration] = []
    seen: dict[str, str] = {}
    for p in sorted(directory.glob("*.sql")):
        m = NAME_RE.match(p.name)
        if not m:
            raise MigrationError(f"bad migration file name {p.name} (want NNNN_lower_name.sql)")
        if m.group(1) in seen:
            raise MigrationError(f"two migrations numbered {m.group(1)}: {seen[m.group(1)]} and {p.name}")
        seen[m.group(1)] = p.name
        out.append(Migration(p.name, p, sha256_file(p)))
    return out


def applied(conn: Conn) -> dict[str, str]:
    rows = conn.execute("SELECT name, sha256 FROM schema_migrations ORDER BY name").fetchall()
    return {r[0]: r[1] for r in rows}


def verify(files: list[Migration], done: dict[str, str]) -> list[Migration]:
    """Pending migrations; raises when an applied file changed or vanished."""
    by_name = {m.name: m for m in files}
    for name, digest in done.items():
        m = by_name.get(name)
        if m is None:
            raise MigrationError(f"{name} is recorded as applied but its file is missing")
        if m.sha256 != digest:
            raise MigrationError(
                f"{name} changed after it was applied (sha256 {digest[:12]} -> "
                f"{m.sha256[:12]}); add a new migration instead"
            )
    pending = [m for m in files if m.name not in done]
    if pending and done and pending[0].name < max(done):
        raise MigrationError(f"{pending[0].name} sorts before the last applied migration {max(done)}")
    return pending


def run(conn: Conn, directory: Path) -> list[str]:
    """Apply pending migrations, each in its own transaction. Returns the names applied.

    `conn` must be in autocommit mode so each transaction() block commits on its own.
    """
    files = discover(directory)
    conn.execute("SELECT pg_advisory_lock(%s)", (MIGRATE_LOCK,))
    try:
        with conn.transaction():
            conn.execute(CREATE_TABLE)
        pending = verify(files, applied(conn))
        done: list[str] = []
        for m in pending:
            log.info("applying migration %s", m.name)
            with conn.transaction():
                conn.execute(m.sql())
                conn.execute("INSERT INTO schema_migrations (name, sha256) VALUES (%s, %s)", (m.name, m.sha256))
            done.append(m.name)
        return done
    finally:
        conn.execute("SELECT pg_advisory_unlock(%s)", (MIGRATE_LOCK,))


def status(conn: Conn, directory: Path) -> tuple[bool, str]:
    """(at_head, detail) for doctor; never raises MigrationError."""
    files = discover(directory)
    try:
        pending = verify(files, applied(conn))
    except MigrationError as e:
        return False, str(e)
    if pending:
        return False, f"{len(pending)} pending: {', '.join(m.name for m in pending)}"
    return True, f"at head ({files[-1].name})" if files else "no migrations"


def bootstrap_roles(settings: Any) -> list[str]:
    """dev/test/ci only: create the roles and the database when missing, as PG_SUPERUSER_URL.

    Passwords come from PG_MIGRATE_PASSWORD_FILE and PG_APP_PASSWORD_FILE; existing roles are left alone.
    """
    import psycopg
    from psycopg import sql
    from psycopg.conninfo import make_conninfo

    if settings.stack == "prod":
        raise MigrationError("--bootstrap-roles is refused on STACK=prod")
    if settings.pg_superuser_url is None:
        raise MigrationError("--bootstrap-roles needs PG_SUPERUSER_URL")
    url = settings.pg_superuser_url.get_secret_value()
    did: list[str] = []
    with psycopg.connect(url, autocommit=True, application_name="ps-bootstrap") as su:
        for role in ("migrate", "app"):
            user = settings.pg_user(role)
            if su.execute("SELECT 1 FROM pg_roles WHERE rolname = %s", (user,)).fetchone():
                continue
            su.execute(
                sql.SQL("CREATE ROLE {} LOGIN PASSWORD {}").format(
                    sql.Identifier(user), sql.Literal(settings.pg_password(role))
                )
            )
            did.append(f"role {user}")
        db = settings.pg_database
        if not su.execute("SELECT 1 FROM pg_database WHERE datname = %s", (db,)).fetchone():
            su.execute(
                sql.SQL("CREATE DATABASE {} OWNER {}").format(
                    sql.Identifier(db), sql.Identifier(settings.pg_migrate_user)
                )
            )
            did.append(f"database {db}")
    target = make_conninfo(url, dbname=settings.pg_database)
    with psycopg.connect(target, autocommit=True, application_name="ps-bootstrap") as su:
        su.execute(sql.SQL("REVOKE ALL ON DATABASE {} FROM PUBLIC").format(sql.Identifier(settings.pg_database)))
        su.execute(
            sql.SQL("GRANT CONNECT, TEMPORARY ON DATABASE {} TO {}, {}").format(
                sql.Identifier(settings.pg_database),
                sql.Identifier(settings.pg_migrate_user),
                sql.Identifier(settings.pg_app_user),
            )
        )
        su.execute("REVOKE CREATE ON SCHEMA public FROM PUBLIC")
        su.execute(
            sql.SQL("GRANT USAGE, CREATE ON SCHEMA public TO {}").format(sql.Identifier(settings.pg_migrate_user))
        )
    return did