fix: harden backup snapshot integrity

core/backup.py

@@ -3,6 +3,7 @@ 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
@@ -18,7 +19,7 @@ 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_]*\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",
@@ -94,27 +95,25 @@ def create_backup(name=None, *, runner=None, _target=None):
 
     destination = root / name
     manifest_destination = root / f"{name}.json"
-    _ensure_available(destination, manifest_destination)
-
     target = _target or DatabaseTarget.from_default(settings.DATABASES["default"])
-    snapshot = _read_snapshot(target)
     dump_path = None
     manifest_path = None
-    published = []
     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)
-        with os.fdopen(dump_fd, "wb") as stream:
+        os.close(dump_fd)
+        with _exported_snapshot(target) as (snapshot_id, snapshot):
             runner(
-                _pg_dump_argv(target),
+                _pg_dump_argv(target, dump_path, snapshot_id),
                 env=_database_env(target),
                 check=True,
-                stdout=stream,
-                capture_output=False,
+                stdout=None,
+                capture_output=True,
             )
-            stream.flush()
+        with dump_path.open("rb") as stream:
             os.fsync(stream.fileno())
 
         content_digest, byte_size = _file_digest(dump_path)
@@ -148,14 +147,11 @@ def create_backup(name=None, *, runner=None, _target=None):
             stream.flush()
             os.fsync(stream.fileno())
 
-        _publish_no_clobber(manifest_path, manifest_destination)
-        published.append((manifest_destination, manifest_path))
-        _publish_no_clobber(dump_path, destination)
-        published.append((destination, dump_path))
-        _fsync_directory(root)
-        manifest_path.unlink()
+        _publish_pair(
+            root,
+            ((manifest_path, manifest_destination), (dump_path, destination)),
+        )
         manifest_path = None
-        dump_path.unlink()
         dump_path = None
         return BackupResult(
             name=name,
@@ -167,14 +163,16 @@ def create_backup(name=None, *, runner=None, _target=None):
             row_counts=dict(snapshot.row_counts),
         )
     except BackupError:
-        _remove_our_publications(published)
         raise
     except Exception:
-        _remove_our_publications(published)
         raise BackupError("backup_failed") from None
     finally:
-        _unlink_private(manifest_path)
-        _unlink_private(dump_path)
+        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):
@@ -192,6 +190,7 @@ def verify_backup(name, target, *, runner=None):
     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"
@@ -203,8 +202,6 @@ def verify_backup(name, target, *, runner=None):
             raise BackupError("size_mismatch")
         if manifest["sha256"] != digest:
             raise BackupError("digest_mismatch")
-
-    try:
         before = _read_snapshot(target)
         if before.migrations or before.row_counts:
             raise BackupError("target_not_empty")
@@ -229,12 +226,14 @@ def verify_backup(name, target, *, runner=None):
     return VerificationResult(name=name, version=BACKUP_VERSION, row_counts=after.row_counts)
 
 
-def _pg_dump_argv(target):
+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",
@@ -287,8 +286,13 @@ def _validate_name(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):
@@ -300,6 +304,48 @@ def _publish_no_clobber(source, destination):
         raise BackupError("publication_failed") from None
 
 
+def _publish_pair(root, pairs):
+    publications = []
+    guards = []
+    try:
+        for source, destination in pairs:
+            _publish_no_clobber(source, destination)
+            publications.append((destination, source))
+        _fsync_directory(root)
+
+        for source, _destination in pairs:
+            guard = Path(f"{source}.guard")
+            os.link(source, guard)
+            guards.append(guard)
+        for source, _destination in pairs:
+            source.unlink()
+        for guard in guards:
+            guard.unlink()
+        _fsync_directory(root)
+    except Exception:
+        ownership_sources = [*guards, *(source for source, _destination in pairs if source.exists())]
+        _remove_owned_publications(publications, ownership_sources)
+        for source, _destination in pairs:
+            _unlink_private(source)
+        for guard in guards:
+            _unlink_private(guard)
+        raise BackupError("publication_failed") from None
+
+
+def _remove_owned_publications(publications, owned_sources):
+    for destination, original_source in reversed(publications):
+        sources = [source for source in owned_sources if source.exists()]
+        if original_source.exists():
+            sources.append(original_source)
+        try:
+            if destination.exists() and any(
+                os.path.samefile(destination, source) for source in sources
+            ):
+                destination.unlink()
+        except OSError:
+            pass
+
+
 def _remove_our_publications(publications):
     for destination, source in reversed(publications):
         try:
