gizmoduck/core/backup.py

import json
import os
import re
import subprocess
import tempfile
from contextlib import contextmanager
from dataclasses import dataclass
from datetime import datetime, timezone as datetime_timezone
from hashlib import sha256
from pathlib import Path

import psycopg
from psycopg import sql
from django.conf import settings


BACKUP_FORMAT = "gizmoduck-postgresql-backup"
BACKUP_VERSION = 1
APPLICATION_VERSION = "0.1.0"

_PLAIN_DUMP_NAME = re.compile(r"\A[A-Za-z0-9][A-Za-z0-9._-]{0,198}\.dump\Z")
_TABLE_NAME = re.compile(r"\A[A-Za-z_][A-Za-z0-9_]*(?:\.[A-Za-z_][A-Za-z0-9_]*)?\Z")
_MANIFEST_KEYS = {
    "application_version",
    "byte_size",
    "created_at",
    "database_name",
    "format",
    "migrations",
    "postgresql_client_version",
    "row_counts",
    "sha256",
    "version",
}


class BackupError(Exception):
    def __init__(self, code):
        self.code = code
        super().__init__(code)


@dataclass(frozen=True)
class DatabaseTarget:
    host: str
    port: str
    name: str
    user: str
    password: str

    @classmethod
    def from_default(cls, database):
        return cls(
            host=str(database.get("HOST", "")),
            port=str(database.get("PORT") or "5432"),
            name=str(database.get("NAME", "")),
            user=str(database.get("USER", "")),
            password=str(database.get("PASSWORD", "")),
        )


@dataclass(frozen=True)
class DatabaseSnapshot:
    migrations: tuple[tuple[str, str], ...]
    row_counts: dict[str, int]


@dataclass(frozen=True)
class BackupResult:
    name: str
    path: Path
    manifest_path: Path
    version: int
    created_at: datetime
    postgresql_client_version: str
    row_counts: dict[str, int]


@dataclass(frozen=True)
class VerificationResult:
    name: str
    version: int
    row_counts: dict[str, int]


def create_backup(name=None, *, runner=None, _target=None):
    runner = runner or subprocess.run
    root = _backup_root()
    if name is None:
        now = datetime.now(datetime_timezone.utc)
        name = f"gizmoduck-backup-{now:%Y%m%dT%H%M%SZ}.dump"
    else:
        now = datetime.now(datetime_timezone.utc)
    _validate_name(name)

    destination = root / name
    manifest_destination = root / f"{name}.json"
    target = _target or DatabaseTarget.from_default(settings.DATABASES["default"])
    dump_path = None
    manifest_path = None
    try:
        _ensure_available(destination, manifest_destination)
        dump_fd, dump_name = tempfile.mkstemp(
            prefix=".gizmoduck-backup-", suffix=".dump.tmp", dir=root
        )
        dump_path = Path(dump_name)
        os.close(dump_fd)
        with _exported_snapshot(target) as (snapshot_id, snapshot):
            runner(
                _pg_dump_argv(target, dump_path, snapshot_id),
                env=_database_env(target),
                check=True,
                stdout=None,
                capture_output=True,
            )
        with dump_path.open("rb") as stream:
            os.fsync(stream.fileno())

        content_digest, byte_size = _file_digest(dump_path)
        version_process = runner(
            ["pg_dump", "--version"],
            env=_database_env(target),
            check=True,
            stdout=None,
            capture_output=True,
        )
        client_version = _safe_client_version(getattr(version_process, "stdout", b""))
        manifest = {
            "application_version": APPLICATION_VERSION,
            "byte_size": byte_size,
            "created_at": now.isoformat(),
            "database_name": target.name,
            "format": BACKUP_FORMAT,
            "migrations": [list(item) for item in snapshot.migrations],
            "postgresql_client_version": client_version,
            "row_counts": snapshot.row_counts,
            "sha256": content_digest,
            "version": BACKUP_VERSION,
        }
        manifest_fd, manifest_name = tempfile.mkstemp(
            prefix=".gizmoduck-backup-", suffix=".json.tmp", dir=root
        )
        manifest_path = Path(manifest_name)
        with os.fdopen(manifest_fd, "w", encoding="utf-8", newline="\n") as stream:
            json.dump(manifest, stream, ensure_ascii=False, separators=(",", ":"))
            stream.write("\n")
            stream.flush()
            os.fsync(stream.fileno())

        _publish_pair(
            root,
            ((manifest_path, manifest_destination), (dump_path, destination)),
        )
        manifest_path = None
        dump_path = None
        return BackupResult(
            name=name,
            path=destination,
            manifest_path=manifest_destination,
            version=BACKUP_VERSION,
            created_at=now,
            postgresql_client_version=client_version,
            row_counts=dict(snapshot.row_counts),
        )
    except BackupError:
        raise
    except Exception:
        raise BackupError("backup_failed") from None
    finally:
        cleanup_failed = False
        for path in (manifest_path, dump_path):
            if not _unlink_private(path):
                cleanup_failed = True
        if cleanup_failed:
            raise BackupError("cleanup_failed") from None


