FabledCurator/backend/app/services/backup_service.py

"""FC-3h: first-class backup/restore service for FC.

The database dump (pg_dump), nightly via Beat when the setting is on. With an
off-site destination set up, the dump streams straight to that bucket
(backup_db_offsite), and images and attachments go there too (backup_media).
Without one, the dump is written under <images_root>/_backups/:
  fc_db_<ts>.dump                     — the dump (custom format)
  fc_db_<ts>.json                     — manifest (tag, trigger, row counts)

Service functions are sync (subprocess-bound). Celery tasks in
backend.app.tasks.backup wrap each one with task_run-tracked
lifecycle + soft/hard time limits + retention bookkeeping.
"""
from __future__ import annotations

import json
import logging
import os
import shutil
import subprocess
import tempfile
import time
from contextlib import contextmanager
from datetime import UTC, datetime
from pathlib import Path

log = logging.getLogger(__name__)

_BACKUPS_DIRNAME = "_backups"

# Subprocess-level guardrails BEYOND the Celery soft_time_limit. The Celery
# soft limit signals the Python process; subprocess.Popen in a blocking syscall
# ignores that signal, so these bound the worst case directly. Each sits just
# UNDER its task's Celery soft_time_limit so the bounded-kill (_run_bounded) is
# the primary guard and fires cleanly before Celery's soft/hard limits — which
# matters because an NFS D-state hang defeats even Celery's SIGKILL (the failure
# that wedged the maintenance lane for hours, #739).
#   backup_db_task:     soft=1800s / hard=2100s  → 1700s
_DB_SUBPROCESS_TIMEOUT_S = 1700            # ~28 min, under the 30-min DB soft limit
# Grace after SIGKILL to reap the child. If it can't be reaped in this window
# (an uninterruptible NFS D-state — the failure mode that wedged the
# concurrency-1 maintenance lane for hours, operator-flagged 2026-06-07), we
# STOP waiting and fail fast, freeing the worker slot. The orphan is reaped by
# the OS once its blocking syscall clears.
_KILL_REAP_GRACE_S = 10


# Tables whose rows stay out of every dump. `credential` holds account logins,
# sealed with the key in /images/secrets: a restored install keeps the table,
# empty, and you sign in to each account again.
DATA_LEFT_OUT = ("credential",)

# Every database dump, local or off-site: custom format (-Fc, compressed and
# restored with pg_restore), without ownership or grants so it loads into any
# install.
_PG_DUMP = ["pg_dump", "--no-owner", "--no-acl", "-Fc",
            *(f"--exclude-table-data={t}" for t in DATA_LEFT_OUT)]


def table_counts(conn) -> dict[str, int]:
    """How many rows each table in the public schema holds, read on a psycopg
    connection, in the order of the tables' names."""
    from psycopg import sql

    names = [r[0] for r in conn.execute(
        "SELECT tablename FROM pg_tables WHERE schemaname = 'public' ORDER BY tablename")]
    return {
        n: conn.execute(sql.SQL("SELECT count(*) FROM {}").format(sql.Identifier(n))).fetchone()[0]
        for n in names
    }


def database_counts(conn) -> dict:
    """{"alembic": the schema's revision, "tables": {name: rows}}."""
    has_version = conn.execute("SELECT to_regclass('public.alembic_version')").fetchone()[0]
    alembic = (conn.execute("SELECT version_num FROM alembic_version").fetchone()
               if has_version else None)
    return {"alembic": alembic[0] if alembic else None, "tables": table_counts(conn)}


def read_counts(db_url: str) -> dict:
    """database_counts for the database at `db_url`, as it is now."""
    import psycopg

    with psycopg.connect(_libpq_url(db_url)) as conn:
        return database_counts(conn)


class DumpSnapshot:
    """A read-only transaction whose snapshot pg_dump dumps from (it is handed
    over with --snapshot), so rows counted here are exactly the rows the dump
    holds. A test restore checks every table against these counts."""

    def __init__(self, conn):
        self.conn = conn
        self.id = conn.execute("SELECT pg_export_snapshot()").fetchone()[0]

    def counts(self) -> dict | None:
        """The counts, or None when they can't be read: the dump is good
        without them, and a test restore then says none were recorded."""
        try:
            return self._counts()
        except Exception:  # noqa: BLE001 — see above
            log.exception("backup: the dump's row counts couldn't be read")
            return None

    def _counts(self) -> dict:
        counts = database_counts(self.conn)
        for t in DATA_LEFT_OUT:
            if t in counts["tables"]:
                counts["tables"][t] = 0      # the dump holds the table, not its rows
        return counts


