"""Bringing Steward back from a backup file into an empty database. pg_restore runs in its three sections so the rows can be put right between the data and the constraints: a kept table may point at a row of a table whose rows were left out (an alert event naming the Ansible run it started). Each such reference gets what its foreign key's ON DELETE says would happen had that row been deleted: SET NULL clears it, anything else removes the row. Then the constraints and indexes go on. Migrations are not run here: the next start of Steward runs them, as it always does, so a backup made by an older Steward comes up to date. """ from __future__ import annotations import asyncio import logging import tempfile from dataclasses import dataclass from pathlib import Path from sqlalchemy import MetaData, text from sqlalchemy.ext.asyncio import create_async_engine from .contents import EXCLUDED from .dump import BackupError, pg_env, read_backup logger = logging.getLogger(__name__) RESTORE_DEADLINE_SECONDS = 30 * 60 @dataclass class Restored: manifest: dict secret_key: str fixed: int # rows put right between data and constraints def all_metadata() -> MetaData: """Steward's tables, every bundled plugin's included, without an app: the restore runs before anything else would load them.""" import importlib import pkgutil import sys import plugins import steward.models # noqa: F401 — registers core tables from steward.models.base import Base for info in pkgutil.iter_modules(plugins.__path__): # A plugin the app already loaded is under the loader's name; its # models imported a second time would map every table twice. loaded = f"_steward_plugin_{info.name}" package = loaded if loaded in sys.modules else f"plugins.{info.name}" try: importlib.import_module(f"{package}.models") except ImportError: # no models, or an optional dependency missing continue return Base.metadata def dangling_fixes(metadata: MetaData) -> list[str]: """SQL putting right the rows of kept tables that point at a table whose rows were left out, by each foreign key's own ON DELETE.""" out: list[str] = [] for table in sorted(metadata.tables.values(), key=lambda t: t.name): if table.name in EXCLUDED: continue for fk in sorted(table.foreign_keys, key=lambda f: f.parent.name): parent = fk.column.table.name if parent not in EXCLUDED: continue child, col, pcol = table.name, fk.parent.name, fk.column.name gone = (f'"{col}" IS NOT NULL AND NOT EXISTS ' f'(SELECT 1 FROM "{parent}" p WHERE p."{pcol}" = "{child}"."{col}")') if (fk.ondelete or "").upper() == "SET NULL": out.append(f'UPDATE "{child}" SET "{col}" = NULL WHERE {gone}') else: out.append(f'DELETE FROM "{child}" WHERE {gone}') return out async def _pg_restore(database_url: str, dump: Path, section: str) -> None: args = ["pg_restore", "--no-owner", "--no-privileges", "--exit-on-error", f"--section={section}", "--dbname", make_dbname(database_url), str(dump)] try: proc = await asyncio.create_subprocess_exec( *args, env=pg_env(database_url), stdout=asyncio.subprocess.DEVNULL, stderr=asyncio.subprocess.PIPE) except FileNotFoundError: raise BackupError("pg_restore isn't installed here.") from None try: _, err = await asyncio.wait_for(proc.communicate(), RESTORE_DEADLINE_SECONDS) except TimeoutError: proc.kill() await proc.wait() raise BackupError(f"pg_restore ({section}) didn't finish within " f"{RESTORE_DEADLINE_SECONDS // 60} minutes.") from None if proc.returncode != 0: lines = [ln for ln in err.decode(errors="replace").splitlines() if ln.strip()] raise BackupError(f"pg_restore ({section}): " + (lines[-1] if lines else f"exited {proc.returncode}.")) def make_dbname(database_url: str) -> str: """The database name for pg_restore's --dbname; host, user and password come from the environment pg_env builds.""" from sqlalchemy.engine import make_url return make_url(database_url).database or "" def _async_url(database_url: str) -> str: from sqlalchemy.engine import make_url url = make_url(database_url) return str(url.set(drivername="postgresql+asyncpg")) if url.drivername == "postgresql" else database_url async def _tables(database_url: str) -> list[str]: engine = create_async_engine(_async_url(database_url)) try: async with engine.connect() as conn: return list((await conn.execute(text( "SELECT tablename FROM pg_tables WHERE schemaname = 'public' ORDER BY tablename"))).scalars()) finally: await engine.dispose() async def _execute(database_url: str, statements: list[str]) -> int: engine = create_async_engine(_async_url(database_url)) changed = 0 try: async with engine.begin() as conn: for sql in statements: result = await conn.execute(text(sql)) changed += max(result.rowcount or 0, 0) finally: await engine.dispose() return changed async def restore(path: Path, database_url: str, *, wipe: bool = False, metadata: MetaData | None = None) -> Restored: """Restore a backup file into the database at `database_url`. It must be empty unless `wipe`, which drops everything in its public schema first.""" existing = await _tables(database_url) if existing and not wipe: raise BackupError(f"The database already has {len(existing)} tables " f"({', '.join(existing[:3])}{'…' if len(existing) > 3 else ''}). " "Restore into an empty database, or pass --wipe to replace what is there.") if existing: logger.warning("Restore: dropping %d tables in the public schema first", len(existing)) await _execute(database_url, ["DROP SCHEMA public CASCADE", "CREATE SCHEMA public"]) with tempfile.TemporaryDirectory() as tmp: info, dump, key = await asyncio.to_thread(read_backup, path, Path(tmp)) await _pg_restore(database_url, dump, "pre-data") await _pg_restore(database_url, dump, "data") fixed = await _execute(database_url, dangling_fixes(metadata or all_metadata())) await _pg_restore(database_url, dump, "post-data") return Restored(info, key, fixed)