From 34f7209877d63c2575cfd1da2ba75a17c7bc031d Mon Sep 17 00:00:00 2001 From: Peter Rothlaender Date: Mon, 13 Jul 2026 15:52:10 +0200 Subject: [PATCH 1/2] feat: guard pgvector startup DDL --- README.md | 3 + app/services/database.py | 123 ++++++++++++++++++-------------- tests/services/test_database.py | 32 +++++++-- 3 files changed, 97 insertions(+), 61 deletions(-) diff --git a/README.md b/README.md index 9b2b5036..86bfcd5e 100644 --- a/README.md +++ b/README.md @@ -81,6 +81,9 @@ The following environment variables are required to run the application: - `PG_POOL_PRE_PING`: (Optional) Set to "False" to disable SQLAlchemy's pre-ping check. Default is "True". When enabled, the connection pool issues a lightweight `SELECT 1` before handing out a pooled connection, so stale connections dropped by a remote server or middlebox idle timeout are transparently replaced instead of surfacing as query errors. Recommended for any deployment that connects to a remote PostgreSQL instance (managed Postgres, connections that traverse a load balancer, etc.). - `PG_POOL_RECYCLE`: (Optional) Maximum age in seconds of a pooled connection before it is recycled. Default is "-1" (disabled). Set to a positive value when the server enforces a hard idle or max-lifetime limit (e.g. "1800" for a 30-minute cap). - `POSTGRES_SCHEMA`: (Optional) Prepend this schema to the Postgres `search_path` so langchain's pgvector tables live in (and are read from) it. Unset by default (uses the user's default schema, typically `public`). Useful when sharing a database with other services — create the schema out-of-band first (`CREATE SCHEMA IF NOT EXISTS ; GRANT USAGE, CREATE ON SCHEMA TO ;`); the RAG API will not create it for you and fails fast at startup if the schema is missing. `public` is always appended to the resulting search path so the `vector` data type stays resolvable when the extension was installed there (the common case). Multiple schemas may be supplied as a comma-separated list (e.g. `myapp,extensions`) when the `vector` extension lives in a non-`public` schema. +- `PGVECTOR_CREATE_LEGACY_INDEXES`: (Optional) Set to "True" to create the legacy `custom_id` and `cmetadata->>'file_id'` indexes on startup. Default is "False". +- `PGVECTOR_MIGRATE_CMETADATA_JSONB`: (Optional) Set to "True" to migrate `langchain_pg_embedding.cmetadata` from JSON to JSONB on startup. Default is "False". +- `PGVECTOR_CREATE_CMETADATA_GIN_INDEX`: (Optional) Set to "True" to create the `cmetadata` JSONB GIN index on startup. Default is "False". - `RAG_HOST`: (Optional) The hostname or IP address where the API server will run. Defaults to "0.0.0.0" - `RAG_PORT`: (Optional) The port number where the API server will run. Defaults to port 8000. - `JWT_SECRET`: (Optional) The secret key used for verifying JWT tokens for requests. diff --git a/app/services/database.py b/app/services/database.py index 32ba2a76..83042c46 100644 --- a/app/services/database.py +++ b/app/services/database.py @@ -1,8 +1,17 @@ # app/services/database.py +import os + import asyncpg from app.config import DSN, logger +def _env_flag(name: str, default: bool = False) -> bool: + value = os.getenv(name) + if value is None: + return default + return value.lower() in ("true", "1", "yes", "y", "on") + + class PSQLDatabase: pool = None @@ -20,70 +29,74 @@ async def close_pool(cls): async def ensure_vector_indexes(): - """Ensure required indexes on langchain_pg_embedding and migrate cmetadata to JSONB. - - Runs at startup. Idempotent — safe to call repeatedly. - Operations: - 1. B-tree index on custom_id. - 2. Expression index on (cmetadata->>'file_id'). - 3. DDL migration: JSON -> JSONB for cmetadata (skipped if already JSONB). - 4. GIN index (jsonb_path_ops) on cmetadata for containment queries. - """ + """Ensure optional pgvector indexes/migrations when explicitly enabled.""" table_name = "langchain_pg_embedding" column_name = "custom_id" - # You might want to standardize the index naming convention index_name = f"idx_{table_name}_{column_name}" + create_legacy_indexes = _env_flag("PGVECTOR_CREATE_LEGACY_INDEXES") + migrate_cmetadata_jsonb = _env_flag("PGVECTOR_MIGRATE_CMETADATA_JSONB") + create_cmetadata_gin_index = _env_flag("PGVECTOR_CREATE_CMETADATA_GIN_INDEX") pool = await PSQLDatabase.get_pool() async with pool.acquire() as conn: - await conn.execute( - f""" - CREATE INDEX IF NOT EXISTS {index_name} ON {table_name} ({column_name}); - """ - ) - - # Expression index for (cmetadata->>'file_id') — critical for query - # performance. ExtendedPgVector overrides LangChain's default - # jsonb_path_match() to emit cmetadata->>'file_id' = ... which uses - # this B-tree index for fast equality lookups. - await conn.execute( - f""" - CREATE INDEX IF NOT EXISTS idx_{table_name}_file_id - ON {table_name} ((cmetadata->>'file_id')); - """ - ) - - # Migrate cmetadata from JSON to JSONB (idempotent — skipped if already JSONB). - # Rollback: ALTER TABLE langchain_pg_embedding ALTER COLUMN cmetadata TYPE JSON USING cmetadata::json; - # NOTE: table name is hardcoded below (not interpolated) to avoid SQL injection. - await conn.execute( - """ - DO $$ - BEGIN - IF EXISTS ( - SELECT 1 FROM information_schema.columns - WHERE table_name = 'langchain_pg_embedding' - AND table_schema = current_schema() - AND column_name = 'cmetadata' - AND data_type = 'json' - ) THEN - SET LOCAL lock_timeout = '10s'; - ALTER TABLE langchain_pg_embedding - ALTER COLUMN cmetadata TYPE JSONB USING cmetadata::jsonb; - END IF; - END - $$; + if create_legacy_indexes: + await conn.execute( + f""" + CREATE INDEX IF NOT EXISTS {index_name} ON {table_name} ({column_name}); """ - ) + ) - # GIN index on cmetadata for efficient JSONB filtering - await conn.execute( + # Expression index for (cmetadata->>'file_id') - critical for query + # performance. ExtendedPgVector emits cmetadata->>'file_id' = ... + # so this B-tree index supports fast equality lookups. + await conn.execute( + f""" + CREATE INDEX IF NOT EXISTS idx_{table_name}_file_id + ON {table_name} ((cmetadata->>'file_id')); """ - CREATE INDEX IF NOT EXISTS ix_cmetadata_gin - ON langchain_pg_embedding - USING gin (cmetadata jsonb_path_ops); - """ - ) + ) + else: + logger.info( + "Skipping legacy vector indexes; set PGVECTOR_CREATE_LEGACY_INDEXES=true to enable" + ) + + if migrate_cmetadata_jsonb: + await conn.execute( + """ + DO $$ + BEGIN + IF EXISTS ( + SELECT 1 FROM information_schema.columns + WHERE table_name = 'langchain_pg_embedding' + AND table_schema = current_schema() + AND column_name = 'cmetadata' + AND data_type = 'json' + ) THEN + SET LOCAL lock_timeout = '10s'; + ALTER TABLE langchain_pg_embedding + ALTER COLUMN cmetadata TYPE JSONB USING cmetadata::jsonb; + END IF; + END + $$; + """ + ) + else: + logger.info( + "Skipping cmetadata JSONB migration; set PGVECTOR_MIGRATE_CMETADATA_JSONB=true to enable" + ) + + if create_cmetadata_gin_index: + await conn.execute( + """ + CREATE INDEX IF NOT EXISTS ix_cmetadata_gin + ON langchain_pg_embedding + USING gin (cmetadata jsonb_path_ops); + """ + ) + else: + logger.info( + "Skipping cmetadata GIN index; set PGVECTOR_CREATE_CMETADATA_GIN_INDEX=true to enable" + ) logger.info("Vector database indexes ensured") diff --git a/tests/services/test_database.py b/tests/services/test_database.py index 18570482..3f816ca7 100644 --- a/tests/services/test_database.py +++ b/tests/services/test_database.py @@ -37,8 +37,20 @@ def acquire(self): return CapturingAcquire(self._conn) -def _run_with_captured_conn(monkeypatch): +DDL_FLAGS = ( + "PGVECTOR_CREATE_LEGACY_INDEXES", + "PGVECTOR_MIGRATE_CMETADATA_JSONB", + "PGVECTOR_CREATE_CMETADATA_GIN_INDEX", +) + + +def _run_with_captured_conn(monkeypatch, *enabled_flags): """Run ensure_vector_indexes() and return the captured connection.""" + for flag in DDL_FLAGS: + monkeypatch.delenv(flag, raising=False) + for flag in enabled_flags: + monkeypatch.setenv(flag, "true") + conn = CapturingConnection() pool = CapturingPool(conn) @@ -52,19 +64,27 @@ async def fake_get_pool(): def test_ensure_vector_indexes(monkeypatch): conn = _run_with_captured_conn(monkeypatch) - assert len(conn.statements) > 0 + assert conn.statements == [] + + +def test_ensure_vector_indexes_legacy_indexes_opt_in(monkeypatch): + conn = _run_with_captured_conn(monkeypatch, "PGVECTOR_CREATE_LEGACY_INDEXES") + + assert len(conn.statements) == 2 + assert "custom_id" in conn.statements[0] + assert "cmetadata->>'file_id'" in conn.statements[1] def test_ensure_vector_indexes_do_block_dollar_quoting(monkeypatch): """DO block must use $$ dollar-quoting, not single $.""" - conn = _run_with_captured_conn(monkeypatch) + conn = _run_with_captured_conn(monkeypatch, "PGVECTOR_MIGRATE_CMETADATA_JSONB") do_block = next(s for s in conn.statements if "DO" in s) assert "$$" in do_block, "DO block must use $$ dollar-quoting" def test_ensure_vector_indexes_jsonb_migration_sql(monkeypatch): """Migration block contains the correct ALTER COLUMN and schema filter.""" - conn = _run_with_captured_conn(monkeypatch) + conn = _run_with_captured_conn(monkeypatch, "PGVECTOR_MIGRATE_CMETADATA_JSONB") do_block = next(s for s in conn.statements if "DO" in s) assert "TYPE JSONB" in do_block assert "cmetadata::jsonb" in do_block @@ -73,14 +93,14 @@ def test_ensure_vector_indexes_jsonb_migration_sql(monkeypatch): def test_ensure_vector_indexes_lock_timeout(monkeypatch): """Migration sets a lock_timeout before ALTER TABLE.""" - conn = _run_with_captured_conn(monkeypatch) + conn = _run_with_captured_conn(monkeypatch, "PGVECTOR_MIGRATE_CMETADATA_JSONB") do_block = next(s for s in conn.statements if "DO" in s) assert "lock_timeout" in do_block def test_ensure_vector_indexes_gin_index(monkeypatch): """GIN index with jsonb_path_ops is created.""" - conn = _run_with_captured_conn(monkeypatch) + conn = _run_with_captured_conn(monkeypatch, "PGVECTOR_CREATE_CMETADATA_GIN_INDEX") gin_stmt = next(s for s in conn.statements if "ix_cmetadata_gin" in s) assert "jsonb_path_ops" in gin_stmt assert "USING gin" in gin_stmt From 5708a64d243f6cb8876d3256ad9b9181ae6e2300 Mon Sep 17 00:00:00 2001 From: Peter Rothlaender Date: Mon, 13 Jul 2026 18:23:07 +0200 Subject: [PATCH 2/2] chore: retrigger CI