@contextmanager
def dump_snapshot(db_url: str):
    """Open the snapshot a dump is taken from. It stays open until the dump
    is finished: pg_dump can only take it up while it is."""
    import psycopg

    with psycopg.connect(_libpq_url(db_url)) as conn:
        conn.isolation_level = psycopg.IsolationLevel.REPEATABLE_READ
        conn.read_only = True
        yield DumpSnapshot(conn)


def _run_bounded(cmd: list[str], timeout: int) -> None:
    """subprocess.run(check=True, timeout) whose reaper can't itself hang.

    subprocess.run's timeout path SIGKILLs the child then blocks in wait() to
    reap it — but a process stuck in uninterruptible I/O (NFS) can't be reaped,
    so wait() blocks for hours. Here we bound the post-kill reap and re-raise
    TimeoutExpired regardless, so the caller fails fast instead of wedging."""
    proc = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    try:
        out, err = proc.communicate(timeout=timeout)
    except subprocess.TimeoutExpired:
        proc.kill()
        try:
            proc.communicate(timeout=_KILL_REAP_GRACE_S)
        except subprocess.TimeoutExpired:
            pass  # unkillable (D-state) — abandon the reap, fail fast
        raise
    if proc.returncode != 0:
        raise subprocess.CalledProcessError(
            proc.returncode, cmd, output=out, stderr=err
        )


def _libpq_url(sa_url: str) -> str:
    """Strip SQLAlchemy +psycopg/+asyncpg driver suffix for pg_dump/psql."""
    for driver in (
        "postgresql+psycopg",
        "postgresql+asyncpg",
        "postgresql+psycopg2",
    ):
        if sa_url.startswith(driver + "://"):
            return "postgresql://" + sa_url[len(driver) + 3:]
    return sa_url


def _backups_dir(images_root: Path) -> Path:
    p = images_root / _BACKUPS_DIRNAME
    p.mkdir(parents=True, exist_ok=True)
    return p


def _now_ts() -> str:
    return datetime.now(UTC).strftime("%Y%m%dT%H%M%SZ")


def _file_size_or_none(path: Path) -> int | None:
    try:
        return path.stat().st_size
    except OSError:
        return None


def _write_manifest(
    out_dir: Path, *, kind: str, ts: str,
    tag: str | None, triggered_by: str,
    artifact_path: Path, counts: dict | None = None,
) -> Path:
    manifest = {
        "kind": kind,
        "backup_id": f"fc_{kind}_{ts}",
        "tag": tag,
        "triggered_by": triggered_by,
        "created_at": datetime.now(UTC).isoformat(),
        "artifact_path": str(artifact_path),
    }
    if counts is not None:
        manifest["counts"] = counts
    mf = out_dir / f"fc_{kind}_{ts}.json"
    mf.write_text(json.dumps(manifest, indent=2))
    return mf


def backup_db(
    *, db_url: str, images_root: Path,
    tag: str | None = None, triggered_by: str = "manual",
) -> dict:
    """Run pg_dump; write .sql + manifest; return dict for the caller
    to persist into BackupRun. Raises on subprocess failure."""
    ts = _now_ts()
    out_dir = _backups_dir(images_root)
    # Custom format (-Fc): compressed (much smaller on NFS) and restored with
    # pg_restore. The .dump extension marks it as non-SQL. The BackupRun field
    # is still named sql_path — it's just "the db artifact path".
    sql_path = out_dir / f"fc_db_{ts}.dump"
    # Dump to LOCAL disk first, then move the finished file to the (NFS) backups
    # dir. pg_dump's long phase is then a DB-socket wait + local writes — both
    # killable — instead of an NFS write that can hang uninterruptibly. Only the
    # final move touches NFS, and it's a bounded single-file step.
    fd, tmp_name = tempfile.mkstemp(prefix="fc_db_", suffix=".dump")
    os.close(fd)
    tmp_path = Path(tmp_name)
    try:
        with dump_snapshot(db_url) as snap:
            _run_bounded(
                _gently([*_PG_DUMP, f"--snapshot={snap.id}", "-f", str(tmp_path),
                         _libpq_url(db_url)]),
                _DB_SUBPROCESS_TIMEOUT_S,
            )
            counts = snap.counts()
        shutil.move(str(tmp_path), str(sql_path))
    finally:
        if tmp_path.exists():
            tmp_path.unlink(missing_ok=True)
    manifest_path = _write_manifest(
        out_dir, kind="db", ts=ts, tag=tag, triggered_by=triggered_by,
        artifact_path=sql_path, counts=counts,
    )
    return {
        "kind": "db",
        "ts": ts,
        "sql_path": str(sql_path),
        "manifest_path": str(manifest_path),
        "size_bytes": _file_size_or_none(sql_path),
        "counts": counts,
    }


