🗄️ feat: add POSTGRES_SCHEMA to pin pgvector tables to a custom schema

README.md

@@ -80,6 +80,7 @@ The following environment variables are required to run the application:
 - `PGVECTOR_CREATE_EXTENSION`: (Optional) Set to "False" to skip the `CREATE EXTENSION IF NOT EXISTS vector` call on startup. Default is "True". Use this when the `vector` extension is already installed on a managed Postgres (e.g. RDS, Azure Database for PostgreSQL) and the application user is not a superuser.
 - `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.
 - `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/config.py

@@ -69,6 +69,7 @@ def get_env_variable(
     "on",
 )
 PG_POOL_RECYCLE = int(get_env_variable("PG_POOL_RECYCLE", "-1"))
+POSTGRES_SCHEMA = get_env_variable("POSTGRES_SCHEMA", None) or None
 COLLECTION_NAME = get_env_variable("COLLECTION_NAME", "testcollection")
 ATLAS_MONGO_DB_URI = get_env_variable(
     "ATLAS_MONGO_DB_URI", "mongodb://127.0.0.1:27018/LibreChat"
@@ -369,6 +370,7 @@ def init_embeddings(provider, model, dimensions=None):
         create_extension=PGVECTOR_CREATE_EXTENSION,
         pool_pre_ping=PG_POOL_PRE_PING,
         pool_recycle=PG_POOL_RECYCLE,
+        schema=POSTGRES_SCHEMA,
     )
 elif VECTOR_DB_TYPE == VectorDBType.ATLAS_MONGO:
     # Backward compatability check

app/services/vector_store/factory.py

@@ -1,5 +1,5 @@
 import logging
-from typing import Any, Optional
+from typing import Any, List, Optional
 
 from pymongo import MongoClient
 from langchain_core.embeddings import Embeddings
@@ -14,6 +14,96 @@
 _mongo_client: Optional[MongoClient] = None
 
 
