scripts/database_restore.py

#!/usr/bin/env python3
"""Verify or restore a DocDoe PostgreSQL custom-format backup.

Restore is deliberately locked behind both --execute and an exact database
name confirmation. Without them the command only verifies the archive and
prints the intended target. Credentials never appear in the command line.
"""
from __future__ import annotations

import argparse
import json
import os
import shutil
import subprocess
import sys
from pathlib import Path

from dotenv import load_dotenv

from database_backup import BACKEND_DIR, file_sha256, parse_postgres_url


def verify_manifest(backup_path: Path) -> None:
    manifest_path = backup_path.with_suffix(".manifest.json")
    if not manifest_path.exists():
        raise ValueError(f"Backup manifest is missing: {manifest_path}")
    manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
    expected = str(manifest.get("sha256") or "")
    actual = file_sha256(backup_path)
    if not expected or expected != actual:
        raise ValueError("Backup checksum does not match its manifest")


def restore_backup(
    *,
    backup_path: Path,
    database_url: str,
    execute: bool,
    confirm_database: str | None,
    replace_existing: bool,
) -> None:
    target = parse_postgres_url(database_url)
    backup_path = backup_path.expanduser().resolve()
    if not backup_path.is_file():
        raise ValueError(f"Backup file does not exist: {backup_path}")
    verify_manifest(backup_path)

    pg_restore = shutil.which("pg_restore")
    if not pg_restore:
        raise RuntimeError("pg_restore is not installed or not available on PATH")

    env = target.process_environment()
    verification = subprocess.run(
        [pg_restore, "--list", str(backup_path)],
        env=env,
        stdout=subprocess.DEVNULL,
        stderr=subprocess.PIPE,
        text=True,
        check=False,
    )
    if verification.returncode != 0:
        detail = (verification.stderr or "unknown error").strip().splitlines()[-1]
        raise RuntimeError(f"Backup archive verification failed: {detail[:240]}")

    print(f"[OK] Archive and checksum verified: {backup_path.name}")
    print(f"[TARGET] {target.host}:{target.port}/{target.database}")
    if not execute:
        print("[VERIFY ONLY] Add --execute and --confirm-database to restore.")
        return
    if confirm_database != target.database:
        raise ValueError("--confirm-database must exactly match the target database name")

    command = [
        pg_restore,
        "--exit-on-error",
        "--no-owner",
        "--no-privileges",
        "--dbname",
        target.database,
    ]
    if replace_existing:
        command.extend(["--clean", "--if-exists"])
    command.append(str(backup_path))
    restored = subprocess.run(
        command,
        env=env,
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        check=False,
    )
    if restored.returncode != 0:
        detail = (restored.stderr or "unknown error").strip().splitlines()[-1]
        raise RuntimeError(f"Database restore failed: {detail[:240]}")
    print("[OK] Restore completed. Run application smoke tests before reopening traffic.")


def main(argv: list[str] | None = None) -> int:
    load_dotenv(BACKEND_DIR / ".env", override=False)
    parser = argparse.ArgumentParser(description="Verify or restore a DocDoe database backup.")
    parser.add_argument("backup", type=Path)
    parser.add_argument("--database-url", default=os.getenv("RESTORE_DATABASE_URL"))
    parser.add_argument("--execute", action="store_true")
    parser.add_argument("--confirm-database")
    parser.add_argument("--replace-existing", action="store_true")
    args = parser.parse_args(argv)
    try:
        restore_backup(
            backup_path=args.backup,
            database_url=args.database_url or "",
            execute=args.execute,
            confirm_database=args.confirm_database,
            replace_existing=args.replace_existing,
        )
    except (OSError, RuntimeError, ValueError, json.JSONDecodeError) as exc:
        print(f"[ERROR] {exc}", file=sys.stderr)
        return 1
    return 0


if __name__ == "__main__":
    raise SystemExit(main())