scripts/database_backup.py

#!/usr/bin/env python3
"""Create and verify a PostgreSQL backup without exposing credentials.

The script is designed for a scheduler or managed job runner. It writes to a
temporary file, verifies the archive with pg_restore, records a checksum-only
manifest, then atomically promotes the backup. Credentials are passed through
PG* environment variables and never included in the command line or output.

Examples:
    python scripts/database_backup.py --dry-run
    python scripts/database_backup.py --output-dir D:\\docdoe-backups
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import shutil
import subprocess
import sys
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
from pathlib import Path
from urllib.parse import parse_qs, unquote, urlsplit

from dotenv import load_dotenv


BACKEND_DIR = Path(__file__).resolve().parents[1]
DEFAULT_OUTPUT_DIR = BACKEND_DIR / "backups" / "database"
BACKUP_PREFIX = "docdoe-"


@dataclass(frozen=True)
class PostgresTarget:
    host: str
    port: int
    database: str
    username: str
    password: str
    sslmode: str | None

    def process_environment(self) -> dict[str, str]:
        env = os.environ.copy()
        env.update(
            {
                "PGHOST": self.host,
                "PGPORT": str(self.port),
                "PGDATABASE": self.database,
                "PGUSER": self.username,
                "PGPASSWORD": self.password,
                "PGCONNECT_TIMEOUT": env.get("PGCONNECT_TIMEOUT", "15"),
            }
        )
        if self.sslmode:
            env["PGSSLMODE"] = self.sslmode
        return env


def parse_postgres_url(database_url: str) -> PostgresTarget:
    raw = database_url.strip()
    if not raw:
        raise ValueError("DATABASE_URL is missing")

    normalized = raw.replace("postgresql+psycopg://", "postgresql://", 1)
    parsed = urlsplit(normalized)
    if parsed.scheme not in {"postgres", "postgresql"}:
        raise ValueError("Backups require a PostgreSQL DATABASE_URL")
    if not parsed.hostname or not parsed.username:
        raise ValueError("DATABASE_URL must include a host and username")

    database = unquote(parsed.path.lstrip("/"))
    if not database:
        raise ValueError("DATABASE_URL must include a database name")

    query = parse_qs(parsed.query)
    sslmode = query.get("sslmode", [None])[0]
    return PostgresTarget(
        host=parsed.hostname,
        port=parsed.port or 5432,
        database=database,
        username=unquote(parsed.username),
        password=unquote(parsed.password or ""),
        sslmode=sslmode,
    )


def backup_filename(now: datetime) -> str:
    stamp = now.astimezone(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
    return f"{BACKUP_PREFIX}{stamp}.dump"


def file_sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as handle:
        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
            digest.update(chunk)
    return digest.hexdigest()


def prune_expired_backups(output_dir: Path, retention_days: int, now: datetime) -> list[Path]:
    if retention_days < 1:
        raise ValueError("Retention must be at least one day")
    if not output_dir.exists():
        return []

    cutoff = now.astimezone(timezone.utc) - timedelta(days=retention_days)
    removed: list[Path] = []
    for backup in output_dir.glob(f"{BACKUP_PREFIX}*.dump"):
        modified = datetime.fromtimestamp(backup.stat().st_mtime, tz=timezone.utc)
        if modified >= cutoff:
            continue
        manifest = backup.with_suffix(".manifest.json")
        backup.unlink()
        removed.append(backup)
        if manifest.exists():
            manifest.unlink()
            removed.append(manifest)
    return removed


def _resolve_binary(name: str, dry_run: bool) -> str:
    resolved = shutil.which(name)
    if resolved:
        return resolved
    if dry_run:
        return name
    raise RuntimeError(f"{name} is not installed or not available on PATH")


def _run(command: list[str], env: dict[str, str], *, stdout=None) -> None:
    completed = subprocess.run(
        command,
        env=env,
        stdout=stdout,
        stderr=subprocess.PIPE,
        text=True,
        check=False,
    )
    if completed.returncode != 0:
        detail = (completed.stderr or "unknown error").strip().splitlines()[-1]
        raise RuntimeError(f"Database backup command failed: {detail[:240]}")


def create_backup(
    *,
    database_url: str,
    output_dir: Path,
    retention_days: int,
    verify: bool,
    dry_run: bool,
    now: datetime | None = None,
) -> Path | None:
    target = parse_postgres_url(database_url)
    moment = now or datetime.now(timezone.utc)
    output_dir = output_dir.expanduser().resolve()
    final_path = output_dir / backup_filename(moment)
    partial_path = final_path.with_suffix(".dump.partial")
    pg_dump = _resolve_binary("pg_dump", dry_run)
    pg_restore = _resolve_binary("pg_restore", dry_run) if verify else None

    if dry_run:
        print(
            f"[DRY RUN] PostgreSQL backup: {target.host}:{target.port}/{target.database} "
            f"-> {final_path} (verify={verify}, retention={retention_days}d)"
        )
        return None

    output_dir.mkdir(parents=True, exist_ok=True)
    partial_path.unlink(missing_ok=True)
    env = target.process_environment()
    try:
        _run(
            [
                pg_dump,
                "--format=custom",
                "--compress=6",
                "--no-owner",
                "--no-privileges",
                "--file",
                str(partial_path),
            ],
            env,
        )
        if not partial_path.exists() or partial_path.stat().st_size == 0:
            raise RuntimeError("pg_dump completed without creating a non-empty archive")

        if verify and pg_restore:
            _run([pg_restore, "--list", str(partial_path)], env, stdout=subprocess.DEVNULL)

        partial_path.replace(final_path)
        manifest = {
            "version": 1,
            "created_at_utc": moment.astimezone(timezone.utc).isoformat(),
            "database_host": target.host,
            "database_name": target.database,
            "format": "postgresql_custom",
            "backup_file": final_path.name,
            "bytes": final_path.stat().st_size,
            "sha256": file_sha256(final_path),
            "verified_with_pg_restore": verify,
        }
        final_path.with_suffix(".manifest.json").write_text(
            json.dumps(manifest, indent=2) + "\n",
            encoding="utf-8",
        )
        removed = prune_expired_backups(output_dir, retention_days, moment)
        print(f"[OK] Backup verified: {final_path} ({manifest['bytes']} bytes)")
        if removed:
            print(f"[OK] Removed {len(removed)} expired backup file(s)")
        return final_path
    except Exception:
        partial_path.unlink(missing_ok=True)
        raise


def main(argv: list[str] | None = None) -> int:
    load_dotenv(BACKEND_DIR / ".env", override=False)
    parser = argparse.ArgumentParser(description="Create and verify a DocDoe PostgreSQL backup.")
    parser.add_argument("--database-url", default=os.getenv("DATABASE_URL"))
    parser.add_argument(
        "--output-dir",
        type=Path,
        default=Path(os.getenv("DATABASE_BACKUP_DIR", str(DEFAULT_OUTPUT_DIR))),
    )
    parser.add_argument(
        "--retention-days",
        type=int,
        default=int(os.getenv("DATABASE_BACKUP_RETENTION_DAYS", "14")),
    )
    parser.add_argument("--skip-verify", action="store_true")
    parser.add_argument("--dry-run", action="store_true")
    args = parser.parse_args(argv)

    try:
        create_backup(
            database_url=args.database_url or "",
            output_dir=args.output_dir,
            retention_days=args.retention_days,
            verify=not args.skip_verify,
            dry_run=args.dry_run,
        )
    except (OSError, RuntimeError, ValueError) as exc:
        print(f"[ERROR] {exc}", file=sys.stderr)
        return 1
    return 0


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