+def _parse_schemas(schema: str) -> List[str]:
+    """Split POSTGRES_SCHEMA's comma-separated value into a clean list."""
+    return [s.strip() for s in schema.split(",") if s.strip()]
+
+
+def _build_search_path(schemas: List[str]) -> str:
+    """Build a Postgres search_path value that includes every requested
+    schema plus `public` (appended if missing). pgvector installs the
+    `vector` data type into whatever schema its CREATE EXTENSION targets
+    (almost always `public`), and unqualified type names are resolved
+    against search_path — so bare `search_path=myapp` breaks
+    `CREATE TABLE ... vector(...)` with `type "vector" does not exist`.
+    """
+    parts = list(schemas)
+    if "public" not in parts:
+        parts.append("public")
+    return ",".join(parts)
+
+
+def _verify_schemas_exist(connection_string: str, schemas: List[str]) -> None:
+    """Raise if the POSTGRES_SCHEMA config won't let the app write to the
+    target schema or read types from its fallbacks.
+
+    schemas[0] is the write target — the schema pgvector's tables will land
+    in — so the role needs USAGE + CREATE there. schemas[1:] are read-only
+    search_path fallbacks for type/function resolution (e.g. an `extensions`
+    schema that holds the `vector` type), so USAGE alone is sufficient; a
+    role with CREATE here is fine too, but demanding it would reject the
+    common least-privilege setup where writes are intentionally confined to
+    the target schema.
+
+    Silent fallback is worse than failing fast here: PostgreSQL resolves
+    unqualified CREATE TABLE against the first schema in search_path where
+    the role has CREATE privileges, so a typo or a missing grant on the
+    target would land the pgvector tables in `public` (the entry appended
+    by _build_search_path) instead of the intended namespace, silently
+    defeating the isolation this feature is meant to provide.
+    """
+    from sqlalchemy import create_engine, text
+
+    engine = create_engine(connection_string)
+    try:
+        with engine.connect() as conn:
+            rows = conn.execute(
+                text(
+                    "SELECT s.schema_name, "
+                    "       has_schema_privilege(s.schema_name, 'USAGE') AS has_usage, "
+                    "       has_schema_privilege(s.schema_name, 'CREATE') AS has_create "
+                    "FROM information_schema.schemata s "
+                    "WHERE s.schema_name = ANY(:names)"
+                ),
+                {"names": schemas},
+            ).fetchall()
+            found = {row[0]: (row[1], row[2]) for row in rows}
+
+            target_schema = schemas[0]
+            fallback_schemas = schemas[1:]
+
+            missing = [s for s in schemas if s not in found]
+            # USAGE is required on every listed schema (target to write,
+            # fallbacks to look up types/functions at query time).
+            no_usage = [s for s in schemas if s in found and not found[s][0]]
+            # CREATE is only required on the target schema — the fallback
+            # entries exist solely to make type resolution work.
+            no_create_target = (
+                [target_schema]
+                if target_schema in found and not found[target_schema][1]
+                else []
+            )
+
+            problems = []
+            if missing:
+                problems.append(f"does not exist: {missing!r}")
+            if no_usage:
+                problems.append(f"role lacks USAGE on: {no_usage!r}")
+            if no_create_target:
+                problems.append(f"role lacks CREATE on target: {no_create_target!r}")
+
+            if problems:
+                hint_target = (missing or no_usage or no_create_target)[0]
+                raise ValueError(
+                    "POSTGRES_SCHEMA: " + "; ".join(problems) + ". "
+                    "Create/grant out-of-band first (e.g. "
+                    f"`CREATE SCHEMA IF NOT EXISTS {hint_target}; "
+                    f"GRANT USAGE, CREATE ON SCHEMA {hint_target} TO <app_user>;`)."
+                )
+    finally:
+        engine.dispose()
+
+
 def get_vector_store(
     connection_string: str,
     embeddings: Embeddings,
@@ -23,6 +113,7 @@ def get_vector_store(
     create_extension: bool = True,
     pool_pre_ping: bool = True,
     pool_recycle: int = -1,
+    schema: Optional[str] = None,
 ):
     """Create a vector store instance for the given mode.
 
@@ -39,12 +130,27 @@ def get_vector_store(
     error. pool_recycle<=0 disables periodic recycling (SQLAlchemy default);
     set a positive number of seconds when the server enforces an idle or
     max-lifetime limit.
+
+    Set schema to prepend that schema to every connection's search_path so
+    langchain's pgvector tables are created and queried there — use this to
+    keep the vector store logically separated when sharing a database with
+    other services. The schema must already exist (fails fast via
+    information_schema.schemata if missing). Multiple schemas may be
+    supplied as a comma-separated list; `public` is always appended so the
+    `vector` data type stays resolvable when the extension was created
+    there.
     """
     global _mongo_client
 
     engine_args: dict = {"pool_pre_ping": pool_pre_ping}
     if pool_recycle > 0:
         engine_args["pool_recycle"] = pool_recycle
+    if schema:
+        schemas = _parse_schemas(schema)
+        if schemas:
+            _verify_schemas_exist(connection_string, schemas)
+            search_path = _build_search_path(schemas)
+            engine_args["connect_args"] = {"options": f"-csearch_path={search_path}"}
 
     if mode == "sync":
         return ExtendedPgVector(

tests/services/test_vector_store_factory.py

@@ -186,6 +186,102 @@ def test_get_vector_store_propagates_pool_recycle():
         assert "pool_recycle" not in kwargs.get("engine_args", {})
 
 
+def test_get_vector_store_schema_sets_search_path():
+    """schema=<name> must translate into a connect_args search_path option so
+    pgvector's tables are created and queried in the requested schema.
+    `public` must be appended so the `vector` data type — installed by the
+    extension into whichever schema it was created in (usually `public`) —
+    remains resolvable; without it CREATE TABLE fails with
+    `type "vector" does not exist`."""
+    with patch(
+        "app.services.vector_store.factory._verify_schemas_exist"
+    ), patch("app.services.vector_store.factory.AsyncPgVector") as MockPG:
+        mock_embeddings = MagicMock()
+        factory.get_vector_store(
+            "conn", mock_embeddings, "coll", mode="async", schema="myapp"
+        )
+        _, kwargs = MockPG.call_args
+        engine_args = kwargs.get("engine_args")
+        assert engine_args is not None
+        assert engine_args["connect_args"]["options"] == "-csearch_path=myapp,public"
+        # schema must not displace the pool defaults already on engine_args
+        assert engine_args.get("pool_pre_ping") is True
+
+
+def test_get_vector_store_schema_preserves_user_supplied_public():
+    """If the user explicitly lists `public` we must not duplicate it."""
+    with patch(
+        "app.services.vector_store.factory._verify_schemas_exist"
+    ), patch("app.services.vector_store.factory.AsyncPgVector") as MockPG:
+        mock_embeddings = MagicMock()
+        factory.get_vector_store(
+            "conn", mock_embeddings, "coll", mode="async", schema="myapp,public"
+        )
+        _, kwargs = MockPG.call_args
+        assert (
+            kwargs["engine_args"]["connect_args"]["options"]
+            == "-csearch_path=myapp,public"
+        )
+
+
+def test_get_vector_store_schema_accepts_multi_schema_list():
+    """Comma-separated list for callers whose `vector` extension lives in a
+    non-public schema (e.g. a dedicated `extensions` schema)."""
+    with patch(
+        "app.services.vector_store.factory._verify_schemas_exist"
+    ), patch("app.services.vector_store.factory.AsyncPgVector") as MockPG:
+        mock_embeddings = MagicMock()
+        factory.get_vector_store(
+            "conn",
+            mock_embeddings,
+            "coll",
+            mode="async",
+            schema="myapp, extensions",
+        )
+        _, kwargs = MockPG.call_args
+        # whitespace around comma is stripped; public still auto-appended
+        assert (
+            kwargs["engine_args"]["connect_args"]["options"]
+            == "-csearch_path=myapp,extensions,public"
+        )
+
+
+def test_get_vector_store_schema_verifies_existence_before_engine_args():
+    """The factory must call _verify_schemas_exist with the parsed schemas
+    before constructing the vector store, so a typo in POSTGRES_SCHEMA fails
+    fast instead of silently creating tables in `public`."""
+    with patch(
+        "app.services.vector_store.factory._verify_schemas_exist"
+    ) as mock_verify, patch(
+        "app.services.vector_store.factory.AsyncPgVector"
+    ):
+        mock_embeddings = MagicMock()
+        factory.get_vector_store(
+            "conn://test",
+            mock_embeddings,
+            "coll",
+            mode="async",
+            schema="myapp, extensions",
+        )
+        mock_verify.assert_called_once_with("conn://test", ["myapp", "extensions"])
+
+
+def test_get_vector_store_schema_validation_error_propagates():
+    """If any configured schema is missing, _verify_schemas_exist raises and
+    we must surface that — not quietly skip validation and let tables land
+    in `public`."""
+    with patch(
+        "app.services.vector_store.factory._verify_schemas_exist",
+        side_effect=ValueError("schema 'typo' does not exist"),
+    ), patch("app.services.vector_store.factory.AsyncPgVector") as MockPG:
+        mock_embeddings = MagicMock()
+        with pytest.raises(ValueError, match="does not exist"):
+            factory.get_vector_store(
+                "conn", mock_embeddings, "coll", mode="async", schema="typo"
+            )
+        MockPG.assert_not_called()
+
+
 def test_load_file_content_cleans_up_on_lazy_load_failure():
     """cleanup_temp_encoding_file is called even when lazy_load() raises."""
     from app.routes.document_routes import load_file_content