Fix run-db-migration for all-tenant upgrades (#530)
* Add release-scoped migration admin command * Fix run-db-migration for all-tenant upgrades
This commit is contained in:
Chris Bartholomew
GitHub
parent
7bcf26097c
commit
9a694f64b8
8 changed files with 305 additions and 22 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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")]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
```
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
```
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
```
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
```
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue