from __future__ import annotations import contextlib from collections.abc import Iterator from pathlib import Path from typing import Any import pytest from pshelpers import needs_pg from packetstation import db, migrate from packetstation.settings import Settings class FakeResult: def __init__(self, rows: list[tuple[Any, ...]] | None = None): self.rows = rows or [] def fetchall(self) -> list[tuple[Any, ...]]: return self.rows def fetchone(self) -> tuple[Any, ...] | None: return self.rows[0] if self.rows else None class FakeConn: """Records statements; keeps schema_migrations in a dict; rolls back on error.""" def __init__(self, fail_on: str | None = None): self.table: dict[str, str] = {} self.executed: list[str] = [] self.fail_on = fail_on self._pending: dict[str, str] | None = None def execute(self, query: Any, params: Any = None) -> FakeResult: q = str(query) self.executed.append(q) if self.fail_on and self.fail_on in q: raise RuntimeError("boom") if q.startswith("SELECT name, sha256 FROM schema_migrations"): return FakeResult(sorted(self.table.items())) if q.startswith("INSERT INTO schema_migrations"): target = self._pending if self._pending is not None else self.table target[params[0]] = params[1] return FakeResult() @contextlib.contextmanager def transaction(self) -> Iterator[None]: self._pending = {} try: yield except Exception: self._pending = None raise self.table.update(self._pending) self._pending = None def write(d: Path, name: str, sql: str) -> None: (d / name).write_text(sql, encoding="utf-8") def test_applies_in_order_and_records_hashes(tmp_path: Path) -> None: write(tmp_path, "0002_two.sql", "-- two") write(tmp_path, "0001_one.sql", "-- one") write(tmp_path, "0010_ten.sql", "-- ten") conn = FakeConn() assert migrate.run(conn, tmp_path) == ["0001_one.sql", "0002_two.sql", "0010_ten.sql"] bodies = [q for q in conn.executed if q.startswith("-- ")] assert bodies == ["-- one", "-- two", "-- ten"] assert conn.table["0001_one.sql"] == migrate.sha256_file(tmp_path / "0001_one.sql") assert "CREATE TABLE IF NOT EXISTS schema_migrations" in conn.executed[1] assert migrate.run(conn, tmp_path) == [] def test_refuses_changed_file(tmp_path: Path) -> None: write(tmp_path, "0001_one.sql", "-- one") conn = FakeConn() migrate.run(conn, tmp_path) write(tmp_path, "0001_one.sql", "-- one, edited") write(tmp_path, "0002_two.sql", "-- two") with pytest.raises(migrate.MigrationError, match="changed after it was applied"): migrate.run(conn, tmp_path) assert "0002_two.sql" not in conn.table assert "SELECT pg_advisory_unlock(%s)" in conn.executed[-1] def test_refuses_missing_and_out_of_order(tmp_path: Path) -> None: write(tmp_path, "0002_two.sql", "-- two") conn = FakeConn() migrate.run(conn, tmp_path) write(tmp_path, "0001_one.sql", "-- late") with pytest.raises(migrate.MigrationError, match="sorts before"): migrate.run(conn, tmp_path) (tmp_path / "0002_two.sql").unlink() with pytest.raises(migrate.MigrationError, match="missing"): migrate.run(conn, tmp_path) def test_failed_migration_not_recorded(tmp_path: Path) -> None: write(tmp_path, "0001_one.sql", "-- one") write(tmp_path, "0002_bad.sql", "-- bad") conn = FakeConn(fail_on="-- bad") with pytest.raises(RuntimeError): migrate.run(conn, tmp_path) assert list(conn.table) == ["0001_one.sql"] @pytest.mark.parametrize("name", ["1_x.sql", "0001-x.sql", "0001_X.sql"]) def test_bad_names(tmp_path: Path, name: str) -> None: write(tmp_path, name, "") with pytest.raises(migrate.MigrationError, match="bad migration file name"): migrate.discover(tmp_path) def test_duplicate_numbers(tmp_path: Path) -> None: write(tmp_path, "0001_a.sql", "") write(tmp_path, "0001_b.sql", "") with pytest.raises(migrate.MigrationError, match="two migrations"): migrate.discover(tmp_path) def test_status(tmp_path: Path) -> None: write(tmp_path, "0001_one.sql", "-- one") conn = FakeConn() assert migrate.status(conn, tmp_path)[0] is False migrate.run(conn, tmp_path) assert migrate.status(conn, tmp_path) == (True, "at head (0001_one.sql)") def test_shipped_migrations_are_well_formed() -> None: files = migrate.discover(db.MIGRATIONS_DIR) assert files[0].name == "0001_base.sql" sql = files[0].sql() assert "schema_migrations (" not in sql # the runner creates it revoke = sql.index("REVOKE UPDATE, DELETE, TRUNCATE ON audit_events FROM ps_app") assert sql.index("ALTER DEFAULT PRIVILEGES") < sql.index("CREATE TABLE audit_events") < revoke def test_bootstrap_refused_on_prod(tmp_path: Path) -> None: s = Settings(_env_file=None, stack="prod", pg_superuser_url="postgresql://x@y/z") with pytest.raises(migrate.MigrationError, match="refused"): migrate.bootstrap_roles(s) with pytest.raises(migrate.MigrationError, match="PG_SUPERUSER_URL"): migrate.bootstrap_roles(Settings(_env_file=None, stack="dev")) # ---------------------------------------------------------------- real PostgreSQL @needs_pg def test_pg_migrations_recorded(pg_settings: Settings) -> None: with db.connect_as("migrate", pg_settings, autocommit=True) as conn: assert migrate.run(conn, db.MIGRATIONS_DIR) == [] rows = conn.execute("SELECT name, sha256 FROM schema_migrations ORDER BY name").fetchall() assert rows[0] == ("0001_base.sql", migrate.sha256_file(db.MIGRATIONS_DIR / "0001_base.sql")) @needs_pg def test_pg_refuses_changed_file(pg_settings: Settings, tmp_path: Path) -> None: d = tmp_path / "m" d.mkdir() write(d, "0001_base.sql", (db.MIGRATIONS_DIR / "0001_base.sql").read_text() + "\n-- edited\n") with ( db.connect_as("migrate", pg_settings, autocommit=True) as conn, pytest.raises(migrate.MigrationError, match="changed"), ): migrate.run(conn, d) @needs_pg def test_pg_app_cannot_update_or_delete_audit(pg_settings: Settings) -> None: import psycopg from packetstation import audit with db.connect_as("app", pg_settings, autocommit=True) as conn: audit.append(conn, action="test.event", detail={"a": 1}) for stmt in ("UPDATE audit_events SET action = 'x'", "DELETE FROM audit_events", "TRUNCATE audit_events"): with pytest.raises(psycopg.errors.InsufficientPrivilege): conn.execute(stmt) conn.execute("SELECT key FROM settings").fetchall() conn.execute("INSERT INTO settings (key, value) VALUES ('k', '1')") with ( db.connect_as("migrate", pg_settings, autocommit=True) as conn, pytest.raises(psycopg.errors.InsufficientPrivilege), ): conn.execute("UPDATE audit_events SET action = 'x'") # trigger refuses even the owner