class DumpFailed(Exception):
    """pg_dump exited with an error; the message is its last line of output."""


def offsite_dump_key(ts: str) -> str:
    """Where a database dump lands in the off-site bucket. Dated names sit
    side by side under one prefix, which is what lets the bucket's retention
    thin them by age."""
    return f"database/fabledcurator-{ts}.dump"


def _gently(cmd: list[str]) -> list[str]:
    """`cmd` at low CPU priority, and low disk priority where `ionice` is
    installed, so a long dump leaves the server responsive."""
    io = ["ionice", "-c", "2", "-n", "7"] if shutil.which("ionice") else []
    return ["nice", "-n", "10", *io, *cmd]


def _last_line(f) -> str | None:
    f.seek(0)
    lines = f.read().decode(errors="replace").strip().splitlines()
    return lines[-1].strip() if lines else None


def backup_db_offsite(*, db_url: str, target, limit_mbps: int | None = None) -> dict:
    """Stream `pg_dump` straight into the off-site bucket, in parts, with no
    copy staged on this server. The object appears only once pg_dump has
    exited cleanly and every part has arrived; a dump that fails partway
    leaves nothing in the bucket, and raises DumpFailed with pg_dump's own
    last line. Returns a dict for the caller to record in BackupRun, with
    `counts`: each table's rows in the dump, for the caller to store beside
    it."""
    ts = _now_ts()
    with dump_snapshot(db_url) as snap:
        result = _stream_dump(db_url, target, ts, limit_mbps, snap.id)
        result["counts"] = snap.counts()
    return result


def _stream_dump(db_url: str, target, ts: str, limit_mbps: int | None, snapshot: str) -> dict:
    from .backup_s3 import PART_SIZE, Pacer, Upload

    key = offsite_dump_key(ts)
    deadline = time.monotonic() + _DB_SUBPROCESS_TIMEOUT_S
    with tempfile.TemporaryFile() as err:
        proc = subprocess.Popen(
            _gently([*_PG_DUMP, f"--snapshot={snapshot}", _libpq_url(db_url)]),
            stdout=subprocess.PIPE, stderr=err,
        )
        upload = None
        try:
            upload = Upload(target, key, pacer=Pacer(limit_mbps))
            while part := proc.stdout.read(PART_SIZE):
                if time.monotonic() > deadline:
                    raise subprocess.TimeoutExpired(proc.args, _DB_SUBPROCESS_TIMEOUT_S)
                upload.add(part)
            proc.stdout.close()
            proc.wait(timeout=_KILL_REAP_GRACE_S * 6)
            if proc.returncode != 0:
                raise DumpFailed(_last_line(err) or f"pg_dump exited with {proc.returncode}")
            upload.complete()
        except BaseException:
            if proc.poll() is None:
                proc.kill()
                try:
                    proc.wait(timeout=_KILL_REAP_GRACE_S)
                except subprocess.TimeoutExpired:
                    pass  # unkillable (D-state): abandon the reap, fail fast
            if upload is not None:
                upload.abort()
            raise
    return {
        "kind": "db",
        "ts": ts,
        "key": key,
        "bucket": target.bucket,
        "endpoint": target.endpoint,
        "parts": upload.parts,
        "size_bytes": upload.size,
    }


class RestoreFailed(Exception):
    """psql or pg_restore exited with an error; the message is its last line."""


# A restore waits this long for other connections to let go of the tables it
# replaces, then gives up with Postgres' own reason instead of queueing every
# other query behind it.
_RESTORE_LOCK_TIMEOUT = "60s"
# pg_restore of a large library (vector tables, rebuilt indexes) runs well past
# a dump's length; it is bounded on its own.
_RESTORE_SUBPROCESS_TIMEOUT_S = 3 * 3600