def verify_backup(name, target, *, runner=None):
    runner = runner or subprocess.run
    _validate_name(name)
    live = DatabaseTarget(
        host=str(getattr(settings, "LIVE_DATABASE_HOST", None) or settings.DATABASES["default"].get("HOST", "")),
        port=str(settings.DATABASES["default"].get("PORT") or "5432"),
        name=str(getattr(settings, "LIVE_DATABASE_NAME", None) or settings.DATABASES["default"].get("NAME", "")),
        user="",
        password="",
    )
    if _same_database(live, target):
        raise BackupError("live_target")
    if not getattr(settings, "BACKUP_VERIFY_ALLOWED", False):
        raise BackupError("verification_not_allowed")

    try:
        root = _backup_root()
        dump_path = root / name
        manifest_path = root / f"{name}.json"
        manifest = _load_manifest(manifest_path)
        if not dump_path.is_file() or dump_path.is_symlink():
            raise BackupError("missing_backup")
        digest, byte_size = _file_digest(dump_path)
        if manifest["byte_size"] != byte_size:
            raise BackupError("size_mismatch")
        if manifest["sha256"] != digest:
            raise BackupError("digest_mismatch")
        before = _read_snapshot(target)
        if before.migrations or before.row_counts:
            raise BackupError("target_not_empty")
        runner(
            _pg_restore_argv(target, dump_path),
            env=_database_env(target),
            check=True,
            stdout=None,
            capture_output=True,
        )
        after = _read_snapshot(target)
    except BackupError:
        raise
    except Exception:
        raise BackupError("verification_failed") from None

    expected_migrations = tuple(tuple(item) for item in manifest["migrations"])
    if after.migrations != expected_migrations:
        raise BackupError("migration_mismatch")
    if after.row_counts != manifest["row_counts"]:
        raise BackupError("row_count_mismatch")
    return VerificationResult(name=name, version=BACKUP_VERSION, row_counts=after.row_counts)


def _pg_dump_argv(target, dump_path, snapshot_id):
    return [
        "pg_dump",
        "--format=custom",
        "--no-owner",
        "--no-acl",
        f"--snapshot={snapshot_id}",
        f"--file={dump_path}",
        "--host",
        target.host,
        "--port",
        target.port,
        "--username",
        target.user,
        "--dbname",
        target.name,
    ]


def _pg_restore_argv(target, dump_path):
    return [
        "pg_restore",
        "--exit-on-error",
        "--single-transaction",
        "--no-owner",
        "--no-acl",
        "--host",
        target.host,
        "--port",
        target.port,
        "--username",
        target.user,
        "--dbname",
        target.name,
        str(dump_path),
    ]


def _database_env(target):
    env = os.environ.copy()
    env["PGPASSWORD"] = target.password
    return env


def _backup_root():
    try:
        root = Path(settings.BACKUP_ROOT).resolve(strict=True)
    except (AttributeError, OSError, RuntimeError, TypeError):
        raise BackupError("invalid_root") from None
    if not root.is_dir():
        raise BackupError("invalid_root")
    return root


def _validate_name(name):
    if not isinstance(name, str) or not _PLAIN_DUMP_NAME.fullmatch(name):
        raise BackupError("invalid_name")


