src/services/brain_backup/postgres.py

"""Logical PostgreSQL snapshots with an explicit content-table allowlist."""
import hashlib
import json
import os
import subprocess
from contextlib import contextmanager
from pathlib import Path

import psycopg2
from psycopg2 import sql
from psycopg2.extras import Json

from .git import BackupError

TABLES = frozenset({
    "data_text_chunks", "data_observations", "data_structured_data", "data_kg_changes",
    "kg_nodes", "kg_relationships", "kg_hub_bridges", "kg_topic_sessions", "embedding_profiles",
    "vectors_nodes", "vectors_triplets", "vectors_observations", "vectors_data", "vectors_relationships",
})
STORES = {"nodes", "triplets", "observations", "data", "relationships"}


def canonical(value):
    return json.dumps(value, sort_keys=True, separators=(",", ":"), ensure_ascii=False).encode()


class PostgreSQL:
    def __init__(self, settings, profile):
        self.settings = settings
        self.profile = profile
        self.system = settings.get("system_database") or settings["database"]

    @classmethod
    def configured(cls):
        from src.config import config
        from src.lib.embeddings.profile import index_profile
        if config.data_db != "postgresql" or config.vector_db != "postgresql" or config.graph_db != "networkx":
            raise BackupError("BACKUP_REQUIRES_POSTGRESQL_NETWORKX")
        from src.constants.embeddings import EMBEDDING_STORES_SIZES
        if set(EMBEDDING_STORES_SIZES.values()) != {512}:
            raise BackupError("BACKUP_REQUIRES_UNIFORM_512_DIMENSIONS")
        p = config.postgresql
        return cls(dict(host=p.host, port=p.port, user=p.username, password=p.password,
                        database=p.database, system_database=getattr(p, "system_database", None)), index_profile(config, 512))

    @contextmanager
    def connect(self, database=None, autocommit=False):
        args = {k: v for k, v in self.settings.items() if k not in {"database", "system_database"}}
        conn = psycopg2.connect(**args, dbname=database or self.system)
        try:
            conn.autocommit = autocommit
            yield conn
        finally:
            conn.close()

    def tool(self, program, database, *args, stdout=None):
        env = os.environ.copy()
        env.update(PGHOST=str(self.settings["host"]), PGPORT=str(self.settings["port"]),
                   PGUSER=self.settings["user"], PGPASSWORD=self.settings["password"], PGDATABASE=database)
        try:
            result = subprocess.run([program, *args], env=env, stdout=stdout or subprocess.PIPE,
                                    stderr=subprocess.PIPE, timeout=120)
        except (OSError, subprocess.TimeoutExpired):
            raise BackupError("BACKUP_POSTGRES_TOOL_UNAVAILABLE") from None
        if result.returncode:
            raise BackupError("BACKUP_POSTGRES_TOOL_FAILED")
        return result.stdout

    def major_version(self):
        with self.connect() as conn:
            return conn.server_version // 10000

    def inspect_secret_stream(self, stream):
        # Known deployment credentials and minted Brain PATs must never enter
        # registry metadata or content archives, even if pasted into documents.
        import re
        candidates = [os.getenv(key, "").encode() for key in (
            "SYSTEM_PAT", "BRAINPAT_TOKEN", "GITHUB_PAT", "HF_TOKEN",
            "POSTGRES_PASSWORD", "REDIS_PASSWORD")]
        candidates = [v for v in candidates if len(v) >= 8]
        tail = b""
        overlap = max([len(v) for v in candidates] + [256])
        for chunk in iter(lambda: stream.read(1024 * 1024), b""):
            data = tail + chunk
            if any(secret in data for secret in candidates) or re.search(rb"\bbpat_[A-Za-z0-9_-]{16,}", data):
                raise BackupError("BACKUP_SECRET_IN_CONTENT")
            tail = data[-overlap:]

    def inspect_secrets(self, archive, database):
        # Spool logical SQL rather than retain large plaintext in memory.
        import tempfile
        with tempfile.TemporaryFile() as stream:
            self.tool("pg_restore", database, "--file=-", str(archive), stdout=stream)
            stream.seek(0)
            self.inspect_secret_stream(stream)

    def registry(self):
        from src.constants.data import Brain, Workspace
        from src.services.brain_identity import validate_brain_id
        from src.lib.postgresql._naming import brain_db_name
        with self.connect() as conn, conn.cursor() as cur:
            cur.execute("SELECT id,name_key,document,created_at::text FROM data_brains ORDER BY name_key")
            brains = []
            for identity, name, document, created in cur.fetchall():
                validate_brain_id(name)
                # Explicit whitelist, never export legacy PAT/document credentials.
                if document.get("id") != identity or document.get("name_key") != name:
                    raise BackupError("BACKUP_REGISTRY_IDENTITY_MISMATCH")
                brains.append({"id": identity, "name_key": name, "created_at": created,
                               "document": Brain(id=identity, name_key=name).model_dump(mode="json")})
            physical = [brain_db_name(b["name_key"]) for b in brains]
            if len(set(physical)) != len(physical):
                raise BackupError("BACKUP_DATABASE_NAME_COLLISION")
            cur.execute("SELECT id,brain_id,slug,document,created_at::text FROM data_workspaces ORDER BY id")
            workspaces = []
            names = {b["name_key"] for b in brains}
            for identity, name, slug, document, created in cur.fetchall():
                if name not in names:
                    continue
                obj = Workspace.model_validate(document)
                if (obj.id, obj.brain_id, obj.slug) != (identity, name, slug):
                    raise BackupError("BACKUP_WORKSPACE_IDENTITY_MISMATCH")
                workspaces.append({"id": identity, "brain_id": name, "slug": slug,
                                   "document": obj.model_dump(mode="json"), "created_at": created})
            return {"brains": brains, "workspaces": workspaces}

    def tables(self, cur):
        cur.execute("SELECT tablename FROM pg_tables WHERE schemaname='public' ORDER BY tablename")
        tables = [r[0] for r in cur.fetchall()]
        # Unknown tables could contain credentials OR unmodelled content. Refuse
        # instead of silently losing it. Never discover unregistered databases.
        if set(tables) - TABLES:
            raise BackupError("BACKUP_UNKNOWN_BRAIN_TABLE")
        return tables

    def validate_profile(self, cur, tables):
        profiles = {}
        if "embedding_profiles" in tables:
            cur.execute("SELECT store,profile FROM embedding_profiles ORDER BY store")
            profiles = dict(cur.fetchall())
            if set(profiles) - STORES or any(p != self.profile for p in profiles.values()):
                raise BackupError("BACKUP_EMBEDDING_PROFILE_MISMATCH")
        for table in tables:
            if table.startswith("vectors_"):
                store = table.removeprefix("vectors_")
                cur.execute("SELECT atttypmod FROM pg_attribute WHERE attrelid=%s::regclass AND attname='embeddings'", (table,))
                row = cur.fetchone()
                if not row or row[0] != self.profile["stored_dimension"] or store not in profiles:
                    raise BackupError("BACKUP_EMBEDDING_PROFILE_MISMATCH")
        return profiles

    def snapshot(self, output):
        from src.lib.postgresql._naming import brain_db_name
        output = Path(output)
        (output / "system").mkdir(parents=True)
        (output / "brains").mkdir()
        registry = self.registry()
        import io
        self.inspect_secret_stream(io.BytesIO(canonical(registry)))
        (output / "system/registry.json").write_bytes(canonical(registry))
        fingerprints = {}
        for brain in registry["brains"]:
            name = brain["name_key"]
            database = brain_db_name(name)
            with self.connect(database) as conn, conn.cursor() as cur:
                cur.execute("SET TRANSACTION ISOLATION LEVEL REPEATABLE READ READ ONLY")
                cur.execute("SELECT pg_export_snapshot()")
                snapshot = cur.fetchone()[0]
                tables = self.tables(cur)
                self.validate_profile(cur, tables)
                digest = hashlib.sha256()
                for table in tables:
                    digest.update(table.encode())
                    # Canonical row order makes no-op checks independent of VACUUM,
                    # heap order and pg_dump's archive timestamp.
                    with conn.cursor(name=f"backup_{table}") as rows:
                        rows.execute(sql.SQL("SELECT to_jsonb(t)::text FROM {} t ORDER BY to_jsonb(t)::text COLLATE \"C\"").format(sql.Identifier(table)))
                        for (row,) in rows:
                            digest.update(row.encode() + b"\n")
                table_args = [item for table in tables for item in ("--table", f"public.{table}")]
                if tables:
                    archive = output / "brains" / f"{name}.dump"
                    with archive.open("wb") as stream:
                        self.tool("pg_dump", database, "--format=custom", "--no-owner", "--no-privileges",
                                  f"--snapshot={snapshot}", *table_args, stdout=stream)
                    self.tool("pg_restore", database, "--list", str(archive))
                    self.inspect_secrets(archive, database)
                    # Include schema/index changes in semantic identity, without
                    # archive timestamps or pg_dump's random restriction nonce.
                    schema = self.tool("pg_dump", database, "--schema-only", "--no-owner", "--no-privileges",
                                       f"--snapshot={snapshot}", *table_args)
                    stable = b"\n".join(line for line in schema.splitlines()
                                          if not line.startswith((b"--", b"\\restrict", b"\\unrestrict")))
                    digest.update(stable)
                fingerprints[name] = digest.hexdigest()
        if self.registry() != registry:
            raise BackupError("BACKUP_REGISTRY_CHANGED_DURING_SNAPSHOT")
        return registry, fingerprints

    def has_local_content(self):
        return bool(self.registry()["brains"])

    def restore(self, source, registry):
        from src.lib.postgresql._naming import brain_db_name
        if self.has_local_content():
            raise BackupError("BACKUP_RESTORE_REQUIRES_EMPTY_REGISTRY")
        # Preflight ALL destination collisions before writing any DB. Orphans
        # are never resurrected or overwritten, even if their names match.
        names = [brain_db_name(b["name_key"]) for b in registry["brains"]]
        if len(set(names)) != len(names):
            raise BackupError("BACKUP_DATABASE_NAME_COLLISION")
        created = []
        try:
            with self.connect("postgres", autocommit=True) as admin, admin.cursor() as cur:
                cur.execute("SELECT datname FROM pg_database WHERE datname=ANY(%s)", (names,))
                if cur.fetchall():
                    raise BackupError("BACKUP_RESTORE_DATABASE_COLLISION")
                for brain, database in zip(registry["brains"], names):
                    cur.execute(sql.SQL("CREATE DATABASE {} TEMPLATE template1").format(sql.Identifier(database)))
                    created.append(database)
                    # Selective dumps do not include pgvector extension DDL.
                    with self.connect(database, autocommit=True) as conn, conn.cursor() as target:
                        target.execute("CREATE EXTENSION IF NOT EXISTS vector")
                    archive = Path(source) / "brains" / f"{brain['name_key']}.dump"
                    if archive.exists():
                        self.tool("pg_restore", database, "--exit-on-error", "--single-transaction",
                                  "--no-owner", "--no-privileges", "--dbname", database, str(archive))
                    with self.connect(database) as conn, conn.cursor() as target:
                        self.validate_profile(target, self.tables(target))
            # Publish canonical registry LAST in one transaction. API is still
            # stopped; any restore failure leaves no partially registered brain.
            with self.connect() as conn, conn.cursor() as cur:
                cur.execute("LOCK TABLE data_brains,data_workspaces IN ACCESS EXCLUSIVE MODE")
                cur.execute("SELECT count(*) FROM data_brains")
                if cur.fetchone()[0]:
                    raise BackupError("BACKUP_RESTORE_REQUIRES_EMPTY_REGISTRY")
                cur.execute("DELETE FROM data_workspaces")
                for b in registry["brains"]:
                    cur.execute("INSERT INTO data_brains(id,name_key,pat,document,created_at) VALUES(%s,%s,'',%s,%s)",
                                (b["id"], b["name_key"], Json(b["document"]), b["created_at"]))
                for w in registry["workspaces"]:
                    cur.execute("INSERT INTO data_workspaces(id,brain_id,slug,document,created_at) VALUES(%s,%s,%s,%s,%s)",
                                (w["id"], w["brain_id"], w["slug"], Json(w["document"]), w["created_at"]))
                conn.commit()
        except BaseException:
            with self.connect("postgres", autocommit=True) as admin, admin.cursor() as cur:
                for database in reversed(created):
                    cur.execute(sql.SQL("DROP DATABASE {} WITH (FORCE)").format(sql.Identifier(database)))
            raise