feat: guard pgvector startup DDL

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 <name>; GRANT USAGE, CREATE ON SCHEMA <name> TO <app_user>;`); 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.

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")
 

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