Spaces:
Running
Running
File size: 10,071 Bytes
2f1e7ff | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 | """Real logical restore round-trip against an explicitly disposable PG cluster.
Set BACKUP_QUALIFICATION_DISPOSABLE=true and the normal HF CPU environment.
Never run against the production cluster. Only randomly named test DBs are
created/dropped. Git remote is a local bare repository; no deployment occurs.
"""
import asyncio
import json
import os
from pathlib import Path
import subprocess
import sys
import tempfile
from uuid import uuid4
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
def qualify():
if os.getenv("BACKUP_QUALIFICATION_DISPOSABLE") != "true":
raise RuntimeError("Disposable PostgreSQL cluster acknowledgement required")
# Set before application/config imports; never touch a real registry.
suffix = uuid4().hex[:12]
system = "backup_qualification_" + suffix
os.environ["POSTGRES_SYSTEM_DATABASE"] = system
os.environ["BRAIN_BACKUP_ENABLED"] = "true"
from src.config import config
from src.constants.data import TextChunk, Workspace
from src.lib.postgresql.data import PostgreSQLDataClient
from src.lib.postgresql._naming import brain_db_name
from src.lib.postgresql import _provisioning as pools
from src.services.brain_backup.git import GitRepository, BackupError
from src.services.brain_backup.postgres import PostgreSQL
from src.services.brain_backup.manager import Manager, validate
from src.services.brain_backup.state import generation
from src.lib.postgresql.vectors import PostgreSQLVectorStoreClient
from psycopg2 import sql
from psycopg2.extras import Json
pg = PostgreSQL.configured()
names = ["backup-projects-" + suffix, "backup-habits-" + suffix]
orphan = "backup-orphan-" + suffix
facts = ["Asterquartz project launches on October 31 with purple balloons.",
"Asterquartz habit is twelve minutes of sunrise breathing each morning."]
results = {}
with tempfile.TemporaryDirectory(prefix="brainapi-backup-roundtrip-") as temp:
temp = Path(temp)
os.environ["BRAIN_BACKUP_STATE_DIR"] = str(temp / "state")
remote = temp / "remote.git"
subprocess.run(["git", "init", "--bare", "--quiet", str(remote)], check=True)
manager = Manager(GitRepository(temp / "checkout", str(remote)), pg, temp / "local")
client = PostgreSQLDataClient()
# Prove every cleanup target was absent before testing. Never delete a
# preexisting database, even in a cluster acknowledged as disposable.
with pg.connect("postgres") as conn, conn.cursor() as cur:
targets = [system, *(brain_db_name(n) for n in names + [orphan])]
cur.execute("SELECT datname FROM pg_database WHERE datname=ANY(%s)", (targets,))
if cur.fetchall():
raise RuntimeError("Qualification namespace collision; no cleanup attempted")
try:
with pg.connect("template1", autocommit=True) as conn, conn.cursor() as cur:
cur.execute("CREATE EXTENSION IF NOT EXISTS vector")
client._ensure_system_schema()
manager.startup()
identities = {}
for name, fact in zip(names, facts):
brain = client.create_brain(name)
identities[name] = brain.id
client.create_workspace(Workspace(brain_id=name, slug=name, display_name=name))
client.save_text_chunk(TextChunk(text=fact), brain_id=name)
vector = PostgreSQLVectorStoreClient()
for store in ("nodes", "triplets", "observations", "data", "relationships"):
vector._ensure_store(store, name)
# Exercise legacy PAT stripping and separate access-token registry.
with pg.connect() as conn, conn.cursor() as cur:
cur.execute("UPDATE data_brains SET pat='legacy-secret-not-for-backup', document=document || %s::jsonb",
(Json({"pat": "legacy-secret-not-for-backup"}),))
cur.execute("CREATE TABLE access_tokens(token_hash text)")
cur.execute("INSERT INTO access_tokens VALUES('separate-access-token-hash')")
conn.commit()
baseline = generation()
with pools.borrow(pools.get_system_pool()) as conn, conn.cursor() as cur:
cur.execute("INSERT INTO access_tokens VALUES('auth-only-not-content')")
conn.commit()
assert generation() == baseline
with pools.borrow(pools.get_brain_pool(names[0])) as conn, conn.cursor() as cur:
cur.execute("UPDATE data_text_chunks SET text='rolled back'")
conn.rollback()
assert generation() == baseline
with pools.borrow(pools.get_brain_pool(names[0])) as conn:
with conn, conn.cursor() as cur:
cur.execute("UPDATE data_text_chunks SET metadata=metadata")
assert generation() > baseline
results["dirty_tracking"] = "auth-only writes and rollback excluded; context-managed content COMMIT tracked"
with pg.connect("postgres", autocommit=True) as conn, conn.cursor() as cur:
cur.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(brain_db_name(orphan))))
assert generation() > 0
first, changed = manager.backup()
assert changed
manifest, registry = validate(manager.repo.path, pg.profile, pg.major_version())
assert set(manifest["brains"]) == set(names) and orphan not in manifest["brains"]
assert identities == {b["name_key"]: b["id"] for b in registry["brains"]}
assert "legacy-secret" not in json.dumps(registry)
assert not any("access_tokens" in p for p in manifest["checksums"])
second, changed = manager.backup()
assert second == first and not changed
results["logical_snapshot"] = "registered brains only; five 512d profiles; no redundant commit"
results["credential_exclusion"] = "legacy PAT removed; access_tokens table/hash excluded"
# Force profile rejection on a real initialized vector DB.
with pg.connect(brain_db_name(names[0])) as conn, conn.cursor() as cur:
cur.execute("UPDATE embedding_profiles SET profile=profile || '{\"native_dimension\":1024}'::jsonb")
conn.commit()
try:
manager.backup()
except BackupError as exc:
assert "PROFILE_MISMATCH" in str(exc)
else:
raise AssertionError("Mismatched vector metadata was accepted")
results["profile_mismatch"] = "rejected"
# Only names created above, on the acknowledged disposable cluster.
for database in list(pools._registry._pools):
pools._registry.discard(database)
with pg.connect("postgres", autocommit=True) as conn, conn.cursor() as cur:
for name in names:
cur.execute(sql.SQL("DROP DATABASE {} WITH (FORCE)").format(sql.Identifier(brain_db_name(name))))
with pg.connect() as conn, conn.cursor() as cur:
cur.execute("TRUNCATE data_workspaces,data_brains")
conn.commit()
manager.startup()
restored = PostgreSQLDataClient()
for name, fact in zip(names, facts):
assert restored.get_brain(name).id == identities[name]
assert fact in [chunk.text for chunk, _ in restored.search_bm25("Asterquartz", name)]
assert restored.get_workspace_by_brain_id(name).brain_id == name
results["roundtrip"] = "destroyed test brain DBs, restored same IDs/workspaces/facts and BM25 results"
restored.save_text_chunk(TextChunk(text="Asterquartz newer local content retained on restart."), names[0])
manager.startup()
assert any("newer local content" in c.text for c, _ in restored.search_bm25("Asterquartz", names[0]))
results["restart"] = "nonempty local data retained with matching checkpoint lineage"
if os.getenv("BACKUP_QUALIFICATION_RETRIEVAL") == "true":
from src.services.api.controllers.search import search
from src.services.api.controllers.retrieve import get_context
from src.services.api.constants.requests import SearchRequestBody, GetContextRequestBody
for name, fact in zip(names, facts):
hits = asyncio.run(search(SearchRequestBody(query="Asterquartz", brain_id=name)))
context = asyncio.run(get_context(GetContextRequestBody(text="Asterquartz", brain_id=name)))
assert fact in json.dumps(hits.model_dump(mode="json"))
assert fact in json.dumps(context.model_dump(mode="json"))
results["search_context"] = "actual search/context controllers returned both restored unique facts"
# New content changes the remote commit after restoration.
restored.save_text_chunk(TextChunk(text="Asterquartz new milestone after restoration."), names[0])
third, changed = manager.backup()
assert changed and third != first
results["new_mutation"] = "new logical checkpoint commit produced"
export = os.getenv("BACKUP_QUALIFICATION_EXPORT")
if export:
import shutil
shutil.copytree(remote, Path(export))
finally:
for database in list(pools._registry._pools):
pools._registry.discard(database)
with pg.connect("postgres", autocommit=True) as conn, conn.cursor() as cur:
for database in [*(brain_db_name(n) for n in names + [orphan]), system]:
cur.execute(sql.SQL("DROP DATABASE IF EXISTS {} WITH (FORCE)").format(sql.Identifier(database)))
print(json.dumps({"result": "passed", **results}, sort_keys=True))
if __name__ == "__main__":
qualify()
|