def _ensure_available(*paths):
    try:
        if any(path.exists() or path.is_symlink() for path in paths):
            raise BackupError("existing_destination")
    except BackupError:
        raise
    except Exception:
        raise BackupError("publication_failed") from None


def _publish_no_clobber(source, destination):
    try:
        os.link(source, destination)
    except FileExistsError:
        raise BackupError("existing_destination") from None
    except OSError:
        raise BackupError("publication_failed") from None


def _publish_pair(root, pairs):
    publications = []
    try:
        for source, destination in pairs:
            _publish_no_clobber(source, destination)
            publications.append((destination, source))
        _fsync_directory(root)
    except Exception:
        _remove_owned_publications(publications)
        for source, _destination in pairs:
            _unlink_private(source)
        raise BackupError("publication_failed") from None

    cleanup_failed = False
    for source, _destination in pairs:
        if not _unlink_private(source):
            cleanup_failed = True
    if cleanup_failed:
        raise BackupError("cleanup_failed")


def _remove_owned_publications(publications):
    for destination, source in reversed(publications):
        try:
            if source.exists() and destination.exists() and os.path.samefile(
                destination, source
            ):
                destination.unlink()
        except OSError:
            pass


def _unlink_private(path):
    if path is None:
        return True
    try:
        path.unlink(missing_ok=True)
        return True
    except OSError:
        return False


def _fsync_directory(root):
    flags = os.O_RDONLY | getattr(os, "O_DIRECTORY", 0)
    descriptor = os.open(root, flags)
    try:
        os.fsync(descriptor)
    finally:
        os.close(descriptor)


def _file_digest(path):
    digest = sha256()
    size = 0
    with path.open("rb") as stream:
        while chunk := stream.read(1024 * 1024):
            digest.update(chunk)
            size += len(chunk)
    return digest.hexdigest(), size


def _safe_client_version(output):
    if isinstance(output, bytes):
        output = output.decode("ascii", errors="replace")
    version = str(output).strip()
    if not _is_pg17_version(version):
        raise BackupError("client_version_failed")
    return version


def _is_pg17_version(version):
    return (
        isinstance(version, str)
        and version.startswith("pg_dump (PostgreSQL) 17.")
        and len(version) <= 80
    )


def _load_manifest(path):
    try:
        if path.is_symlink():
            raise BackupError("invalid_manifest")
        manifest = json.loads(path.read_text(encoding="utf-8"))
    except BackupError:
        raise
    except (OSError, UnicodeError, json.JSONDecodeError):
        raise BackupError("invalid_manifest") from None
    if not isinstance(manifest, dict) or set(manifest) != _MANIFEST_KEYS:
        raise BackupError("invalid_manifest")
    if (
        manifest["format"] != BACKUP_FORMAT
        or type(manifest["version"]) is not int
        or manifest["version"] != BACKUP_VERSION
    ):
        raise BackupError("invalid_manifest")
    if (
        not isinstance(manifest["application_version"], str)
        or not manifest["application_version"]
        or not isinstance(manifest["created_at"], str)
        or not isinstance(manifest["database_name"], str)
        or not manifest["database_name"]
        or not _is_pg17_version(manifest["postgresql_client_version"])
        or type(manifest["byte_size"]) is not int
        or manifest["byte_size"] < 0
        or not isinstance(manifest["sha256"], str)
        or not re.fullmatch(r"[0-9a-f]{64}", manifest["sha256"])
        or not isinstance(manifest["migrations"], list)
        or not isinstance(manifest["row_counts"], dict)
    ):
        raise BackupError("invalid_manifest")
    try:
        created_at = datetime.fromisoformat(manifest["created_at"])
    except ValueError:
        raise BackupError("invalid_manifest") from None
    if created_at.tzinfo is None:
        raise BackupError("invalid_manifest")
    migrations = manifest["migrations"]
    if any(
        not isinstance(item, list)
        or len(item) != 2
        or not all(isinstance(value, str) and value for value in item)
        for item in migrations
    ):
        raise BackupError("invalid_manifest")
    if any(
        not isinstance(table, str)
        or not _TABLE_NAME.fullmatch(table)
        or type(count) is not int
        or count < 0
        for table, count in manifest["row_counts"].items()
    ):
        raise BackupError("invalid_manifest")
    return manifest


