FabledSteward/steward/backup/restore.py

"""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(plugin_dirs: list[Path] | None = None) -> MetaData:
    """Steward's tables, every plugin's included, without an app: the
    restore runs before anything else would load them. Plugins are found in
    `plugin_dirs` (the bootstrap's, bundled first) and loaded by path under
    the same module names the app gives them, so a plugin the app already
    loaded isn't mapped twice. Without dirs (tests, a checkout): the
    `plugins` package, as the tests import it."""
    import importlib
    import sys

    import steward
    import steward.models  # noqa: F401 — registers core tables
    from steward.core.plugin_manager import _import_plugin, resolve_plugin_path
    from steward.models.base import Base

    roots = [Path(d).resolve() for d in (plugin_dirs or [Path(steward.__file__).parent.parent / "plugins"])]
    by_path = bool(plugin_dirs)
    names = sorted({p.name for root in roots if root.is_dir() for p in root.iterdir()
                    if (p / "__init__.py").exists() and (p / "models.py").exists()})
    for name in names:
        if f"plugins.{name}.models" in sys.modules or f"_steward_plugin_{name}.models" in sys.modules:
            continue          # already mapped, under one name or the other
        try:
            if by_path:
                package = _import_plugin(name, resolve_plugin_path(roots, name)).__name__
            else:
                package = f"plugins.{name}"
            importlib.import_module(f"{package}.models")
        except Exception:     # a plugin that won't load: its tables are left as pg_restore put them
            logger.warning("Restore: couldn't load plugin %r's tables", name, exc_info=True)
    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


# pg_restore 17 opens every script with this; a Postgres 16 server doesn't
# know the setting and refuses it, so the line is dropped on the way through.
_NEWER_THAN_SERVER = {b"SET transaction_timeout = 0;"}
# A COPY row is one line of the script; this bounds the longest one read.
_LINE_LIMIT = 256 * 1024 * 1024


async def _pg_restore(database_url: str, dump: Path, section: str) -> None:
    """One section of the dump into the database: pg_restore writes it as a
    script, psql runs it and stops at the first error."""
    env = pg_env(database_url)
    try:
        script = await asyncio.create_subprocess_exec(
            "pg_restore", "--no-owner", "--no-privileges", f"--section={section}", "--file=-", str(dump),
            env=env, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE, limit=_LINE_LIMIT)
        psql = await asyncio.create_subprocess_exec(
            "psql", "-X", "-q", "-v", "ON_ERROR_STOP=1",
            env=env, stdin=asyncio.subprocess.PIPE,
            stdout=asyncio.subprocess.DEVNULL, stderr=asyncio.subprocess.PIPE)
    except FileNotFoundError:
        raise BackupError("pg_restore or psql isn't installed here.") from None

    async def pump() -> None:
        try:
            while line := await script.stdout.readline():
                if line.strip() in _NEWER_THAN_SERVER:
                    continue
                psql.stdin.write(line)
                await psql.stdin.drain()
        except (BrokenPipeError, ConnectionResetError):
            # psql stopped on an error (its stderr says which); pg_restore
            # would block on a pipe nobody reads.
            script.kill()
        finally:
            psql.stdin.close()

    async def run() -> tuple[bytes, bytes]:
        _, script_err, psql_err = await asyncio.gather(
            pump(), script.stderr.read(), psql.stderr.read())
        await asyncio.gather(script.wait(), psql.wait())
        return script_err, psql_err

    try:
        script_err, psql_err = await asyncio.wait_for(run(), RESTORE_DEADLINE_SECONDS)
    except TimeoutError:
        for proc in (script, psql):
            if proc.returncode is None:
                proc.kill()
                await proc.wait()
        raise BackupError(f"Restoring {section} didn't finish within "
                          f"{RESTORE_DEADLINE_SECONDS // 60} minutes.") from None
    for proc, err, name in ((psql, psql_err, "psql"), (script, script_err, "pg_restore")):
        if proc.returncode != 0:
            lines = [ln for ln in err.decode(errors="replace").splitlines() if ln.strip()]
            raise BackupError(f"Restoring {section}: " + (lines[-1] if lines else f"{name} exited {proc.returncode}."))


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, plugin_dirs: list[Path] | 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(plugin_dirs)))
        await _pg_restore(database_url, dump, "post-data")
    return Restored(info, key, fixed)