diff --git a/hindsight-api/hindsight_api/admin/__init__.py b/hindsight-api/hindsight_api/admin/__init__.py new file mode 100644 index 00000000..db871af1 --- /dev/null +++ b/hindsight-api/hindsight_api/admin/__init__.py @@ -0,0 +1 @@ +# Admin CLI for Hindsight diff --git a/hindsight-api/hindsight_api/admin/cli.py b/hindsight-api/hindsight_api/admin/cli.py new file mode 100644 index 00000000..dfd74115 --- /dev/null +++ b/hindsight-api/hindsight_api/admin/cli.py @@ -0,0 +1,222 @@ +""" +Hindsight Admin CLI - backup and restore operations. +""" + +import asyncio +import io +import json +import logging +import zipfile +from datetime import datetime, timezone +from pathlib import Path +from typing import Any + +import asyncpg +import typer + +from ..config import HindsightConfig +from ..pg0 import parse_pg0_url, resolve_database_url + + +def _fq_table(table: str, schema: str) -> str: + """Get fully-qualified table name with schema prefix.""" + return f"{schema}.{table}" + + +# Setup logging +logging.basicConfig( + level=logging.INFO, + format="%(message)s", +) +logger = logging.getLogger(__name__) + +app = typer.Typer(name="hindsight-admin", help="Hindsight administrative commands") + +# Tables to backup/restore in dependency order +# Import must happen in this order due to foreign key constraints +BACKUP_TABLES = [ + "banks", + "documents", + "entities", + "chunks", + "memory_units", + "unit_entities", + "entity_cooccurrences", + "memory_links", +] + +MANIFEST_VERSION = "1" + + +async def _backup(database_url: str, output_path: Path, schema: str = "public") -> dict[str, Any]: + """Backup all tables to a zip file using binary COPY protocol.""" + conn = await asyncpg.connect(database_url) + try: + tables: dict[str, Any] = {} + manifest: dict[str, Any] = { + "version": MANIFEST_VERSION, + "created_at": datetime.now(timezone.utc).isoformat(), + "schema": schema, + "tables": tables, + } + + # Use a transaction with REPEATABLE READ isolation to get a consistent + # snapshot across all tables. This prevents race conditions where + # entity_cooccurrences could reference entities created after the + # entities table was backed up. + async with conn.transaction(isolation="repeatable_read"): + with zipfile.ZipFile(output_path, "w", zipfile.ZIP_DEFLATED) as zf: + for i, table in enumerate(BACKUP_TABLES, 1): + typer.echo(f" [{i}/{len(BACKUP_TABLES)}] Backing up {table}...", nl=False) + + buffer = io.BytesIO() + + # Use binary COPY for exact type preservation + # asyncpg requires schema_name as separate parameter + await conn.copy_from_table(table, schema_name=schema, output=buffer, format="binary") + + data = buffer.getvalue() + zf.writestr(f"{table}.bin", data) + + # Get row count for manifest + qualified_table = _fq_table(table, schema) + row_count = await conn.fetchval(f"SELECT COUNT(*) FROM {qualified_table}") + tables[table] = { + "rows": row_count, + "size_bytes": len(data), + } + + typer.echo(f" {row_count} rows") + + zf.writestr("manifest.json", json.dumps(manifest, indent=2)) + + return manifest + finally: + await conn.close() + + +async def _restore(database_url: str, input_path: Path, schema: str = "public") -> dict[str, Any]: + """Restore all tables from a zip file using binary COPY protocol.""" + conn = await asyncpg.connect(database_url) + try: + with zipfile.ZipFile(input_path, "r") as zf: + # Read and validate manifest + manifest: dict[str, Any] = json.loads(zf.read("manifest.json")) + if manifest.get("version") != MANIFEST_VERSION: + raise ValueError(f"Unsupported backup version: {manifest.get('version')}") + + # Use a transaction for atomic restore - either all tables are + # restored or none are, preventing partial/inconsistent state. + async with conn.transaction(): + typer.echo(" Clearing existing data...") + # Truncate tables in reverse order (respects FK constraints) + for table in reversed(BACKUP_TABLES): + qualified_table = _fq_table(table, schema) + await conn.execute(f"TRUNCATE TABLE {qualified_table} CASCADE") + + # Restore tables in forward order + for i, table in enumerate(BACKUP_TABLES, 1): + filename = f"{table}.bin" + if filename not in zf.namelist(): + typer.echo(f" [{i}/{len(BACKUP_TABLES)}] {table}: skipped (not in backup)") + continue + + expected_rows = manifest["tables"].get(table, {}).get("rows", "?") + typer.echo(f" [{i}/{len(BACKUP_TABLES)}] Restoring {table}... {expected_rows} rows") + + data = zf.read(filename) + buffer = io.BytesIO(data) + # asyncpg requires schema_name as separate parameter + await conn.copy_to_table(table, schema_name=schema, source=buffer, format="binary") + + # Refresh materialized view + typer.echo(" Refreshing materialized views...") + await conn.execute(f"REFRESH MATERIALIZED VIEW {_fq_table('memory_units_bm25', schema)}") + + return manifest + finally: + await conn.close() + + +async def _run_backup(db_url: str, output: Path, schema: str = "public") -> dict[str, Any]: + """Resolve database URL and run backup.""" + is_pg0, instance_name, _ = parse_pg0_url(db_url) + if is_pg0: + typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...") + resolved_url = await resolve_database_url(db_url) + return await _backup(resolved_url, output, schema) + + +async def _run_restore(db_url: str, input_file: Path, schema: str = "public") -> dict[str, Any]: + """Resolve database URL and run restore.""" + is_pg0, instance_name, _ = parse_pg0_url(db_url) + if is_pg0: + typer.echo(f"Starting embedded PostgreSQL (instance: {instance_name})...") + resolved_url = await resolve_database_url(db_url) + return await _restore(resolved_url, input_file, schema) + + +@app.command() +def backup( + output: Path = typer.Argument(..., help="Output file path (.zip)"), + schema: str = typer.Option("public", "--schema", "-s", help="Database schema to backup"), +): + """Backup the Hindsight database to a zip file.""" + config = HindsightConfig.from_env() + + if not config.database_url: + typer.echo("Error: Database URL not configured.", err=True) + typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True) + raise typer.Exit(1) + + if output.suffix != ".zip": + output = output.with_suffix(".zip") + + typer.echo(f"Backing up database (schema: {schema}) to {output}...") + + manifest = asyncio.run(_run_backup(config.database_url, output, schema)) + + total_rows = sum(t["rows"] for t in manifest["tables"].values()) + typer.echo(f"Backed up {total_rows} rows across {len(BACKUP_TABLES)} tables") + typer.echo(f"Backup saved to {output}") + + +@app.command() +def restore( + input_file: Path = typer.Argument(..., help="Input backup file (.zip)"), + schema: str = typer.Option("public", "--schema", "-s", help="Database schema to restore to"), + yes: bool = typer.Option(False, "--yes", "-y", help="Skip confirmation prompt"), +): + """Restore the database from a backup file. WARNING: This deletes all existing data.""" + config = HindsightConfig.from_env() + + if not config.database_url: + typer.echo("Error: Database URL not configured.", err=True) + typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True) + raise typer.Exit(1) + + if not input_file.exists(): + typer.echo(f"Error: File not found: {input_file}", err=True) + raise typer.Exit(1) + + if not yes: + typer.confirm( + "This will DELETE all existing data and replace it with the backup. Continue?", + abort=True, + ) + + typer.echo(f"Restoring database (schema: {schema}) from {input_file}...") + + manifest = asyncio.run(_run_restore(config.database_url, input_file, schema)) + + total_rows = sum(t["rows"] for t in manifest["tables"].values()) + typer.echo(f"Restored {total_rows} rows across {len(BACKUP_TABLES)} tables") + typer.echo("Restore complete") + + +def main(): + app() + + +if __name__ == "__main__": + main() diff --git a/hindsight-api/hindsight_api/engine/memory_engine.py b/hindsight-api/hindsight_api/engine/memory_engine.py index 7c8e536a..a94f5dcd 100644 --- a/hindsight-api/hindsight_api/engine/memory_engine.py +++ b/hindsight-api/hindsight_api/engine/memory_engine.py @@ -132,7 +132,7 @@ if TYPE_CHECKING: from enum import Enum -from ..pg0 import EmbeddedPostgres +from ..pg0 import EmbeddedPostgres, parse_pg0_url from .entity_resolver import EntityResolver from .llm_wrapper import LLMConfig from .query_analyzer import QueryAnalyzer @@ -259,31 +259,14 @@ class MemoryEngine(MemoryEngineInterface): memory_llm_base_url = memory_llm_base_url or config.get_llm_base_url() or None # Track pg0 instance (if used) self._pg0: EmbeddedPostgres | None = None - self._pg0_instance_name: str | None = None # Initialize PostgreSQL connection URL # The actual URL will be set during initialize() after starting the server # Supports: "pg0" (default instance), "pg0://instance-name" (named instance), or regular postgresql:// URL - if db_url == "pg0": - self._use_pg0 = True - self._pg0_instance_name = "hindsight" - self._pg0_port = None # Use default port - self.db_url = None - elif db_url.startswith("pg0://"): - self._use_pg0 = True - # Parse instance name and optional port: pg0://instance-name or pg0://instance-name:port - url_part = db_url[6:] # Remove "pg0://" - if ":" in url_part: - self._pg0_instance_name, port_str = url_part.rsplit(":", 1) - self._pg0_port = int(port_str) - else: - self._pg0_instance_name = url_part or "hindsight" - self._pg0_port = None # Use default port + self._use_pg0, self._pg0_instance_name, self._pg0_port = parse_pg0_url(db_url) + if self._use_pg0: self.db_url = None else: - self._use_pg0 = False - self._pg0_instance_name = None - self._pg0_port = None self.db_url = db_url # Set default base URL if not provided diff --git a/hindsight-api/hindsight_api/pg0.py b/hindsight-api/hindsight_api/pg0.py index 797d9465..ec9ef443 100644 --- a/hindsight-api/hindsight_api/pg0.py +++ b/hindsight-api/hindsight_api/pg0.py @@ -132,3 +132,56 @@ async def stop_embedded_postgres() -> None: global _default_instance if _default_instance: await _default_instance.stop() + + +def parse_pg0_url(db_url: str) -> tuple[bool, str | None, int | None]: + """ + Parse a database URL and check if it's a pg0:// embedded database URL. + + Supports: + - "pg0" -> default instance "hindsight" + - "pg0://instance-name" -> named instance + - "pg0://instance-name:port" -> named instance with explicit port + - Any other URL (e.g., postgresql://) -> not a pg0 URL + + Args: + db_url: The database URL to parse + + Returns: + Tuple of (is_pg0, instance_name, port) + - is_pg0: True if this is a pg0 URL + - instance_name: The instance name (or None if not pg0) + - port: The explicit port (or None for auto-assign) + """ + if db_url == "pg0": + return True, "hindsight", None + + if db_url.startswith("pg0://"): + url_part = db_url[6:] # Remove "pg0://" + if ":" in url_part: + instance_name, port_str = url_part.rsplit(":", 1) + return True, instance_name or "hindsight", int(port_str) + else: + return True, url_part or "hindsight", None + + return False, None, None + + +async def resolve_database_url(db_url: str) -> str: + """ + Resolve a database URL, handling pg0:// embedded database URLs. + + If the URL is a pg0:// URL, starts the embedded PostgreSQL and returns + the actual postgresql:// connection URL. Otherwise, returns the URL unchanged. + + Args: + db_url: Database URL (pg0://, pg0, or postgresql://) + + Returns: + The resolved postgresql:// connection URL + """ + is_pg0, instance_name, port = parse_pg0_url(db_url) + if is_pg0: + pg0 = EmbeddedPostgres(name=instance_name, port=port) + return await pg0.ensure_running() + return db_url diff --git a/hindsight-api/pyproject.toml b/hindsight-api/pyproject.toml index caca2640..8d3c5d3e 100644 --- a/hindsight-api/pyproject.toml +++ b/hindsight-api/pyproject.toml @@ -38,6 +38,7 @@ dependencies = [ "dateparser>=1.2.2", "google-genai>=1.0.0", "anthropic>=0.40.0", + "typer>=0.9.0", ] [project.optional-dependencies] @@ -52,6 +53,7 @@ test = [ [project.scripts] hindsight-api = "hindsight_api.main:main" hindsight-local-mcp = "hindsight_api.mcp_local:main" +hindsight-admin = "hindsight_api.admin.cli:main" [tool.hatch.build.targets.wheel] packages = ["hindsight_api"] @@ -75,7 +77,7 @@ log_cli = true log_cli_level = "INFO" log_cli_format = "%(asctime)s - %(levelname)s - %(name)s - %(message)s" log_cli_date_format = "%Y-%m-%d %H:%M:%S" -addopts = "--timeout 120 -n 8 --durations=10 -v" +addopts = "--timeout 120 -n 8 --dist loadgroup --durations=10 -v" asyncio_mode = "auto" asyncio_default_fixture_loop_scope = "function" log_auto_indent = true diff --git a/hindsight-api/tests/test_admin_backup_restore.py b/hindsight-api/tests/test_admin_backup_restore.py new file mode 100644 index 00000000..37b0a4ce --- /dev/null +++ b/hindsight-api/tests/test_admin_backup_restore.py @@ -0,0 +1,208 @@ +""" +Tests for admin backup and restore functionality. + +Note: These tests run sequentially (not in parallel) because they all +manipulate the same database and do full backup/restore operations. +""" + +import tempfile +import uuid +import zipfile +from pathlib import Path + +import pytest + +from hindsight_api import RequestContext +from hindsight_api.admin.cli import _backup, _restore, BACKUP_TABLES + + +# Run these tests sequentially since they do full DB backup/restore +pytestmark = pytest.mark.xdist_group(name="backup_restore") + + +@pytest.mark.asyncio +async def test_backup_restore_roundtrip(memory, pg0_db_url, request_context): + """Test that backup and restore preserves all data correctly.""" + # Use unique bank ID to avoid conflicts + bank_id = f"test-backup-{uuid.uuid4().hex[:8]}" + + # Create some test data + await memory.retain_batch_async( + bank_id=bank_id, + contents=[ + {"content": "Alice is a software engineer who loves Python."}, + {"content": "Bob works with Alice on the backend team."}, + {"content": "The team uses PostgreSQL for their database."}, + ], + request_context=request_context, + ) + + # Get counts before backup + async with memory._pool.acquire() as conn: + counts_before = {} + for table in BACKUP_TABLES: + counts_before[table] = await conn.fetchval(f"SELECT COUNT(*) FROM {table}") + + # Verify we have data + assert counts_before["banks"] > 0 + assert counts_before["memory_units"] > 0 + + # Backup to a temp file + with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as f: + backup_path = Path(f.name) + + try: + manifest = await _backup(pg0_db_url, backup_path) + + # Verify backup file exists and is valid + assert backup_path.exists() + assert backup_path.stat().st_size > 0 + + # Verify manifest + assert manifest["version"] == "1" + assert "created_at" in manifest + for table in BACKUP_TABLES: + assert table in manifest["tables"] + assert manifest["tables"][table]["rows"] == counts_before[table] + + # Verify zip contents + with zipfile.ZipFile(backup_path, "r") as zf: + assert "manifest.json" in zf.namelist() + for table in BACKUP_TABLES: + assert f"{table}.bin" in zf.namelist() + + # Clear all data + async with memory._pool.acquire() as conn: + for table in reversed(BACKUP_TABLES): + await conn.execute(f"TRUNCATE TABLE {table} CASCADE") + + # Verify data is gone + async with memory._pool.acquire() as conn: + for table in BACKUP_TABLES: + count = await conn.fetchval(f"SELECT COUNT(*) FROM {table}") + assert count == 0, f"Table {table} should be empty after truncate" + + # Restore from backup + await _restore(pg0_db_url, backup_path) + + # Verify counts match original + async with memory._pool.acquire() as conn: + for table in BACKUP_TABLES: + count = await conn.fetchval(f"SELECT COUNT(*) FROM {table}") + assert count == counts_before[table], f"Table {table} count mismatch after restore" + + # Verify data content is preserved + async with memory._pool.acquire() as conn: + texts = await conn.fetch( + "SELECT text FROM memory_units WHERE bank_id = $1", + bank_id, + ) + text_content = " ".join(r["text"] for r in texts) + assert "Alice" in text_content or "software" in text_content + + finally: + # Cleanup + if backup_path.exists(): + backup_path.unlink() + + +@pytest.mark.asyncio +async def test_backup_restore_preserves_all_column_types(memory, pg0_db_url, request_context): + """Test that all column types are preserved: vectors, UUIDs, timestamps, JSONB.""" + # Use unique bank ID + bank_id = f"test-types-{uuid.uuid4().hex[:8]}" + + # Create data with meaningful content that will produce facts + await memory.retain_batch_async( + bank_id=bank_id, + contents=[ + {"content": "John Smith is a senior engineer at Acme Corp since 2020."}, + {"content": "The project deadline is December 15th 2024."}, + ], + request_context=request_context, + ) + + # Get original data with all important column types + async with memory._pool.acquire() as conn: + # memory_units: UUID (id), Vector (embedding), Timestamp (event_date, created_at), JSONB (metadata) + original_unit = await conn.fetchrow( + """SELECT id, embedding, event_date, created_at, metadata, text + FROM memory_units WHERE bank_id = $1 LIMIT 1""", + bank_id, + ) + + # entities: UUID (id), Timestamp (first_seen, last_seen), JSONB (metadata) + original_entity = await conn.fetchrow( + """SELECT id, first_seen, last_seen, metadata, canonical_name + FROM entities WHERE bank_id = $1 LIMIT 1""", + bank_id, + ) + + # banks: JSONB (personality/disposition) + original_bank = await conn.fetchrow( + "SELECT bank_id, created_at, updated_at FROM banks WHERE bank_id = $1", + bank_id, + ) + + assert original_unit is not None, "Should have created memory units" + assert original_unit["embedding"] is not None, "Should have embedding" + assert original_unit["id"] is not None, "Should have UUID" + assert original_entity is not None, "Should have created entities" + + with tempfile.NamedTemporaryFile(suffix=".zip", delete=False) as f: + backup_path = Path(f.name) + + try: + await _backup(pg0_db_url, backup_path) + + # Clear all data + async with memory._pool.acquire() as conn: + for table in reversed(BACKUP_TABLES): + await conn.execute(f"TRUNCATE TABLE {table} CASCADE") + + await _restore(pg0_db_url, backup_path) + + # Verify all column types are preserved exactly + async with memory._pool.acquire() as conn: + restored_unit = await conn.fetchrow( + """SELECT id, embedding, event_date, created_at, metadata, text + FROM memory_units WHERE bank_id = $1 LIMIT 1""", + bank_id, + ) + + restored_entity = await conn.fetchrow( + """SELECT id, first_seen, last_seen, metadata, canonical_name + FROM entities WHERE bank_id = $1 LIMIT 1""", + bank_id, + ) + + restored_bank = await conn.fetchrow( + "SELECT bank_id, created_at, updated_at FROM banks WHERE bank_id = $1", + bank_id, + ) + + # Verify memory_units + assert restored_unit is not None, "Should have restored memory unit" + assert restored_unit["id"] == original_unit["id"], "UUID should match exactly" + assert restored_unit["text"] == original_unit["text"], "Text should match" + assert list(restored_unit["embedding"]) == list(original_unit["embedding"]), "Vector embedding should match exactly" + assert restored_unit["event_date"] == original_unit["event_date"], "Timestamp should match exactly" + assert restored_unit["created_at"] == original_unit["created_at"], "Created timestamp should match" + assert restored_unit["metadata"] == original_unit["metadata"], "JSONB metadata should match" + + # Verify entities + assert restored_entity is not None, "Should have restored entity" + assert restored_entity["id"] == original_entity["id"], "Entity UUID should match" + assert restored_entity["canonical_name"] == original_entity["canonical_name"], "Entity name should match" + assert restored_entity["first_seen"] == original_entity["first_seen"], "Entity first_seen should match" + assert restored_entity["last_seen"] == original_entity["last_seen"], "Entity last_seen should match" + assert restored_entity["metadata"] == original_entity["metadata"], "Entity metadata should match" + + # Verify banks + assert restored_bank is not None, "Should have restored bank" + assert restored_bank["bank_id"] == original_bank["bank_id"], "Bank ID should match" + assert restored_bank["created_at"] == original_bank["created_at"], "Bank created_at should match" + + finally: + if backup_path.exists(): + backup_path.unlink() diff --git a/uv.lock b/uv.lock index 4fbc858f..b64b4072 100644 --- a/uv.lock +++ b/uv.lock @@ -1215,6 +1215,7 @@ dependencies = [ { name = "tiktoken" }, { name = "torch" }, { name = "transformers" }, + { name = "typer" }, { name = "uvicorn" }, { name = "wsproto" }, ] @@ -1274,6 +1275,7 @@ requires-dist = [ { name = "tiktoken", specifier = ">=0.12.0" }, { name = "torch", specifier = ">=2.0.0" }, { name = "transformers", specifier = ">=4.30.0,<4.46.0" }, + { name = "typer", specifier = ">=0.9.0" }, { name = "uvicorn", specifier = ">=0.38.0" }, { name = "wsproto", specifier = ">=1.0.0" }, ]