Spaces:
Running
Running
| """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"] | |
| 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)) | |
| 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 | |