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