diff --git a/README.md b/README.md index 05c4cb94..0b18db33 100644 --- a/README.md +++ b/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 ; 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. - `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/config.py b/app/config.py index 8ffab937..f518f5cb 100644 --- a/app/config.py +++ b/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 diff --git a/app/services/vector_store/factory.py b/app/services/vector_store/factory.py index c1f7a887..9320101d 100644 --- a/app/services/vector_store/factory.py +++ b/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 ;`)." + ) + 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( diff --git a/tests/services/test_vector_store_factory.py b/tests/services/test_vector_store_factory.py index 5f8ff220..56a80977 100644 --- a/tests/services/test_vector_store_factory.py +++ b/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= 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