From 9a694f64b895f24d428449d83ef1dde4886cebd1 Mon Sep 17 00:00:00 2001 From: Chris Bartholomew Date: Tue, 10 Mar 2026 05:10:03 -0400 Subject: [PATCH] Fix run-db-migration for all-tenant upgrades (#530) * Add release-scoped migration admin command * Fix run-db-migration for all-tenant upgrades --- CLAUDE.md | 2 +- docker/standalone/start-all.sh | 18 +- hindsight-api/hindsight_api/admin/cli.py | 90 +++++++- .../tests/test_admin_backup_restore.py | 195 ++++++++++++++++++ hindsight-docs/docs/developer/admin-cli.md | 7 +- .../docs/developer/configuration.md | 4 +- .../references/developer/admin-cli.md | 7 +- .../references/developer/configuration.md | 4 +- 8 files changed, 305 insertions(+), 22 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 940f5417..8fea5af0 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -154,7 +154,7 @@ Key tables: `banks`, `memory_units`, `documents`, `entities`, `entity_links` 3. **Run migrations locally**: ```bash - # Set database URL and run migrations + # Set database URL and run migrations for the base schema plus all tenants uv run hindsight-admin run-db-migration # Run on a specific tenant schema diff --git a/docker/standalone/start-all.sh b/docker/standalone/start-all.sh index e1bf9443..561ad386 100755 --- a/docker/standalone/start-all.sh +++ b/docker/standalone/start-all.sh @@ -77,18 +77,32 @@ PIDS=() # Start API if enabled if [ "$ENABLE_API" = "true" ]; then cd /app/api + API_HEALTH_URL="${HINDSIGHT_API_HEALTH_URL:-http://localhost:8888/health}" + API_STARTUP_WAIT_SECONDS="${HINDSIGHT_API_STARTUP_WAIT_SECONDS:-300}" + # Run API directly - Python's PYTHONUNBUFFERED=1 handles output buffering hindsight-api & API_PID=$! PIDS+=($API_PID) # Wait for API to be ready - for i in {1..60}; do - if curl -sf http://localhost:8888/health &>/dev/null; then + api_ready=false + for ((i=1; i<=API_STARTUP_WAIT_SECONDS; i++)); do + if ! kill -0 "$API_PID" 2>/dev/null; then + wait "$API_PID" + exit $? + fi + if curl -sf "$API_HEALTH_URL" &>/dev/null; then + api_ready=true break fi sleep 1 done + + if [ "$api_ready" != "true" ]; then + echo "❌ API did not become healthy within ${API_STARTUP_WAIT_SECONDS}s" + exit 1 + fi else echo "API disabled (HINDSIGHT_ENABLE_API=false)" fi diff --git a/hindsight-api/hindsight_api/admin/cli.py b/hindsight-api/hindsight_api/admin/cli.py index d5009602..8b80dca1 100644 --- a/hindsight-api/hindsight_api/admin/cli.py +++ b/hindsight-api/hindsight_api/admin/cli.py @@ -14,7 +14,8 @@ from typing import Any import asyncpg import typer -from ..config import HindsightConfig +from ..config import DEFAULT_DATABASE_SCHEMA, HindsightConfig +from ..extensions import TenantExtension, load_extension from ..pg0 import parse_pg0_url, resolve_database_url @@ -214,20 +215,81 @@ def restore( typer.echo("Restore complete") -async def _run_migration(db_url: str, schema: str = "public") -> None: - """Resolve database URL and run migrations.""" - from ..migrations import run_migrations +async def _run_migration( + db_url: str, + schema: str | None = None, + base_schema: str = DEFAULT_DATABASE_SCHEMA, + embedding_dimension: int | None = None, +) -> list[str]: + """Resolve database URL and run migrations for one schema or all discovered schemas.""" + from ..migrations import ( + ensure_embedding_dimension, + ensure_text_search_extension, + ensure_vector_extension, + run_migrations, + ) 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) - run_migrations(resolved_url, schema=schema) + + config = HindsightConfig.from_env() + if schema: + schemas = [schema] + else: + tenant_extension = load_extension("TENANT", TenantExtension) + + schemas = [base_schema or DEFAULT_DATABASE_SCHEMA] + if tenant_extension: + tenants = await tenant_extension.list_tenants() + schemas.extend(tenant.schema for tenant in tenants if tenant.schema) + + # Preserve order while removing duplicates. + schemas = list(dict.fromkeys(schemas)) + + for schema in schemas: + run_migrations(resolved_url, schema=schema) + + if embedding_dimension is not None: + for schema in schemas: + ensure_embedding_dimension( + resolved_url, + embedding_dimension, + schema=schema, + vector_extension=config.vector_extension, + ) + + for schema in schemas: + ensure_vector_extension( + resolved_url, + vector_extension=config.vector_extension, + schema=schema, + ) + + for schema in schemas: + ensure_text_search_extension( + resolved_url, + text_search_extension=config.text_search_extension, + schema=schema, + ) + + return schemas @app.command(name="run-db-migration") def run_db_migration( - schema: str = typer.Option("public", "--schema", "-s", help="Database schema to run migrations on"), + schema: str | None = typer.Option( + None, + "--schema", + "-s", + help="Database schema to run migrations on. If omitted, migrate the base schema and all discovered tenant schemas.", + ), + embedding_dimension: int | None = typer.Option( + None, + "--embedding-dimension", + help="Expected embedding dimension to enforce after migrations. Omit to skip dimension sync.", + ), ): """Run database migrations to the latest version.""" config = HindsightConfig.from_env() @@ -237,11 +299,21 @@ def run_db_migration( typer.echo("Set HINDSIGHT_API_DATABASE_URL environment variable.", err=True) raise typer.Exit(1) - typer.echo(f"Running database migrations (schema: {schema})...") + if schema: + typer.echo(f"Running database migrations for schema: {schema}...") + else: + typer.echo("Running database migrations for base schema and all discovered tenant schemas...") - asyncio.run(_run_migration(config.database_url, schema)) + schemas = asyncio.run( + _run_migration( + config.database_url, + schema=schema, + base_schema=config.database_schema, + embedding_dimension=embedding_dimension, + ) + ) - typer.echo("Database migrations completed successfully") + typer.echo(f"Database migrations completed successfully for {len(schemas)} schema(s)") async def _decommission_worker(db_url: str, worker_id: str, schema: str = "public") -> int: diff --git a/hindsight-api/tests/test_admin_backup_restore.py b/hindsight-api/tests/test_admin_backup_restore.py index 32e31e0b..2e9cd107 100644 --- a/hindsight-api/tests/test_admin_backup_restore.py +++ b/hindsight-api/tests/test_admin_backup_restore.py @@ -15,7 +15,9 @@ import asyncpg import pytest import pytest_asyncio +import hindsight_api.admin.cli as admin_cli from hindsight_api.admin.cli import _backup, _restore, BACKUP_TABLES +from hindsight_api.extensions import Tenant from hindsight_api.migrations import run_migrations @@ -290,3 +292,196 @@ async def test_backup_restore_preserves_all_column_types(backup_test_schema): finally: if backup_path.exists(): backup_path.unlink() + + +@pytest.mark.asyncio +async def test_run_migration_without_schema_discovers_and_deduplicates_schemas(monkeypatch): + """run-db-migration without --schema should include the base schema and deduplicate tenant schemas.""" + calls: dict[str, list] = { + "run_migrations": [], + "ensure_vector_extension": [], + "ensure_text_search_extension": [], + } + + class MockTenantExtension: + async def list_tenants(self): + return [ + Tenant(schema="public"), + Tenant(schema="tenant_demo"), + Tenant(schema="tenant_demo"), + ] + + async def fake_resolve_database_url(db_url: str) -> str: + return f"resolved::{db_url}" + + def fake_run_migrations(database_url: str, schema: str | None = None) -> None: + calls["run_migrations"].append((database_url, schema)) + + def fake_ensure_vector_extension( + database_url: str, + vector_extension: str = "pgvector", + schema: str | None = None, + ) -> None: + calls["ensure_vector_extension"].append((database_url, vector_extension, schema)) + + def fake_ensure_text_search_extension( + database_url: str, + text_search_extension: str = "native", + schema: str | None = None, + ) -> None: + calls["ensure_text_search_extension"].append((database_url, text_search_extension, schema)) + + monkeypatch.setenv("HINDSIGHT_API_DATABASE_URL", "postgresql://test") + monkeypatch.setattr(admin_cli, "load_extension", lambda *args, **kwargs: MockTenantExtension()) + monkeypatch.setattr(admin_cli, "resolve_database_url", fake_resolve_database_url) + + from hindsight_api import migrations as migrations_module + + monkeypatch.setattr(migrations_module, "run_migrations", fake_run_migrations) + monkeypatch.setattr(migrations_module, "ensure_vector_extension", fake_ensure_vector_extension) + monkeypatch.setattr(migrations_module, "ensure_text_search_extension", fake_ensure_text_search_extension) + + schemas = await admin_cli._run_migration("postgresql://test") + + assert schemas == ["public", "tenant_demo"] + assert calls["run_migrations"] == [ + ("resolved::postgresql://test", "public"), + ("resolved::postgresql://test", "tenant_demo"), + ] + assert calls["ensure_vector_extension"] == [ + ("resolved::postgresql://test", "pgvector", "public"), + ("resolved::postgresql://test", "pgvector", "tenant_demo"), + ] + assert calls["ensure_text_search_extension"] == [ + ("resolved::postgresql://test", "native", "public"), + ("resolved::postgresql://test", "native", "tenant_demo"), + ] + + +@pytest.mark.asyncio +async def test_run_migration_without_schema_runs_optional_post_migration_hooks(monkeypatch): + """Embedding dimension sync should be optional, while vector/text checks always run.""" + monkeypatch.setenv("HINDSIGHT_API_DATABASE_URL", "postgresql://test") + calls: dict[str, list] = { + "run_migrations": [], + "ensure_embedding_dimension": [], + "ensure_vector_extension": [], + "ensure_text_search_extension": [], + } + + class MockTenantExtension: + async def list_tenants(self): + return [Tenant(schema="tenant_demo")] + + async def fake_resolve_database_url(db_url: str) -> str: + return f"resolved::{db_url}" + + def fake_run_migrations(database_url: str, schema: str | None = None) -> None: + calls["run_migrations"].append((database_url, schema)) + + def fake_ensure_embedding_dimension( + database_url: str, + dimension: int, + schema: str | None = None, + vector_extension: str = "pgvector", + ) -> None: + calls["ensure_embedding_dimension"].append((database_url, dimension, schema, vector_extension)) + + def fake_ensure_vector_extension( + database_url: str, + vector_extension: str = "pgvector", + schema: str | None = None, + ) -> None: + calls["ensure_vector_extension"].append((database_url, vector_extension, schema)) + + def fake_ensure_text_search_extension( + database_url: str, + text_search_extension: str = "native", + schema: str | None = None, + ) -> None: + calls["ensure_text_search_extension"].append((database_url, text_search_extension, schema)) + + monkeypatch.setattr(admin_cli, "load_extension", lambda *args, **kwargs: MockTenantExtension()) + monkeypatch.setattr(admin_cli, "resolve_database_url", fake_resolve_database_url) + + from hindsight_api import migrations as migrations_module + + monkeypatch.setattr(migrations_module, "run_migrations", fake_run_migrations) + monkeypatch.setattr(migrations_module, "ensure_embedding_dimension", fake_ensure_embedding_dimension) + monkeypatch.setattr(migrations_module, "ensure_vector_extension", fake_ensure_vector_extension) + monkeypatch.setattr(migrations_module, "ensure_text_search_extension", fake_ensure_text_search_extension) + + schemas = await admin_cli._run_migration( + "postgresql://test", + base_schema="public", + embedding_dimension=384, + ) + + assert schemas == ["public", "tenant_demo"] + assert calls["run_migrations"] == [ + ("resolved::postgresql://test", "public"), + ("resolved::postgresql://test", "tenant_demo"), + ] + assert calls["ensure_embedding_dimension"] == [ + ("resolved::postgresql://test", 384, "public", "pgvector"), + ("resolved::postgresql://test", 384, "tenant_demo", "pgvector"), + ] + assert calls["ensure_vector_extension"] == [ + ("resolved::postgresql://test", "pgvector", "public"), + ("resolved::postgresql://test", "pgvector", "tenant_demo"), + ] + assert calls["ensure_text_search_extension"] == [ + ("resolved::postgresql://test", "native", "public"), + ("resolved::postgresql://test", "native", "tenant_demo"), + ] + + +@pytest.mark.asyncio +async def test_run_migration_with_schema_only_runs_requested_schema(monkeypatch): + """run-db-migration with --schema should only migrate the requested schema.""" + monkeypatch.setenv("HINDSIGHT_API_DATABASE_URL", "postgresql://test") + calls: dict[str, list] = { + "run_migrations": [], + "ensure_vector_extension": [], + "ensure_text_search_extension": [], + } + + class MockTenantExtension: + async def list_tenants(self): + return [Tenant(schema="tenant_demo"), Tenant(schema="tenant_other")] + + async def fake_resolve_database_url(db_url: str) -> str: + return f"resolved::{db_url}" + + def fake_run_migrations(database_url: str, schema: str | None = None) -> None: + calls["run_migrations"].append((database_url, schema)) + + def fake_ensure_vector_extension( + database_url: str, + vector_extension: str = "pgvector", + schema: str | None = None, + ) -> None: + calls["ensure_vector_extension"].append((database_url, vector_extension, schema)) + + def fake_ensure_text_search_extension( + database_url: str, + text_search_extension: str = "native", + schema: str | None = None, + ) -> None: + calls["ensure_text_search_extension"].append((database_url, text_search_extension, schema)) + + monkeypatch.setattr(admin_cli, "load_extension", lambda *args, **kwargs: MockTenantExtension()) + monkeypatch.setattr(admin_cli, "resolve_database_url", fake_resolve_database_url) + + from hindsight_api import migrations as migrations_module + + monkeypatch.setattr(migrations_module, "run_migrations", fake_run_migrations) + monkeypatch.setattr(migrations_module, "ensure_vector_extension", fake_ensure_vector_extension) + monkeypatch.setattr(migrations_module, "ensure_text_search_extension", fake_ensure_text_search_extension) + + schemas = await admin_cli._run_migration("postgresql://test", schema="tenant_demo") + + assert schemas == ["tenant_demo"] + assert calls["run_migrations"] == [("resolved::postgresql://test", "tenant_demo")] + assert calls["ensure_vector_extension"] == [("resolved::postgresql://test", "pgvector", "tenant_demo")] + assert calls["ensure_text_search_extension"] == [("resolved::postgresql://test", "native", "tenant_demo")] diff --git a/hindsight-docs/docs/developer/admin-cli.md b/hindsight-docs/docs/developer/admin-cli.md index f2d59c29..e6d91e5f 100644 --- a/hindsight-docs/docs/developer/admin-cli.md +++ b/hindsight-docs/docs/developer/admin-cli.md @@ -16,7 +16,7 @@ uv add hindsight-api ### run-db-migration -Run database migrations to the latest version. This is useful when you want to run migrations separately from API startup (e.g., in CI/CD pipelines or before deploying a new version). +Run database migrations to the latest version. By default this migrates the base schema plus all tenant schemas discovered by the tenant extension. Use `--schema` for targeted migration of one schema. This is useful when you want to run migrations separately from API startup (e.g., in CI/CD pipelines or before deploying a new version). ```bash hindsight-admin run-db-migration [OPTIONS] @@ -26,12 +26,12 @@ hindsight-admin run-db-migration [OPTIONS] | Option | Description | Default | |--------|-------------|---------| -| `--schema`, `-s` | Database schema to run migrations on | `public` | +| `--schema`, `-s` | Database schema to run migrations on. If omitted, migrate the base schema plus all discovered tenant schemas. | All schemas | **Examples:** ```bash -# Run migrations on the default public schema +# Run migrations on the base schema plus all discovered tenant schemas hindsight-admin run-db-migration # Run migrations on a specific tenant schema @@ -189,4 +189,3 @@ The admin CLI uses the same environment variables as the API service. The most i export HINDSIGHT_API_DATABASE_URL=postgresql://user:pass@localhost:5432/hindsight hindsight-admin backup /backups/mybackup.zip ``` - diff --git a/hindsight-docs/docs/developer/configuration.md b/hindsight-docs/docs/developer/configuration.md index eb3d23e6..43f3368c 100644 --- a/hindsight-docs/docs/developer/configuration.md +++ b/hindsight-docs/docs/developer/configuration.md @@ -52,8 +52,10 @@ For high-concurrency workloads, increase `DB_POOL_MAX_SIZE`. Each concurrent rec To run migrations manually (e.g., before starting the API), use the admin CLI: ```bash +# Migrate the base schema plus all discovered tenant schemas hindsight-admin run-db-migration -# Or for a specific schema: + +# Or migrate a specific schema only: hindsight-admin run-db-migration --schema tenant_acme ``` diff --git a/skills/hindsight-docs/references/developer/admin-cli.md b/skills/hindsight-docs/references/developer/admin-cli.md index f2d59c29..e6d91e5f 100644 --- a/skills/hindsight-docs/references/developer/admin-cli.md +++ b/skills/hindsight-docs/references/developer/admin-cli.md @@ -16,7 +16,7 @@ uv add hindsight-api ### run-db-migration -Run database migrations to the latest version. This is useful when you want to run migrations separately from API startup (e.g., in CI/CD pipelines or before deploying a new version). +Run database migrations to the latest version. By default this migrates the base schema plus all tenant schemas discovered by the tenant extension. Use `--schema` for targeted migration of one schema. This is useful when you want to run migrations separately from API startup (e.g., in CI/CD pipelines or before deploying a new version). ```bash hindsight-admin run-db-migration [OPTIONS] @@ -26,12 +26,12 @@ hindsight-admin run-db-migration [OPTIONS] | Option | Description | Default | |--------|-------------|---------| -| `--schema`, `-s` | Database schema to run migrations on | `public` | +| `--schema`, `-s` | Database schema to run migrations on. If omitted, migrate the base schema plus all discovered tenant schemas. | All schemas | **Examples:** ```bash -# Run migrations on the default public schema +# Run migrations on the base schema plus all discovered tenant schemas hindsight-admin run-db-migration # Run migrations on a specific tenant schema @@ -189,4 +189,3 @@ The admin CLI uses the same environment variables as the API service. The most i export HINDSIGHT_API_DATABASE_URL=postgresql://user:pass@localhost:5432/hindsight hindsight-admin backup /backups/mybackup.zip ``` - diff --git a/skills/hindsight-docs/references/developer/configuration.md b/skills/hindsight-docs/references/developer/configuration.md index 4e642092..1c25c481 100644 --- a/skills/hindsight-docs/references/developer/configuration.md +++ b/skills/hindsight-docs/references/developer/configuration.md @@ -52,8 +52,10 @@ For high-concurrency workloads, increase `DB_POOL_MAX_SIZE`. Each concurrent rec To run migrations manually (e.g., before starting the API), use the admin CLI: ```bash +# Migrate the base schema plus all discovered tenant schemas hindsight-admin run-db-migration -# Or for a specific schema: + +# Or migrate a specific schema only: hindsight-admin run-db-migration --schema tenant_acme ```