@@ -310,11 +356,13 @@ def _remove_our_publications(publications):
 
 
 def _unlink_private(path):
-    if path is not None:
+    if path is None:
+        return True
     try:
         path.unlink(missing_ok=True)
+        return True
     except OSError:
-            pass
+        return False
 
 
 def _fsync_directory(root):
@@ -340,11 +388,19 @@ def _safe_client_version(output):
     if isinstance(output, bytes):
         output = output.decode("ascii", errors="replace")
     version = str(output).strip()
-    if not version.startswith("pg_dump (PostgreSQL) 17.") or len(version) > 80:
+    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():
@@ -364,9 +420,8 @@ def _load_manifest(path):
         or not isinstance(manifest["created_at"], str)
         or not isinstance(manifest["database_name"], str)
         or not manifest["database_name"]
-        or not isinstance(manifest["postgresql_client_version"], str)
-        or not manifest["postgresql_client_version"].startswith("pg_dump (PostgreSQL) 17.")
-        or not isinstance(manifest["byte_size"], int)
+        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"])
@@ -391,7 +446,7 @@ def _load_manifest(path):
     if any(
         not isinstance(table, str)
         or not _TABLE_NAME.fullmatch(table)
-        or not isinstance(count, int)
+        or type(count) is not int
         or count < 0
         for table, count in manifest["row_counts"].items()
     ):
@@ -399,6 +454,26 @@ def _load_manifest(path):
     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,
@@ -409,27 +484,34 @@ def _read_snapshot(target):
         connect_timeout=3,
     ) as database:
         with database.cursor() as cursor:
