FabledCurator/backend/app/services/backup_restore.py

"""Restoring from the off-site bucket: the database from one of its dated
dumps, then every library file the restored rows name, fetched back to the
path each row records.

The database comes first because it is the map: each image and attachment row
holds its file's path and sha256, which is also the file's name in the bucket
(`media/ab/cd/<sha256><ext>`). A file is checked against that sha256 on
arrival and only then moved into place, so a restore never leaves a wrong or
half-written file in the library. A file already on disk at the right size is
left alone, which makes a restore safe to run again and lets it bring back
only what a damaged library is missing.
"""
from __future__ import annotations

import os
import re
import shutil
import subprocess
import sys
import tempfile
from dataclasses import dataclass
from pathlib import Path

from sqlalchemy import select, text

from ..models import BackupBlob, ImportSettings
from . import backup_destination
from .backup_media import _LIBRARY, blob_key
from .backup_s3 import Pacer, S3Error, S3Target, get_file, list_objects

DUMP_PREFIX = "database/"
_DUMP_NAME = re.compile(r"^database/fabledcurator-(\d{8}T\d{6}Z)\.dump$")
_ALEMBIC_INI = Path(__file__).resolve().parents[3] / "alembic.ini"
_MIGRATE_TIMEOUT_S = 1800


@dataclass(frozen=True)
class Dump:
    key: str
    ts: str           # 20261107T134258Z, UTC
    size_bytes: int
    counted: bool = False   # its row counts sit beside it, for a test restore


def list_dumps(target: S3Target) -> list[Dump]:
    """Every database dump in the bucket, newest first."""
    found, keys, token = [], set(), None
    while True:
        page = list_objects(target, DUMP_PREFIX, 1000, continuation=token)
        for o in page.objects:
            keys.add(o.key)
            if m := _DUMP_NAME.match(o.key):
                found.append((o.key, m.group(1), o.size))
        token = page.next_token
        if not token:
            break
    dumps = [Dump(k, ts, size, counted=k.removesuffix(".dump") + ".counts.json" in keys)
             for k, ts, size in found]
    return sorted(dumps, key=lambda d: d.ts, reverse=True)


def find_dump(target: S3Target, key: str | None) -> Dump:
    """The dump at `key`, or the newest when `key` is None or "latest"."""
    dumps = list_dumps(target)
    if not dumps:
        raise S3Error(f"The {target.bucket} bucket holds no database dumps.")
    if key in (None, "", "latest"):
        return dumps[0]
    for d in dumps:
        if d.key == key:
            return d
    raise S3Error(f"{key} is not among the dumps in the {target.bucket} bucket.")


def fetch_dump(target: S3Target, dump: Dump, into: Path, pacer: Pacer | None) -> Path:
    """Download a dump into the folder `into`, checked against its listed
    size. pg_restore reads it from there."""
    into.mkdir(parents=True, exist_ok=True)
    path = into / Path(dump.key).name
    get_file(target, dump.key, path, size=dump.size_bytes, pacer=pacer)
    return path


def migrate(db_name: str | None = None) -> None:
    """Bring a restored database up to this install's schema. A dump taken by
    an older FabledCurator gains the migrations it predates; one from a newer
    FabledCurator fails here, with alembic's reason, because this install
    doesn't know its schema. `db_name` names another database on the same
    server to bring up instead of this install's own."""
    env = {**os.environ, "DB_NAME": db_name} if db_name else None
    with tempfile.TemporaryFile() as err:
        proc = subprocess.run(
            [sys.executable, "-m", "alembic", "-c", str(_ALEMBIC_INI), "upgrade", "head"],
            cwd=_ALEMBIC_INI.parent, stdout=subprocess.DEVNULL, stderr=err,
            timeout=_MIGRATE_TIMEOUT_S, env=env,
        )
        if proc.returncode != 0:
            err.seek(0)
            lines = err.read().decode(errors="replace").strip().splitlines()
            raise S3Error("The restored database couldn't be brought up to this "
                          f"version: {lines[-1] if lines else 'alembic failed'}")


def keep_destination(session, target: S3Target, limit_mbps: int | None) -> None:
    """Write this install's destination into the restored settings. The
    dump's own copy of the secret is sealed with the old server's credential
    key, which never leaves that server, so it can't be opened here."""
    settings = ImportSettings.load_sync(session)
    settings.backup_s3_endpoint = target.endpoint
    settings.backup_s3_bucket = target.bucket
    settings.backup_s3_access_key = target.access_key
    settings.backup_s3_region = target.region
    settings.backup_s3_secret_sealed = backup_destination.seal(target.secret_key)
    settings.backup_upload_limit_mbps = limit_mbps
    session.commit()


# ── files ───────────────────────────────────────────────────────────────────


@dataclass(frozen=True)
class LibraryFile:
    sha256: str
    paths: tuple[str, ...]     # every row's path for this sha, sorted
    size_bytes: int


def library_files(session, *, after: str = "", limit: int = 200) -> list[LibraryFile]:
    """The next `limit` distinct files the library names, in sha order after
    `after`, each with every path a row gives it."""
    rows = session.execute(text(f"""
        SELECT lib.sha256, array_agg(DISTINCT lib.path ORDER BY lib.path) AS paths,
               max(lib.size_bytes) AS size_bytes
        FROM ({_LIBRARY}) AS lib
        WHERE lib.sha256 > :after
        GROUP BY lib.sha256
        ORDER BY lib.sha256
        LIMIT :limit
    """), {"after": after, "limit": limit}).all()
    return [LibraryFile(r.sha256, tuple(r.paths), r.size_bytes) for r in rows]


def sent_keys(session, tid: str, shas: list[str]) -> dict[str, str]:
    """The bucket key recorded for each sha that was sent to this destination."""
    if not shas:
        return {}
    return dict(session.execute(
        select(BackupBlob.sha256, BackupBlob.key)
        .where(BackupBlob.target == tid, BackupBlob.sha256.in_(shas))
    ).all())


def _resolve(path: str, images_root: Path) -> Path:
    """Where a row's path lands on this server. A path outside the library
    folder is refused: a restore writes only inside it."""
    p = Path(path)
    if not p.is_absolute():
        p = images_root / p
    root = Path(*images_root.parts)
    if ".." in p.parts or not (p == root or root in p.parents):
        raise S3Error(f"{path} is outside {images_root}, so it isn't restored there.")
    return p


def _present(p: Path, size: int) -> bool:
    try:
        return p.is_file() and p.stat().st_size == size
    except OSError:
        return False


def restore_file(target: S3Target, f: LibraryFile, key: str | None,
                 images_root: Path, pacer: Pacer | None) -> str:
    """Put one file back at every path the library gives it. Returns
    "present" (nothing was missing), "copied" (a path already on disk
    supplied the others) or "restored" (fetched from the bucket). Raises
    ObjectMissing when the bucket doesn't hold it, S3Error or OSError
    otherwise."""
    places = [_resolve(p, images_root) for p in f.paths]
    have = [p for p in places if _present(p, f.size_bytes)]
    need = [p for p in places if p not in have]
    if not need:
        return "present"
    outcome = "copied"
    if not have:
        first = need.pop(0)
        first.parent.mkdir(parents=True, exist_ok=True)
        get_file(target, key or blob_key(f.sha256, f.paths[0]), first,
                 sha256=f.sha256, size=f.size_bytes, pacer=pacer)
        have.append(first)
        outcome = "restored"
    for p in need:
        p.parent.mkdir(parents=True, exist_ok=True)
        partial = p.with_name(p.name + ".partial")
        shutil.copyfile(have[0], partial)
        partial.replace(p)
    return outcome