feat(migrations): parallelize tenant schema migrations (#2203)

Add HINDSIGHT_API_MIGRATION_CONCURRENCY to migrate tenant schemas concurrently (each in its own spawn process; per-schema work stays sequential). NullPool on migration engines bounds per-worker connections. Validated at 20k schemas: no-op resweep ~60min->~11min (5x) at concurrency=12, peak 29 connections, 0 errors. Default 1 (sequential).

Author: nicoloboschi

.env.example

@@ -59,6 +59,7 @@ HINDSIGHT_API_LOG_LEVEL=info
 # HINDSIGHT_API_READ_DATABASE_URL=  # Optional read-replica URL. When set, recall queries (semantic, BM25, graph, temporal) flow through a separate pool against this URL, offloading the primary. Typically points to a read-only endpoint (CNPG's <cluster>-ro service or Aurora reader endpoint).
 # HINDSIGHT_API_MIGRATION_DATABASE_URL=  # Direct PostgreSQL URL for migrations (bypasses PgBouncer). Falls back to DATABASE_URL.
 # HINDSIGHT_API_DATABASE_SCHEMA=public  # PostgreSQL schema name (default: public)
+# HINDSIGHT_API_MIGRATION_CONCURRENCY=1  # Tenant schemas to migrate concurrently (PG only, each in its own process; per-schema work stays sequential). Each worker has ~1-2s startup cost + uses ~3 DB connections, so it only pays off with many schemas (tens+) or slow migrations; keep concurrency*3 <= spare max_connections. Default: 1 (sequential).
 
 # Vector Extension (Optional - uses pgvector by default)
 # Options: "pgvector" (default), "vchord", "pgvectorscale" (DiskANN)

hindsight-api-slim/hindsight_api/admin/cli.py

@@ -258,12 +258,7 @@ async def _run_migration(
     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,
-    )
+    from ..migrations import run_migrations_for_schemas
 
     is_pg0, instance_name, _ = parse_pg0_url(db_url)
     if is_pg0:
@@ -284,32 +279,21 @@ async def _run_migration(
         # Preserve order while removing duplicates.
         schemas = list(dict.fromkeys(schemas))
 
-    for schema in schemas:
-        run_migrations(resolved_url, schema=schema, migration_database_url=config.migration_database_url)
-
-    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,
-            pg_search_tokenizer=config.text_search_extension_pg_search_tokenizer,
-            schema=schema,
-        )
+    # Migrate up to `migration_concurrency` schemas at once (each in its own
+    # process); within a schema the work stays sequential. Run off the event
+    # loop so the process pool's blocking joins don't stall it.
+    await asyncio.to_thread(
+        run_migrations_for_schemas,
+        resolved_url,
+        schemas,
+        concurrency=config.migration_concurrency,
+        migration_database_url=config.migration_database_url,
+        embedding_dimension=embedding_dimension,
+        vector_extension=config.vector_extension,
+        text_search_extension=config.text_search_extension,
+        pg_search_tokenizer=config.text_search_extension_pg_search_tokenizer,
+        ensure_extensions=True,
+    )
 
     return schemas
 

hindsight-api-slim/hindsight_api/config.py

@@ -479,6 +479,7 @@ def normalize_config_dict(config: dict[str, Any]) -> dict[str, Any]:
 
 # Database migrations
 ENV_RUN_MIGRATIONS_ON_STARTUP = "HINDSIGHT_API_RUN_MIGRATIONS_ON_STARTUP"
+ENV_MIGRATION_CONCURRENCY = "HINDSIGHT_API_MIGRATION_CONCURRENCY"
 
 # Database connection pool
 ENV_DB_POOL_MIN_SIZE = "HINDSIGHT_API_DB_POOL_MIN_SIZE"
@@ -900,6 +901,10 @@ def _parse_strategy_boosts(raw: str | None) -> dict[str, str]:
 
 # Database migrations
 DEFAULT_RUN_MIGRATIONS_ON_STARTUP = True
+# Number of tenant schemas to migrate concurrently. Each schema runs in its own
+# process (Alembic's command.upgrade() is not thread-safe); within a schema the
+# work is always sequential. 1 = fully sequential (the safe default).
+DEFAULT_MIGRATION_CONCURRENCY = 1
 
 # Database connection pool
 DEFAULT_DB_POOL_MIN_SIZE = 5
@@ -1538,6 +1543,7 @@ class HindsightConfig:
 
     # Database migrations
     run_migrations_on_startup: bool
+    migration_concurrency: int
 
     # Database connection pool
     db_pool_min_size: int
@@ -2449,6 +2455,7 @@ def from_env(cls) -> "HindsightConfig":
             memory_defense=None,
             # Database migrations
             run_migrations_on_startup=os.getenv(ENV_RUN_MIGRATIONS_ON_STARTUP, "true").lower() == "true",
+            migration_concurrency=int(os.getenv(ENV_MIGRATION_CONCURRENCY, str(DEFAULT_MIGRATION_CONCURRENCY))),
             # Database connection pool
             db_pool_min_size=int(os.getenv(ENV_DB_POOL_MIN_SIZE, str(DEFAULT_DB_POOL_MIN_SIZE))),
             db_pool_max_size=int(os.getenv(ENV_DB_POOL_MAX_SIZE, str(DEFAULT_DB_POOL_MAX_SIZE))),

hindsight-api-slim/hindsight_api/engine/memory_engine.py

@@ -2605,50 +2605,39 @@ async def verify_llm():
             tenants = await self._tenant_extension.list_tenants()
             if tenants:
                 logger.info(f"Running migrations on {len(tenants)} schema(s)...")
-                for tenant in tenants:
-                    schema = tenant.schema
-                    if schema:
-                        schema = self._backend.normalize_schema(schema)
-                        self._backend.run_migrations(self.db_url, schema=schema)
-                logger.info("Schema migrations completed")
-
-            # PG-specific post-migration steps: ensure vector/text search extensions
-            # and embedding dimensions match configuration. These are no-ops for
-            # non-PG backends since they use different indexing strategies.
-            if self._backend.supports_bm25:
-                from ..migrations import (
-                    ensure_embedding_dimension,
-                    ensure_text_search_extension,
-                    ensure_vector_extension,
-                )
-
-                if tenants:
-                    for tenant in tenants:
-                        schema = tenant.schema
-                        if schema:
-                            ensure_embedding_dimension(
-                                self.db_url,
-                                self.embeddings.dimension,
-                                schema=schema,
-                                vector_extension=config.vector_extension,
-                            )
-
-                    for tenant in tenants:
-                        schema = tenant.schema
-                        if schema:
-                            ensure_vector_extension(
-                                self.db_url, vector_extension=config.vector_extension, schema=schema
-                            )
-
+                if self._database_backend_type == "postgresql":
+                    # PG: fan out across schemas (up to migration_concurrency, each
+                    # in its own process) and fold the PG-specific post-migration
+                    # extension/dimension sync into the same per-schema unit. Run
+                    # off the event loop so the process pool's blocking joins don't
+                    # stall it.
+                    from ..migrations import run_migrations_for_schemas
+
+                    schemas = [tenant.schema for tenant in tenants if tenant.schema]
+                    await asyncio.to_thread(
+                        run_migrations_for_schemas,
+                        self.db_url,
+                        schemas,
+                        concurrency=config.migration_concurrency,
+                        migration_database_url=config.migration_database_url,
+                        embedding_dimension=self.embeddings.dimension,
+                        vector_extension=config.vector_extension,
+                        text_search_extension=config.text_search_extension,
+                        pg_search_tokenizer=config.text_search_extension_pg_search_tokenizer,
+                        ensure_extensions=self._backend.supports_bm25,
+                    )
+                else:
+                    # Oracle and other backends: Alembic's non-thread-safe globals
+                    # and the absence of per-schema extension steps make parallelism
+                    # unnecessary; run sequentially via the backend's own runner.
+                    # normalize_schema() maps PG's "public" default to None (the
+                    # connecting user's schema) on Oracle.
                     for tenant in tenants:
-                        schema = tenant.schema
-                        if schema:
-                            ensure_text_search_extension(
-                                self.db_url,
-                                text_search_extension=config.text_search_extension,
-                                pg_search_tokenizer=config.text_search_extension_pg_search_tokenizer,
-                                schema=schema,
+                        if tenant.schema:
+                            self._backend.run_migrations(
+                                self.db_url, schema=self._backend.normalize_schema(tenant.schema)
                             )
+                logger.info("Schema migrations completed")
 
         logger.info(f"Connecting to database at {mask_network_location(self.db_url)}")
 

hindsight-api-slim/hindsight_api/migrations.py

@@ -27,6 +27,7 @@
 from alembic.script.revision import ResolutionError
 from alembic.util.exc import CommandError
 from sqlalchemy import Connection, create_engine, text
+from sqlalchemy.pool import NullPool
 
 from ._pg_search import normalize_pg_search_tokenizer, pg_search_bm25_columns
 from ._vector_index import (
@@ -247,7 +248,14 @@ def run_migrations(
         #   2. After acquiring the lock, COMMIT the transaction on the advisory-lock
         #      connection itself before running migrations.  pg_advisory_lock is
         #      session-level, so the lock survives the COMMIT.
-        engine = create_engine(migration_url)
+        # NullPool: do not retain the connection in a pool after the migration.
+        # Each schema migration opens a few short-lived engines (here plus the
+        # ensure_* steps); with the default QueuePool those connections linger
+        # until GC, and running many schemas in parallel (migration_concurrency)
+        # multiplies that footprint and exhausts max_connections — observed as
+        # "FATAL: sorry, too many clients already" sweeping 20k schemas at
+        # concurrency 12. NullPool closes the connection on return.
+        engine = create_engine(migration_url, poolclass=NullPool)
         with engine.connect() as conn:
             logger.debug(f"Acquiring migration advisory lock for schema '{schema_name}' (id={lock_id})...")
             while True:
@@ -400,7 +408,7 @@ def check_migration_status(
             return None, None
 
         # Get current revision from database
-        engine = create_engine(to_libpq_url(database_url))
+        engine = create_engine(to_libpq_url(database_url), poolclass=NullPool)
         with engine.connect() as connection:
             context = MigrationContext.configure(connection)
             current_rev = context.get_current_revision()
@@ -573,7 +581,7 @@ def ensure_embedding_dimension(
     """
     schema_name = schema or "public"
 
-    engine = create_engine(to_libpq_url(database_url))
+    engine = create_engine(to_libpq_url(database_url), poolclass=NullPool)
     with engine.connect() as conn:
         # Check if memory_units table exists (proxy for schema being initialized)
         table_exists = conn.execute(
@@ -622,7 +630,7 @@ def ensure_vector_extension(
     """
     schema_name = schema or "public"
 
-    engine = create_engine(to_libpq_url(database_url))
+    engine = create_engine(to_libpq_url(database_url), poolclass=NullPool)
     with engine.connect() as conn:
         # Detect which vector extension should be used
         target_ext = _detect_vector_extension(conn, vector_extension)
@@ -836,7 +844,7 @@ def ensure_text_search_extension(
     schema_name = schema or "public"
     pg_search_tokenizer = normalize_pg_search_tokenizer(pg_search_tokenizer)
 
-    engine = create_engine(to_libpq_url(database_url))
+    engine = create_engine(to_libpq_url(database_url), poolclass=NullPool)
     with engine.connect() as conn:
         # Tables with search_vector columns to check
         tables_to_check = [
@@ -1129,3 +1137,129 @@ def ensure_text_search_extension(
 
         conn.commit()
         logger.info(f"Successfully migrated text search to {text_search_extension}")
+
+
+def _migrate_one_schema_pg(
+    database_url: str,
+    schema: str,
+    *,
+    migration_database_url: str | None,
+    embedding_dimension: int | None,
+    vector_extension: str,
+    text_search_extension: str,
+    pg_search_tokenizer: str | None,
+    ensure_extensions: bool,
+) -> str:
+    """Run migrations + post-migration extension setup for a SINGLE PG schema.
+
+    Module-level (not a closure) so it is picklable and can run inside a
+    ``ProcessPoolExecutor`` worker. The steps run strictly in order — this is
+    the per-tenant sequential unit; parallelism happens only *across* schemas.
+    Returns the schema name on success; raises on the first failing step so the
+    caller can attribute the failure back to this schema.
+    """
+    run_migrations(database_url, schema=schema, migration_database_url=migration_database_url)
+    if embedding_dimension is not None:
+        ensure_embedding_dimension(
+            database_url,
+            embedding_dimension,
+            schema=schema,
+            vector_extension=vector_extension,
+        )
+    if ensure_extensions:
+        ensure_vector_extension(database_url, vector_extension=vector_extension, schema=schema)
+        ensure_text_search_extension(
+            database_url,
+            text_search_extension=text_search_extension,
+            schema=schema,
+            pg_search_tokenizer=pg_search_tokenizer,
+        )
+    return schema
+
+
+def _make_migration_executor(max_workers: int):
+    """Build the executor that runs per-schema migrations in parallel.
+
+    Each schema must run in its OWN process — Alembic's ``command.upgrade()``
+    uses non-thread-safe module globals (serialized in-process by
+    ``_alembic_lock``), so a thread pool would not actually run two upgrades at
+    once. ``spawn`` gives every worker a clean interpreter on all platforms,
+    avoiding the fork-of-a-multithreaded-process deadlock hazard (the API server
+    holds threads/pools when migrations run on startup).
+
+    Factored out so tests can substitute an in-process executor.
+    """
+    import multiprocessing
+    from concurrent.futures import ProcessPoolExecutor
+
+    return ProcessPoolExecutor(max_workers=max_workers, mp_context=multiprocessing.get_context("spawn"))
+
+
+def run_migrations_for_schemas(
+    database_url: str,
+    schemas: list[str],
+    *,
+    concurrency: int = 1,
+    migration_database_url: str | None = None,
+    embedding_dimension: int | None = None,
+    vector_extension: str = "pgvector",
+    text_search_extension: str = "native",
+    pg_search_tokenizer: str | None = None,
+    ensure_extensions: bool = True,
+) -> None:
+    """Run PostgreSQL migrations for many schemas, up to ``concurrency`` at once.
+
+    Within a schema the work is always sequential (migrate → embedding dim →
+    vector ext → text-search ext). Across schemas, when ``concurrency > 1`` each
+    schema is migrated in its OWN process: Alembic's ``command.upgrade()`` relies
+    on non-thread-safe module-level globals (serialized in-process by
+    ``_alembic_lock``), so threads would gain nothing — separate interpreters
+    each get a clean Alembic context. Per-schema advisory locks
+    (``_get_schema_lock_id``) keep concurrent processes from colliding on the
+    same schema across replicas.
+
+    ``database_url`` must already be resolved (e.g. an embedded ``pg0`` instance
+    started in the parent) — workers receive it verbatim and only connect.
+
+    Failures are collected per schema and re-raised together so one bad tenant
+    does not hide the status of the others.
+    """
+    if not schemas:
+        return
+
+    worker_kwargs = dict(
+        migration_database_url=migration_database_url,
+        embedding_dimension=embedding_dimension,
+        vector_extension=vector_extension,
+        text_search_extension=text_search_extension,
+        pg_search_tokenizer=pg_search_tokenizer,
+        ensure_extensions=ensure_extensions,
+    )
+
+    effective = max(1, min(concurrency, len(schemas)))
+    if effective == 1:
+        # Inline, in-process — no subprocess overhead for the common single
+        # tenant / sequential case (and keeps embedded pg0 dev simple).
+        for schema in schemas:
+            _migrate_one_schema_pg(database_url, schema, **worker_kwargs)
+        return
+
+    logger.info("Migrating %d schema(s) with concurrency=%d", len(schemas), effective)
+    errors: dict[str, BaseException] = {}
+    with _make_migration_executor(effective) as executor:
+        futures = {
+            executor.submit(_migrate_one_schema_pg, database_url, schema, **worker_kwargs): schema for schema in schemas
+        }
+        for future in futures:
+            schema = futures[future]
+            try:
+                future.result()
+            except Exception as exc:  # noqa: BLE001 — aggregate per-schema, re-raise below
+                errors[schema] = exc
+                logger.error("Migration failed for schema '%s': %s", schema, exc)
+
+    if errors:
+        failed = ", ".join(sorted(errors))
+        raise RuntimeError(
+            f"Database migrations failed for {len(errors)} of {len(schemas)} schema(s): {failed}"
+        ) from next(iter(errors.values()))