+            return _read_snapshot_cursor(cursor)
+
+
+def _read_snapshot_cursor(cursor):
     cursor.execute(
         """
-                SELECT table_name
+        SELECT table_schema, table_name
         FROM information_schema.tables
-                WHERE table_schema = 'public' AND table_type = 'BASE TABLE'
-                ORDER BY table_name
+        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] for row in cursor.fetchall()]
+    tables = [(row[0], row[1]) for row in cursor.fetchall()]
     migrations = ()
-            if "django_migrations" in tables:
-                cursor.execute(
-                    "SELECT app, name FROM django_migrations ORDER BY app, name"
-                )
+    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 table in tables:
+    for schema, table in tables:
         cursor.execute(
-                    sql.SQL("SELECT count(*) FROM {}").format(sql.Identifier(table))
+            sql.SQL("SELECT count(*) FROM {}.{}").format(
+                sql.Identifier(schema), sql.Identifier(table)
+            )
         )
-                row_counts[table] = cursor.fetchone()[0]
+        key = table if schema == "public" else f"{schema}.{table}"
+        row_counts[key] = cursor.fetchone()[0]
     return DatabaseSnapshot(migrations=migrations, row_counts=row_counts)
 
 
@@ -563,14 +645,19 @@ def run_synthetic_restore_verification():
     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)
-        if artifact is not None:
-            artifact.path.unlink(missing_ok=True)
-            artifact.manifest_path.unlink(missing_ok=True)
-        _clear_disposable_database(source_target)
-        _clear_disposable_database(restore_target)
+        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):
@@ -599,5 +686,22 @@ def _clear_disposable_database(target):
             with database.cursor() as cursor:
                 cursor.execute("DROP SCHEMA public CASCADE")
                 cursor.execute("CREATE SCHEMA public")
-    except psycopg.Error:
-        pass
+    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")

tests/test_backup.py

@@ -1,9 +1,14 @@
 import json
+import os
+import subprocess
+from contextlib import contextmanager
 from hashlib import sha256
 from pathlib import Path
 from types import SimpleNamespace
 
 import pytest
+from django.db import connection
+from psycopg import connect, sql
 
 from core.backup import (
     BACKUP_FORMAT,
@@ -23,8 +28,9 @@ class RecordingRunner:
 
     def __call__(self, argv, *, env, check, stdout=None, capture_output=False):
         self.calls.append((list(argv), dict(env), check, stdout, capture_output))
-        if stdout is not None:
-            stdout.write(self.output)
+        if argv[0] == "pg_dump" and argv != ["pg_dump", "--version"]:
+            output_name = next(value.removeprefix("--file=") for value in argv if value.startswith("--file="))
+            Path(output_name).write_bytes(self.output)
         return SimpleNamespace(stdout=b"pg_dump (PostgreSQL) 17.6\n")
 
 
@@ -53,6 +59,11 @@ def synthetic_snapshot(monkeypatch):
         migrations=(("finance", "0001_initial"), ("ingestion", "0001_initial")),
         row_counts={"finance_account": 2, "ingestion_importbatch": 1},
     )
+    @contextmanager
+    def exported_snapshot(_target):
+        yield "synthetic-snapshot", snapshot
+
+    monkeypatch.setattr("core.backup._exported_snapshot", exported_snapshot)
     monkeypatch.setattr("core.backup._read_snapshot", lambda _target: snapshot)
     return snapshot
 
@@ -69,11 +80,22 @@ def test_pg_dump_receives_password_only_in_environment(
     assert database_settings["PASSWORD"] not in argv
     assert env["PGPASSWORD"] == database_settings["PASSWORD"]
     assert check is True
-    assert stdout is not None
-    assert capture_output is False
+    assert stdout is None
+    assert capture_output is True
     assert result.path.read_bytes() == runner.output
 
 
+def test_pg_dump_uses_the_exported_snapshot(
+    backup_root, database_settings, synthetic_snapshot
+):
+    runner = RecordingRunner()
+
+    create_backup("synthetic.dump", runner=runner)
+
+    assert "--snapshot=synthetic-snapshot" in runner.calls[0][0]
+    assert runner.calls[0][4] is True
+
+
 def test_backup_writes_versioned_checksum_manifest(
     backup_root, database_settings, synthetic_snapshot
 ):
@@ -142,6 +164,81 @@ def test_backup_cleans_private_temporary_files_on_dump_failure(
     assert list(backup_root.iterdir()) == []
 
 
+def test_backup_sanitizes_subprocess_failure_and_captures_diagnostics(
+    backup_root, database_settings, synthetic_snapshot
+):
+    def failing_runner(argv, *, capture_output, **kwargs):
+        assert capture_output is True
+        raise subprocess.CalledProcessError(
+            1, argv, output=b"SYNTHETIC PRIVATE STDOUT", stderr=b"SYNTHETIC PRIVATE STDERR"
+        )
+
+    with pytest.raises(BackupError, match="backup_failed") as error:
+        create_backup("synthetic.dump", runner=failing_runner)
+
+    assert "PRIVATE" not in str(error.value)
+    assert list(backup_root.iterdir()) == []
+
+
+def test_second_temporary_unlink_failure_leaves_no_partial_final_pair(
+    backup_root, database_settings, synthetic_snapshot, monkeypatch
+):
+    original_unlink = Path.unlink
+    failed = False
+
+    def fail_once(path, *args, **kwargs):
+        nonlocal failed
+        if not failed and path.name.endswith(".dump.tmp"):
+            failed = True
+            raise OSError("synthetic unlink failure")
+        return original_unlink(path, *args, **kwargs)
+
+    monkeypatch.setattr(Path, "unlink", fail_once)
+
+    with pytest.raises(BackupError, match="publication_failed"):
+        create_backup("synthetic.dump", runner=RecordingRunner())
+
+    assert list(backup_root.iterdir()) == []
+
+
+def test_publication_never_deletes_concurrent_final_files(tmp_path, monkeypatch):
+    import core.backup as backup
+
+    manifest_temp = tmp_path / ".manifest.tmp"
+    dump_temp = tmp_path / ".dump.tmp"
+    manifest_final = tmp_path / "synthetic.dump.json"
+    dump_final = tmp_path / "synthetic.dump"
+    manifest_temp.write_bytes(b"our manifest")
+    dump_temp.write_bytes(b"our dump")
+    original_link = os.link
+    calls = 0
+
+    def competing_link(source, destination):
+        nonlocal calls
+        calls += 1
+        if calls == 2:
+            manifest_final.unlink()
+            manifest_final.write_bytes(b"other manifest")
+            dump_final.write_bytes(b"other dump")
+            raise FileExistsError
+        return original_link(source, destination)
+
+    monkeypatch.setattr(os, "link", competing_link)
+
+    with pytest.raises(BackupError, match="publication_failed"):
+        backup._publish_pair(
+            tmp_path,
+            ((manifest_temp, manifest_final), (dump_temp, dump_final)),
+        )
+
+    assert manifest_final.read_bytes() == b"other manifest"
+    assert dump_final.read_bytes() == b"other dump"
+    assert sorted(path.name for path in tmp_path.iterdir()) == [
+        "synthetic.dump",
+        "synthetic.dump.json",
+    ]
+
+
 def test_verification_rejects_live_database(
     backup_root, database_settings, synthetic_snapshot
 ):
@@ -188,6 +285,25 @@ def test_verification_rejects_changed_dump_before_restore(
     assert restore_runner.calls == []
 
 
+def test_verification_sanitizes_digest_read_failure(
+    backup_root, database_settings, synthetic_snapshot, settings, monkeypatch
+):
+    settings.BACKUP_VERIFY_ALLOWED = True
+    result = create_backup("synthetic.dump", runner=RecordingRunner())
+    monkeypatch.setattr(
+        "core.backup._file_digest",
+        lambda _path: (_ for _ in ()).throw(OSError("SYNTHETIC PRIVATE PATH")),
+    )
+    target = DatabaseTarget(
+        host="restore-db", port="5432", name="restore", user="user", password="secret"
+    )
+
+    with pytest.raises(BackupError, match="verification_failed") as error:
+        verify_backup(result.name, target, runner=RecordingRunner())
+
+    assert "PRIVATE" not in str(error.value)
+
+
 def test_verification_rejects_manifest_with_wrong_field_type(
     backup_root, database_settings, synthetic_snapshot, settings
 ):
@@ -258,3 +374,130 @@ def test_pg_restore_receives_password_only_in_environment(
     assert stdout is None
     assert capture_output is True
     assert verification.row_counts == synthetic_snapshot.row_counts
+
+
+@pytest.mark.django_db(transaction=True)
+def test_snapshot_counts_tables_in_non_system_schemas(settings):
+    from core.backup import _read_snapshot
+
+    target = DatabaseTarget.from_default(settings.DATABASES["default"])
+    with connection.cursor() as cursor:
+        cursor.execute("CREATE SCHEMA synthetic_extra")
+        cursor.execute("CREATE TABLE synthetic_extra.marker (id integer)")
+        cursor.execute("INSERT INTO synthetic_extra.marker VALUES (1)")
+    try:
+        snapshot = _read_snapshot(target)
+        assert snapshot.row_counts["synthetic_extra.marker"] == 1
+    finally:
+        with connection.cursor() as cursor:
+            cursor.execute("DROP SCHEMA synthetic_extra CASCADE")
+
+
+def test_manifest_rejects_boolean_integer_fields(
+    backup_root, database_settings, synthetic_snapshot, settings
+):
+    settings.BACKUP_VERIFY_ALLOWED = True
+    result = create_backup("synthetic.dump", runner=RecordingRunner())
+    manifest = json.loads(result.manifest_path.read_text(encoding="utf-8"))
+    manifest["byte_size"] = True
+    result.manifest_path.write_text(json.dumps(manifest), encoding="utf-8")
+    target = DatabaseTarget(
+        host="restore-db", port="5432", name="restore", user="user", password="secret"
+    )
+
+    with pytest.raises(BackupError, match="invalid_manifest"):
+        verify_backup(result.name, target, runner=RecordingRunner())
+
+
+def test_cleanup_attempts_every_resource_and_reports_safe_failure(monkeypatch, tmp_path):
+    import core.backup as backup
+
+    attempts = []
+    dump = tmp_path / "synthetic.dump"
+    manifest = tmp_path / "synthetic.dump.json"
+    dump.write_bytes(b"dump")
+    manifest.write_text("{}", encoding="utf-8")
+    artifact = SimpleNamespace(path=dump, manifest_path=manifest)
+    source = DatabaseTarget("source", "5432", "source", "user", "password")
+    target = DatabaseTarget("target", "5432", "target", "user", "password")
+
+    original_unlink = Path.unlink
+
+    def unlink(path, *args, **kwargs):
+        attempts.append(f"unlink:{path.name}")
+        if path == dump:
+            raise OSError("private path")
+        return original_unlink(path, *args, **kwargs)
+
+    def clear(database):
+        attempts.append(f"clear:{database.name}")
+        if database == source:
+            raise BackupError("cleanup_failed")
+
+    monkeypatch.setattr(Path, "unlink", unlink)
+    monkeypatch.setattr(backup, "_clear_disposable_database", clear)
+
+    with pytest.raises(BackupError, match="cleanup_failed"):
+        backup._cleanup_synthetic_resources(artifact, source, target)
+
+    assert attempts == [
+        "unlink:synthetic.dump",
+        "unlink:synthetic.dump.json",
+        "clear:source",
+        "clear:target",
+    ]
+
+
+@pytest.mark.django_db(transaction=True)
+def test_concurrent_insert_is_excluded_from_both_manifest_and_dump(tmp_path, settings):
+    from finance.models import Account
+
+    settings.BACKUP_ROOT = tmp_path
+    settings.BACKUP_VERIFY_ALLOWED = True
+    source = DatabaseTarget.from_default(settings.DATABASES["default"])
+    Account.objects.create(name="Synthetic Snapshot First", kind="BANK", currency="BRL")
+    inserted = False
+
+    def runner(argv, **kwargs):
+        nonlocal inserted
+        if argv[0] == "pg_dump" and argv != ["pg_dump", "--version"]:
+            assert any(value.startswith("--snapshot=") for value in argv)
+            Account.objects.create(
+                name="Synthetic Snapshot Concurrent", kind="BANK", currency="BRL"
+            )
+            inserted = True
+        return subprocess.run(argv, **kwargs)
+
+    database_name = f"snapshot_restore_{os.getpid()}"
+    admin = dict(
+        host=source.host,
+        port=source.port,
+        dbname="postgres",
+        user=source.user,
+        password=source.password,
+        autocommit=True,
+    )
+    with connect(**admin) as database:
+        database.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(database_name)))
+    target = DatabaseTarget(
+        host=source.host,
+        port=source.port,
+        name=database_name,
+        user=source.user,
+        password=source.password,
+    )
+    try:
+        result = create_backup("concurrent.dump", runner=runner)
+        verification = verify_backup(result.name, target)
+
+        assert inserted is True
+        assert result.row_counts["finance_account"] == 1
+        assert verification.row_counts["finance_account"] == 1
+        assert Account.objects.count() == 2
+    finally:
+        with connect(**admin) as database:
+            database.execute(
+                sql.SQL("DROP DATABASE {} WITH (FORCE)").format(
+                    sql.Identifier(database_name)
+                )
+            )