From 130b49884b728165b2b4d02ab48a5d56d8f257a8 Mon Sep 17 00:00:00 2001 From: Vineeth Voruganti <13438633+VVoruganti@users.noreply.github.com> Date: Tue, 12 May 2026 18:08:28 -0400 Subject: [PATCH] feat(scripts): add configure_embeddings bootstrap CLI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds scripts/configure_embeddings.py alongside the other one-off scripts (provision_db, migrate_db, generate_jwt_secret, etc.). Invoked as `uv run python scripts/configure_embeddings.py` — same convention as the existing scripts in that directory, including the sys.path shim that lets src.* imports resolve when run directly. Bootstrap step for self-hosted installs at a non-default EMBEDDING_VECTOR_DIMENSIONS — runs between `alembic upgrade head` and starting the API/deriver. pgvector ALTER safety (single transaction): - LOCK TABLE {schema}.documents, {schema}.message_embeddings IN ACCESS EXCLUSIVE MODE — closes the TOCTOU window between population check and ALTER - COUNT(*) WHERE embedding IS NOT NULL on both tables; refuse with a non-zero exit if either is populated (ALTER ... USING NULL would silently wipe those vectors) - Snapshot HNSW index DDL from pg_indexes; drop, ALTER, recreate from the captured DDL so operator-set HNSW params (m, ef_construction) survive the round trip External vector stores (turbopuffer, lancedb) are never created or modified — namespaces are per-workspace and lazy-created on first write. The --report mode enumerates workspaces and collections from the application DB, derives the expected namespaces via get_vector_namespace(), and prints a per-namespace status table. CLI modes (mutually exclusive): - (default) interactive: print plan, prompt to confirm - --dry-run: print plan and exit 0 without touching the DB - --yes: apply without prompt - --report: print external-store namespace inventory and exit Also updates src/startup/embedding_validator.py error-message paths and docs/v3/contributing/configuration.mdx invocations to point at the new script location. Tests cover plan no-op, plan needs-alter, plan raises on missing column, ALTER + HNSW round-trip, refuse-when-populated (monkeypatched count to avoid wiring the full workspace/peer/collection/document FK chain just to land one vector row), and idempotency. --- docs/v3/contributing/configuration.mdx | 44 ++- scripts/configure_embeddings.py | 438 +++++++++++++++++++++ src/startup/embedding_validator.py | 9 +- tests/scripts/__init__.py | 0 tests/scripts/test_configure_embeddings.py | 178 +++++++++ 5 files changed, 662 insertions(+), 7 deletions(-) create mode 100644 scripts/configure_embeddings.py create mode 100644 tests/scripts/__init__.py create mode 100644 tests/scripts/test_configure_embeddings.py diff --git a/docs/v3/contributing/configuration.mdx b/docs/v3/contributing/configuration.mdx index 02916ab6..fb21f3f2 100644 --- a/docs/v3/contributing/configuration.mdx +++ b/docs/v3/contributing/configuration.mdx @@ -218,8 +218,48 @@ EMBEDDING_MODEL_CONFIG__OVERRIDES__BASE_URL=http://localhost:8000/v1 EMBEDDING_MODEL_CONFIG__OVERRIDES__API_KEY_ENV=EMBEDDING_CUSTOM_API_KEY ``` -Current constraint: -- `EMBEDDING_VECTOR_DIMENSIONS` can be changed for fully migrated external vector stores, but pgvector and dual-write mode still require `1536` until the schema migration lands. +Forwarding `dimensions=` to OpenAI-compatible providers is controlled by `EMBEDDING_MODEL_CONFIG__DIMENSIONS_MODE`: + +- `auto` (default): forwards `dimensions=` when you have explicitly set `EMBEDDING_VECTOR_DIMENSIONS` to a non-default value, and the configured model is not on the known-rejecting list (currently `text-embedding-ada-002`). Deployments that stay on the default keep their existing behavior. +- `always`: always forward. Use for OpenAI-compatible self-hosted providers that require it, and for same-as-default truncation (e.g. `text-embedding-3-large` truncated to 1536 — `auto` would skip the parameter because 1536 is the default). +- `never`: never forward. Explicit opt-out for providers that reject the parameter. + +#### Bootstrapping non-default dimensions + +`EMBEDDING_VECTOR_DIMENSIONS` is treated as immutable for the life of a deployment. The pgvector schema is dim-pinned by Alembic at `1536` by default; if you want a different dim, you must ALTER the empty columns once at bootstrap time. + +Install order for a non-default dim: + +```bash +# 1. Apply migrations (creates default vector(1536) schema) +alembic upgrade head + +# 2. Set the dim you want +export EMBEDDING_VECTOR_DIMENSIONS=768 + +# 3. ALTER the empty columns to the target dim +uv run python scripts/configure_embeddings.py --dry-run # preview +uv run python scripts/configure_embeddings.py --yes # apply + +# 4. Start API and deriver — both run the startup validator and refuse +# to serve traffic if the schema and EMBEDDING_VECTOR_DIMENSIONS disagree. +``` + +Existing deployments at 1536 with `text-embedding-3-small` need no action — step 3 detects matching dims and skips. + +The script refuses to ALTER tables that already contain non-null embeddings. To switch dim or model on a populated deployment, stand up a new deployment at the new configuration and migrate data out of band; there is no in-place re-embedding affordance. + +External vector stores (Turbopuffer, LanceDB) do not need bootstrap setup. Namespaces are per-workspace and lazy-created on first write at whatever dim the embedding client returns. Use `--report` to inventory the existing namespaces against the configured dim: + +```bash +uv run python scripts/configure_embeddings.py --report +``` + +The startup validator at `src/startup/embedding_validator.py` enforces the dim invariant at boot for both the API (`src/main.py` lifespan) and the deriver (`src/deriver/__main__.py`). A mismatch crashes the process with an actionable error before any HTTP route is served or any queue task is processed. + +`VECTOR_STORE_DIMENSIONS` is **deprecated**. `EMBEDDING_VECTOR_DIMENSIONS` is the single source of truth; setting `VECTOR_STORE_DIMENSIONS` explicitly emits a startup warning and is otherwise ignored. The field will be removed in a future release; drop it from your `.env` to silence the warning. + +The `VECTOR_STORE_MIGRATED` flag still exists and still controls dual-write / cutover semantics for legacy tenants moving between storage backends (pgvector ↔ turbopuffer ↔ lancedb). It is unrelated to dimension configuration after this release. ### Feature-Specific Model Configuration diff --git a/scripts/configure_embeddings.py b/scripts/configure_embeddings.py new file mode 100644 index 00000000..bde9a373 --- /dev/null +++ b/scripts/configure_embeddings.py @@ -0,0 +1,438 @@ +"""Configure pgvector schema dim to match EMBEDDING_VECTOR_DIMENSIONS. + +Usage:: + + uv run python scripts/configure_embeddings.py # interactive + uv run python scripts/configure_embeddings.py --dry-run # print intent, no DB write + uv run python scripts/configure_embeddings.py --yes # apply without prompt + uv run python scripts/configure_embeddings.py --report # full external-store inventory + +The bootstrap sequence for a self-hosted install is: + + 1. alembic upgrade head # creates default vector(1536) schema + 2. uv run python scripts/configure_embeddings.py # ALTER columns to target dim + 3. start the API and deriver # validators refuse to start on mismatch + +Existing 1536 deployments need no action — step 2 is a no-op when settings +already match the schema. + +This script never creates or modifies external-store namespaces. Turbopuffer +and LanceDB namespaces are per-workspace and lazy-created on first write by +application code; their dim is implicitly pinned at that point. Use +``--report`` to enumerate existing namespaces against the configured dim. +""" + +from __future__ import annotations + +import argparse +import asyncio +import logging +import os +import sys +from dataclasses import dataclass + +# Match the path-shim convention used by the other scripts in this directory +# so `src.*` imports resolve when the script is run directly. +_PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +if _PROJECT_ROOT not in sys.path: + sys.path.insert(0, _PROJECT_ROOT) + +from sqlalchemy import text # noqa: E402 +from sqlalchemy.ext.asyncio import AsyncConnection, AsyncEngine # noqa: E402 + +from src.config import settings # noqa: E402 +from src.db import engine # noqa: E402 + +logger = logging.getLogger(__name__) + +_EMBEDDING_TABLES: tuple[str, ...] = ("documents", "message_embeddings") + + +# --------------------------------------------------------------------------- +# Result types +# --------------------------------------------------------------------------- + + +@dataclass(frozen=True) +class _PgvectorPlan: + target_dim: int + schema: str + current_dims: dict[str, int] + needs_alter: bool + + +@dataclass(frozen=True) +class _NamespaceRecord: + """A single row in the --report output.""" + + namespace: str + status: str # one of: "ok", "missing", "mismatch", "unknown" + actual_dim: int | None + target_dim: int + + +# --------------------------------------------------------------------------- +# pgvector phase +# --------------------------------------------------------------------------- + + +async def _introspect_pgvector(conn: AsyncConnection, schema: str) -> dict[str, int]: + """Return ``{table_name: atttypmod}`` for embedding columns in ``schema``. + + Tables not present in the result dict are absent from the schema. + pgvector stores the declared dim directly in ``atttypmod`` (no VARHDRSZ). + """ + query = text( + """ + SELECT c.relname AS table_name, a.atttypmod AS typmod + FROM pg_attribute a + JOIN pg_class c ON a.attrelid = c.oid + JOIN pg_namespace n ON c.relnamespace = n.oid + WHERE n.nspname = :schema + AND c.relname = ANY(:tables) + AND a.attname = 'embedding' + """ + ) + result = await conn.execute( + query, + {"schema": schema, "tables": list(_EMBEDDING_TABLES)}, + ) + return {row.table_name: row.typmod for row in result} + + +async def _build_pgvector_plan( + engine: AsyncEngine, target_dim: int, schema: str +) -> _PgvectorPlan: + """Build a plan describing what (if anything) the script will change.""" + async with engine.connect() as conn: + current = await _introspect_pgvector(conn, schema) + + missing = set(_EMBEDDING_TABLES) - current.keys() + if missing: + listing = ", ".join(sorted(f"{schema}.{t}.embedding" for t in missing)) + raise SystemExit( + f"error: required vector columns missing: {listing}." + + " Run `alembic upgrade head` first." + ) + for table, typmod in current.items(): + if typmod == -1: + raise SystemExit( + f"error: {schema}.{table}.embedding has no declared vector" + + " dimension (unbounded typmod). Drop and recreate the column" + + " or restore from a versioned migration before re-running." + ) + + needs_alter = any(typmod != target_dim for typmod in current.values()) + return _PgvectorPlan( + target_dim=target_dim, + schema=schema, + current_dims=current, + needs_alter=needs_alter, + ) + + +async def _count_non_null_embeddings( + conn: AsyncConnection, schema: str, table: str +) -> int: + query = text( + f'SELECT COUNT(*) AS n FROM "{schema}"."{table}" WHERE embedding IS NOT NULL' + ) + result = await conn.execute(query) + row = result.first() + return int(row.n) if row is not None else 0 + + +async def _fetch_hnsw_index_defs( + conn: AsyncConnection, schema: str +) -> list[tuple[str, str]]: + """Return ``(index_name, CREATE INDEX ...)`` for HNSW indices on the + embedding columns. We re-CREATE them after the ALTER using these exact + definitions, preserving operator-set params (m, ef_construction, etc.).""" + query = text( + """ + SELECT indexname AS name, indexdef AS ddl + FROM pg_indexes + WHERE schemaname = :schema + AND tablename = ANY(:tables) + AND indexdef ILIKE '%USING hnsw%' + """ + ) + result = await conn.execute( + query, {"schema": schema, "tables": list(_EMBEDDING_TABLES)} + ) + return [(row.name, row.ddl) for row in result] + + +async def _apply_pgvector_alter(engine: AsyncEngine, plan: _PgvectorPlan) -> None: + """ALTER the embedding columns to ``plan.target_dim`` in a single + transaction. Refuses to proceed if any non-null embeddings exist. + + Sequence (inside the transaction): + 1. LOCK TABLE ... IN ACCESS EXCLUSIVE MODE — closes the TOCTOU window + between the population check and the ALTER. + 2. SELECT COUNT(embedding IS NOT NULL) per table — refuse if any > 0. + 3. Save HNSW index definitions, then DROP them (cannot ALTER under HNSW). + 4. ALTER ... ALTER COLUMN embedding TYPE vector(N) USING NULL. + 5. Recreate HNSW indices from saved definitions. + """ + async with engine.begin() as conn: + # Step 1: lock both tables for the duration of the transaction. + for table in _EMBEDDING_TABLES: + await conn.execute( + text(f'LOCK TABLE "{plan.schema}"."{table}" IN ACCESS EXCLUSIVE MODE') + ) + + # Step 2: population check. + counts: dict[str, int] = {} + for table in _EMBEDDING_TABLES: + counts[table] = await _count_non_null_embeddings(conn, plan.schema, table) + populated = {t: n for t, n in counts.items() if n > 0} + if populated: + detail = ", ".join(f"{t}: {n} rows" for t, n in sorted(populated.items())) + raise SystemExit( + f"error: refusing to ALTER populated embedding tables ({detail})." + + " This script only configures empty tables. Re-embed out-of-band" + + " into a fresh deployment, then cut over." + ) + + # Step 3: snapshot + drop HNSW indices. + index_defs = await _fetch_hnsw_index_defs(conn, plan.schema) + for index_name, _ddl in index_defs: + logger.info("dropping HNSW index %s", index_name) + await conn.execute(text(f'DROP INDEX "{plan.schema}"."{index_name}"')) + + # Step 4: ALTER columns. + for table in _EMBEDDING_TABLES: + logger.info( + "altering %s.%s.embedding to vector(%d)", + plan.schema, + table, + plan.target_dim, + ) + await conn.execute( + text( + f'ALTER TABLE "{plan.schema}"."{table}"' + + f" ALTER COLUMN embedding TYPE vector({plan.target_dim})" + + " USING NULL" + ) + ) + + # Step 5: recreate HNSW indices from the saved definitions. + for index_name, ddl in index_defs: + logger.info("recreating HNSW index %s", index_name) + await conn.execute(text(ddl)) + + +# --------------------------------------------------------------------------- +# External-store report +# --------------------------------------------------------------------------- + + +async def _enumerate_workspaces(conn: AsyncConnection) -> list[str]: + result = await conn.execute(text("SELECT name FROM workspaces ORDER BY created_at")) + return [row.name for row in result] + + +async def _enumerate_collections( + conn: AsyncConnection, +) -> list[tuple[str, str, str]]: + result = await conn.execute( + text("SELECT workspace_name, observer, observed FROM collections") + ) + return [(row.workspace_name, row.observer, row.observed) for row in result] + + +async def _build_external_namespace_inventory( + engine: AsyncEngine, +) -> list[tuple[str, str]]: + """Return ``(namespace_type, namespace_name)`` pairs for every namespace + that should exist based on the application DB. Message namespaces are + derived per workspace, document namespaces per collection row. + """ + from src.vector_store import get_external_vector_store + + store = get_external_vector_store() + if store is None: + return [] + + async with engine.connect() as conn: + workspace_names = await _enumerate_workspaces(conn) + collection_keys = await _enumerate_collections(conn) + + pairs: list[tuple[str, str]] = [] + for workspace_name in workspace_names: + pairs.append(("message", store.get_vector_namespace("message", workspace_name))) + for workspace_name, observer, observed in collection_keys: + pairs.append( + ( + "document", + store.get_vector_namespace( + "document", workspace_name, observer=observer, observed=observed + ), + ) + ) + return pairs + + +async def _probe_namespace_dim(store: object, namespace: str) -> int | None: + """Best-effort: return the namespace's declared dim if introspectable. + + Returns ``None`` if the SDK does not expose a uniform dim accessor for + the configured store. Future work can specialize per store; today the + pgvector validator is the load-bearing dim safety. + """ + _ = (store, namespace) + return None + + +async def _emit_report(engine: AsyncEngine, target_dim: int) -> int: + """Print the per-namespace inventory and return an exit code. 0 on a + clean report (all matching or missing); non-zero on any mismatch. + """ + if settings.VECTOR_STORE.TYPE == "pgvector": + print("--report has no effect with VECTOR_STORE_TYPE=pgvector") + return 0 + + inventory = await _build_external_namespace_inventory(engine) + if not inventory: + print( + "no external namespaces to inventory" + + " (no workspaces/collections exist yet, or no external store configured)" + ) + return 0 + + from src.vector_store import get_external_vector_store + + store = get_external_vector_store() + + records: list[_NamespaceRecord] = [] + for _ns_type, namespace in inventory: + actual = await _probe_namespace_dim(store, namespace) if store else None + if actual is None: + status = "unknown" + elif actual == target_dim: + status = "ok" + else: + status = "mismatch" + records.append( + _NamespaceRecord( + namespace=namespace, + status=status, + actual_dim=actual, + target_dim=target_dim, + ) + ) + + width = max(len(r.namespace) for r in records) + print(f"{'namespace'.ljust(width)} status dim") + print(f"{'-' * width} --------- ------") + for r in records: + dim_str = "?" if r.actual_dim is None else str(r.actual_dim) + print(f"{r.namespace.ljust(width)} {r.status:<9} {dim_str}") + + mismatches = [r for r in records if r.status == "mismatch"] + if mismatches: + print( + f"\nerror: {len(mismatches)} namespace(s) have dim != {target_dim}", + file=sys.stderr, + ) + return 1 + return 0 + + +# --------------------------------------------------------------------------- +# CLI +# --------------------------------------------------------------------------- + + +def _build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + prog="configure_embeddings", + description=( + "Configure pgvector schema dim to match EMBEDDING_VECTOR_DIMENSIONS." + ), + ) + mode = parser.add_mutually_exclusive_group() + mode.add_argument( + "--dry-run", + action="store_true", + help="print intended changes and exit without touching the DB", + ) + mode.add_argument( + "--yes", + action="store_true", + help="apply changes without an interactive prompt", + ) + mode.add_argument( + "--report", + action="store_true", + help="print external-store namespace inventory and exit", + ) + return parser + + +def _confirm(prompt: str) -> bool: + response = input(f"{prompt} [y/N]: ").strip().lower() + return response in {"y", "yes"} + + +async def _async_main(args: argparse.Namespace) -> int: + target_dim = settings.EMBEDDING.VECTOR_DIMENSIONS + schema = settings.DB.SCHEMA + + if args.report: + return await _emit_report(engine, target_dim) + + plan = await _build_pgvector_plan(engine, target_dim, schema) + if not plan.needs_alter: + print( + f"pgvector: {schema}.documents.embedding and" + + f" {schema}.message_embeddings.embedding already at dim {target_dim}," + + " skipping ALTER" + ) + return await _emit_report(engine, target_dim) + + current_summary = ", ".join( + f"{schema}.{t}.embedding={plan.current_dims[t]}" for t in _EMBEDDING_TABLES + ) + print(f"target dim: {target_dim}") + print(f"current: {current_summary}") + print("planned operations (single transaction):") + print(f" - LOCK TABLE {schema}.documents IN ACCESS EXCLUSIVE MODE") + print(f" - LOCK TABLE {schema}.message_embeddings IN ACCESS EXCLUSIVE MODE") + print(" - refuse if any non-null embeddings exist") + print(" - DROP existing HNSW indices on the embedding columns") + print( + f" - ALTER COLUMN embedding TYPE vector({target_dim}) USING NULL" + + " on both tables" + ) + print(" - CREATE HNSW indices from snapshotted definitions") + + if args.dry_run: + print("\n--dry-run: no changes applied") + return 0 + + if not args.yes and not _confirm("apply?"): + print("aborted") + return 1 + + await _apply_pgvector_alter(engine, plan) + print(f"\npgvector schema is now at dim {target_dim}") + return await _emit_report(engine, target_dim) + + +def main(argv: list[str] | None = None) -> int: + logging.basicConfig(level=logging.INFO, format="%(message)s") + parser = _build_parser() + args = parser.parse_args(argv) + try: + return asyncio.run(_async_main(args)) + finally: + # Best-effort cleanup; the script is short-lived so a leak here is + # harmless but keeping the dispose explicit avoids "engine not + # disposed" warnings during tests. + asyncio.run(engine.dispose()) + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/startup/embedding_validator.py b/src/startup/embedding_validator.py index a9982cfe..82c6f53e 100644 --- a/src/startup/embedding_validator.py +++ b/src/startup/embedding_validator.py @@ -8,8 +8,7 @@ actual state instead of operator-asserted state. For external stores (turbopuffer, lancedb) the check is best-effort: namespaces are per-workspace and lazy-created, so this validator can only sample existing -ones. Full enumeration is available via `uv run python -m -src.scripts.configure_embeddings --report` (Phase 3). +ones. Full enumeration is available via `uv run python scripts/configure_embeddings.py --report`. """ from __future__ import annotations @@ -139,7 +138,7 @@ def _assert_pgvector_dims_match( raise StartupValidationError( f"{schema}.{table}.embedding has no declared vector dimension" + " (unbounded typmod). Run" - + " `uv run python -m src.scripts.configure_embeddings`." + + " `uv run python scripts/configure_embeddings.py`." ) # pgvector stores the declared dim directly in atttypmod (no VARHDRSZ). actual = atttypmod @@ -147,7 +146,7 @@ def _assert_pgvector_dims_match( raise StartupValidationError( f"{schema}.{table}.embedding dim ({actual}) does not match" + f" EMBEDDING_VECTOR_DIMENSIONS ({target_dim}). Run" - + " `uv run python -m src.scripts.configure_embeddings`" + + " `uv run python scripts/configure_embeddings.py`" + " or fix EMBEDDING_VECTOR_DIMENSIONS." ) @@ -191,7 +190,7 @@ async def _sample_external_namespaces(engine: AsyncEngine, *, target_dim: int) - raise StartupValidationError( f"Existing external-store namespaces have dim != {target_dim}:" + f" {formatted}. Run" - + " `uv run python -m src.scripts.configure_embeddings --report`." + + " `uv run python scripts/configure_embeddings.py --report`." ) diff --git a/tests/scripts/__init__.py b/tests/scripts/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/scripts/test_configure_embeddings.py b/tests/scripts/test_configure_embeddings.py new file mode 100644 index 00000000..c35414f9 --- /dev/null +++ b/tests/scripts/test_configure_embeddings.py @@ -0,0 +1,178 @@ +"""Phase 3: configure_embeddings script tests.""" + +from __future__ import annotations + +from collections.abc import AsyncGenerator +from contextlib import asynccontextmanager + +import pytest +from sqlalchemy import text +from sqlalchemy.ext.asyncio import AsyncEngine + +from scripts.configure_embeddings import ( + _apply_pgvector_alter, # pyright: ignore[reportPrivateUsage] + _build_pgvector_plan, # pyright: ignore[reportPrivateUsage] +) + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@asynccontextmanager +async def _restore_schema_to(db_engine: AsyncEngine, dim: int) -> AsyncGenerator[None]: + """ALTER both embedding columns back to ``dim`` on exit so this test + leaves the shared test DB in a consistent state for subsequent tests.""" + try: + yield + finally: + async with db_engine.begin() as conn: + for table in ("documents", "message_embeddings"): + await conn.execute( + text( + f"ALTER TABLE {table} ALTER COLUMN embedding" + + f" TYPE vector({dim}) USING NULL" + ) + ) + + +async def _current_dims(db_engine: AsyncEngine) -> dict[str, int]: + async with db_engine.connect() as conn: + result = await conn.execute( + text( + """ + SELECT c.relname AS table_name, a.atttypmod AS typmod + FROM pg_attribute a + JOIN pg_class c ON a.attrelid = c.oid + JOIN pg_namespace n ON c.relnamespace = n.oid + WHERE n.nspname = 'public' + AND c.relname = ANY(:tables) + AND a.attname = 'embedding' + """ + ), + {"tables": ["documents", "message_embeddings"]}, + ) + return {row.table_name: row.typmod for row in result} + + +async def _hnsw_indexes(db_engine: AsyncEngine) -> set[str]: + async with db_engine.connect() as conn: + result = await conn.execute( + text( + """ + SELECT indexname + FROM pg_indexes + WHERE schemaname = 'public' + AND tablename IN ('documents', 'message_embeddings') + AND indexdef ILIKE '%USING hnsw%' + """ + ) + ) + return {row.indexname for row in result} + + +# --------------------------------------------------------------------------- +# Plan +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_plan_no_alter_needed_when_dims_already_match( + db_engine: AsyncEngine, +) -> None: + plan = await _build_pgvector_plan(db_engine, target_dim=1536, schema="public") + assert plan.needs_alter is False + assert plan.current_dims == {"documents": 1536, "message_embeddings": 1536} + + +@pytest.mark.asyncio +async def test_plan_needs_alter_when_target_differs(db_engine: AsyncEngine) -> None: + plan = await _build_pgvector_plan(db_engine, target_dim=768, schema="public") + assert plan.needs_alter is True + + +@pytest.mark.asyncio +async def test_plan_raises_on_missing_column(db_engine: AsyncEngine) -> None: + with pytest.raises(SystemExit, match="required vector columns missing"): + await _build_pgvector_plan(db_engine, target_dim=1536, schema="no_such_schema") + + +# --------------------------------------------------------------------------- +# Apply +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_apply_alters_dims_and_recreates_hnsw_indexes( + db_engine: AsyncEngine, +) -> None: + # 768 is the canonical "small" dim used in non-1536 deployments and is + # well below pgvector's 2000-dim HNSW limit. + target = 768 + async with _restore_schema_to(db_engine, dim=1536): + before_indexes = await _hnsw_indexes(db_engine) + assert before_indexes, "test fixture should have HNSW indexes pre-alter" + + plan = await _build_pgvector_plan(db_engine, target_dim=target, schema="public") + assert plan.needs_alter is True + await _apply_pgvector_alter(db_engine, plan) + + after_dims = await _current_dims(db_engine) + assert after_dims == {"documents": target, "message_embeddings": target} + + after_indexes = await _hnsw_indexes(db_engine) + assert ( + after_indexes == before_indexes + ), "HNSW indexes should be recreated with the same names" + + +@pytest.mark.asyncio +async def test_apply_refuses_when_embeddings_populated( + db_engine: AsyncEngine, + monkeypatch: pytest.MonkeyPatch, +) -> None: + """ALTER ... USING NULL would silently wipe non-null embeddings, so the + pre-check must abort the transaction before any destructive action. + + We patch the count helper to simulate populated tables rather than wire + up the full FK chain of workspace/peer/collection/document just to land + one vector row. + """ + + async def fake_count(_conn: object, _schema: str, table: str) -> int: + return 7 if table == "documents" else 0 + + monkeypatch.setattr( + "scripts.configure_embeddings._count_non_null_embeddings", + fake_count, + ) + + async with _restore_schema_to(db_engine, dim=1536): + plan = await _build_pgvector_plan(db_engine, target_dim=768, schema="public") + with pytest.raises( + SystemExit, match="refusing to ALTER populated embedding tables" + ): + await _apply_pgvector_alter(db_engine, plan) + + # The SystemExit aborts the transaction; nothing should have changed. + dims_after_refuse = await _current_dims(db_engine) + assert dims_after_refuse == { + "documents": 1536, + "message_embeddings": 1536, + } + + +# --------------------------------------------------------------------------- +# Idempotency +# --------------------------------------------------------------------------- + + +@pytest.mark.asyncio +async def test_idempotent_apply_is_a_noop(db_engine: AsyncEngine) -> None: + """Build plan twice with the matching dim — second call should still + return needs_alter=False without raising or making any changes.""" + plan_a = await _build_pgvector_plan(db_engine, target_dim=1536, schema="public") + plan_b = await _build_pgvector_plan(db_engine, target_dim=1536, schema="public") + assert plan_a.needs_alter is False + assert plan_b.needs_alter is False + assert plan_a.current_dims == plan_b.current_dims