@contextmanager
def _exported_snapshot(target):
    with psycopg.connect(
        host=target.host,
        port=target.port,
        dbname=target.name,
        user=target.user,
        password=target.password,
        connect_timeout=3,
    ) as database:
        with database.cursor() as cursor:
            cursor.execute(
                "SET TRANSACTION ISOLATION LEVEL REPEATABLE READ READ ONLY"
            )
            cursor.execute("SELECT pg_export_snapshot()")
            snapshot_id = cursor.fetchone()[0]
            snapshot = _read_snapshot_cursor(cursor)
            yield snapshot_id, snapshot


def _read_snapshot(target):
    with psycopg.connect(
        host=target.host,
        port=target.port,
        dbname=target.name,
        user=target.user,
        password=target.password,
        connect_timeout=3,
    ) as database:
        with database.cursor() as cursor:
            return _read_snapshot_cursor(cursor)


def _read_snapshot_cursor(cursor):
    cursor.execute(
        """
        SELECT table_schema, table_name
        FROM information_schema.tables
        WHERE table_type = 'BASE TABLE'
          AND table_schema NOT IN ('pg_catalog', 'information_schema', 'pg_toast')
          AND table_schema NOT LIKE 'pg\\_%' ESCAPE '\\'
        ORDER BY table_schema, table_name
        """
    )
    tables = [(row[0], row[1]) for row in cursor.fetchall()]
    migrations = ()
    if ("public", "django_migrations") in tables:
        cursor.execute("SELECT app, name FROM django_migrations ORDER BY app, name")
        migrations = tuple((row[0], row[1]) for row in cursor.fetchall())
    row_counts = {}
    for schema, table in tables:
        cursor.execute(
            sql.SQL("SELECT count(*) FROM {}.{}").format(
                sql.Identifier(schema), sql.Identifier(table)
            )
        )
        key = table if schema == "public" else f"{schema}.{table}"
        row_counts[key] = cursor.fetchone()[0]
    return DatabaseSnapshot(migrations=migrations, row_counts=row_counts)


def _same_database(first, second):
    return (
        first.host.strip().casefold(),
        first.name.strip().casefold(),
    ) == (
        second.host.strip().casefold(),
        second.name.strip().casefold(),
    )