def _run_reporting(cmd: list[str], timeout: int, env: dict | None = None) -> None:
    """Run `cmd`; raise RestoreFailed with the last line it printed when it
    exits non-zero."""
    with tempfile.TemporaryFile() as err:
        proc = subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=err, timeout=timeout,
                              env=env)
        if proc.returncode != 0:
            raise RestoreFailed(_last_line(err) or f"{cmd[0]} exited with {proc.returncode}")


# Every other connection to this database: the web server's, the other
# workers'. Ended before the tables are dropped, so none holds a lock the drop
# waits on, and again once the restore is in, so none keeps statements it
# prepared against the old tables. Each reconnects on its next query.
_END_OTHER_CONNECTIONS = (
    "SELECT pg_terminate_backend(pid) FROM pg_stat_activity "
    "WHERE datname = current_database() AND pid <> pg_backend_pid();"
)


def end_other_connections(*, db_url: str) -> None:
    _run_reporting(["psql", _libpq_url(db_url), "-v", "ON_ERROR_STOP=1",
                    "-c", _END_OTHER_CONNECTIONS], 60)


def wipe_db(*, db_url: str) -> None:
    """Drop and recreate the public schema: an empty database, ready for a
    dump or for `alembic upgrade head`."""
    _run_reporting(
        [
            "psql", _libpq_url(db_url), "-v", "ON_ERROR_STOP=1", "-c",
            f"{_END_OTHER_CONNECTIONS} DROP SCHEMA IF EXISTS public CASCADE; "
            "CREATE SCHEMA public;",
        ],
        120,
        # Set for the session, not as a SET inside the same -c string.
        env={**os.environ, "PGOPTIONS": f"-c lock_timeout={_RESTORE_LOCK_TIMEOUT}"},
    )


# A pg_restore newer than the server sets parameters the server doesn't have
# (pg_restore 17 against Postgres 16: `SET transaction_timeout = 0`). The
# statement fails and nothing depends on it, so its error is not a failed
# restore.
_HARMLESS_RESTORE_ERRORS = (
    'unrecognized configuration parameter "transaction_timeout"',
)


def _restore_errors(stderr: str) -> list[str]:
    """pg_restore's error lines, less the harmless ones."""
    return [
        line.strip() for line in stderr.splitlines()
        if "error:" in line and not any(h in line for h in _HARMLESS_RESTORE_ERRORS)
    ]


def restore_db(*, db_url: str, sql_path: Path) -> None:
    """Wipe the public schema, then load the custom-format dump. Any error
    pg_restore reports, other than a harmless one, raises RestoreFailed with
    the first; the database is then partly loaded, and the caller decides what
    to put back."""
    wipe_db(db_url=db_url)
    # Custom-format (-Fc) dumps are restored with pg_restore, not psql. It runs
    # on past errors (it can't skip the harmless one and stop at the rest), so
    # its errors are judged once it is done.
    cmd = _gently(["pg_restore", "--no-owner", "--no-acl",
                   "-d", _libpq_url(db_url), str(sql_path)])
    with tempfile.TemporaryFile() as err:
        proc = subprocess.run(cmd, stdout=subprocess.DEVNULL, stderr=err,
                              timeout=_RESTORE_SUBPROCESS_TIMEOUT_S)
        if proc.returncode == 0:
            return
        err.seek(0)
        said = err.read().decode(errors="replace")
    errors = _restore_errors(said)
    if errors:
        raise RestoreFailed(errors[0] if len(errors) == 1
                            else f"{errors[0]} ({len(errors)} errors in all)")
    if not any(h in said for h in _HARMLESS_RESTORE_ERRORS):
        # A failure with no error line to show for it.
        lines = said.strip().splitlines()
        raise RestoreFailed(lines[-1].strip() if lines
                            else f"pg_restore exited with {proc.returncode}")


def unlink_artifact_files(
    *,
    sql_path: str | None,
    manifest_path: str | None,
) -> dict:
    """Best-effort unlink of all on-disk files for a BackupRun row.
    Returns dict keyed by label with True/False per file. Missing
    files count as success (missing_ok semantics)."""
    deleted: dict = {}
    for label, p in (
        ("sql", sql_path),
        ("manifest", manifest_path),
    ):
        if not p:
            continue
        path = Path(p)
        try:
            path.unlink(missing_ok=True)
            deleted[label] = True
        except OSError:
            deleted[label] = False
    return deleted