def run_synthetic_restore_verification():
    """Exercise backup and restore only against disposable synthetic databases."""
    from datetime import date
    from decimal import Decimal
    from uuid import uuid4

    from django.core.management import call_command
    from django.db import connections

    from finance.models import Account, AuditEvent, Category, Payee, SourceRecord, Transaction
    from ingestion.models import ImportBatch, ImportItem

    if not getattr(settings, "BACKUP_VERIFY_ALLOWED", False):
        raise BackupError("verification_not_allowed")
    required = (
        "BACKUP_SOURCE_HOST",
        "BACKUP_SOURCE_NAME",
        "BACKUP_SOURCE_USER",
        "BACKUP_SOURCE_PASSWORD",
    )
    if any(not os.environ.get(name) for name in required):
        raise BackupError("source_not_configured")

    restore_target = DatabaseTarget.from_default(settings.DATABASES["default"])
    source_target = DatabaseTarget(
        host=os.environ["BACKUP_SOURCE_HOST"],
        port=os.environ.get("BACKUP_SOURCE_PORT", "5432"),
        name=os.environ["BACKUP_SOURCE_NAME"],
        user=os.environ["BACKUP_SOURCE_USER"],
        password=os.environ["BACKUP_SOURCE_PASSWORD"],
    )
    live_target = DatabaseTarget(
        host=str(getattr(settings, "LIVE_DATABASE_HOST", None) or ""),
        port="5432",
        name=str(getattr(settings, "LIVE_DATABASE_NAME", None) or ""),
        user="",
        password="",
    )
    if _same_database(source_target, restore_target):
        raise BackupError("source_is_target")
    if live_target.host and (
        _same_database(source_target, live_target)
        or _same_database(restore_target, live_target)
    ):
        raise BackupError("live_target")
    try:
        if _read_snapshot(source_target).row_counts:
            raise BackupError("source_not_empty")
        restore_before = _read_snapshot(restore_target)
        if restore_before.migrations or restore_before.row_counts:
            raise BackupError("target_not_empty")
    except BackupError:
        raise
    except Exception:
        raise BackupError("integration_verification_failed") from None

    default_connection = connections["default"]
    original_database = dict(default_connection.settings_dict)
    artifact = None
    try:
        default_connection.close()
        _configure_connection(default_connection, source_target)
        call_command("migrate", interactive=False, verbosity=0)

        account = Account.objects.create(
            name="Synthetic Restore Bank", kind=Account.Kind.BANK, currency="BRL"
        )
        payee = Payee.objects.create(name="Synthetic Restore Merchant")
        category = Category.objects.create(
            name="Synthetic Restore Expense", kind=Category.Kind.EXPENSE
        )
        transaction = Transaction.objects.create(
            account=account,
            amount=Decimal("-12.34"),
            posting_date=date(2026, 9, 8),
            description="Synthetic restore purchase",
            kind=Transaction.Kind.EXPENSE,
            payee=payee,
            category=category,
        )
        SourceRecord.objects.create(
            transaction=transaction,
            document_hash="a" * 64,
            source_format=SourceRecord.Format.BANK_CSV,
            parser_version="synthetic-v1",
            item_locator="row:1",
        )
        AuditEvent.objects.create(
            actor_kind=AuditEvent.ActorKind.AGENT,
            actor_identifier="synthetic-backup-verification",
            operation="synthetic_backup_verification",
            target_id=transaction.id,
            outcome=AuditEvent.Outcome.SUCCESS,
        )
        batch = ImportBatch.objects.create(
            document_hash="b" * 64,
            source_format="BANK_CSV",
            profile="SYNTHETIC",
            encoding="utf-8",
            actor_kind="AGENT",
            actor_identifier="synthetic-backup-verification",
        )
        ImportItem.objects.create(
            batch=batch,
            locator="row:1",
            candidate={"description": "Synthetic restore candidate"},
        )
        before = _read_snapshot(source_target)
        artifact = create_backup(
            f"synthetic-restore-check-{uuid4().hex}.dump", _target=source_target
        )

        default_connection.close()
        default_connection.settings_dict.clear()
        default_connection.settings_dict.update(original_database)
        result = verify_backup(artifact.name, restore_target)
        call_command("check", databases=["default"], verbosity=0)
        if _read_snapshot(source_target) != before:
            raise BackupError("source_changed")
        return result
    except BackupError:
        raise
    except Exception:
        raise BackupError("integration_verification_failed") from None
    finally:
        cleanup_failed = False
        try:
            default_connection.close()
            default_connection.settings_dict.clear()
            default_connection.settings_dict.update(original_database)
        except Exception:
            cleanup_failed = True
        try:
            _cleanup_synthetic_resources(artifact, source_target, restore_target)
        except BackupError:
            cleanup_failed = True
        if cleanup_failed:
            raise BackupError("cleanup_failed") from None


def _configure_connection(connection, target):
    connection.settings_dict.update(
        {
            "HOST": target.host,
            "PORT": target.port,
            "NAME": target.name,
            "USER": target.user,
            "PASSWORD": target.password,
        }
    )


def _clear_disposable_database(target):
    try:
        with psycopg.connect(
            host=target.host,
            port=target.port,
            dbname=target.name,
            user=target.user,
            password=target.password,
            connect_timeout=3,
            autocommit=True,
        ) as database:
            with database.cursor() as cursor:
                cursor.execute("SET lock_timeout = '5s'")
                cursor.execute("SET statement_timeout = '15s'")
                cursor.execute("DROP SCHEMA public CASCADE")
                cursor.execute("CREATE SCHEMA public")
    except Exception:
        raise BackupError("cleanup_failed") from None


def _cleanup_synthetic_resources(artifact, source_target, restore_target):
    failed = False
    if artifact is not None:
        for path in (artifact.path, artifact.manifest_path):
            try:
                path.unlink(missing_ok=True)
            except Exception:
                failed = True
    for target in (source_target, restore_target):
        try:
            _clear_disposable_database(target)
        except BackupError:
            failed = True
    if failed:
        raise BackupError("cleanup_failed")