Fix non-default schema in KV store (#4655)

* Fix non-default schema in KV store

* Fix custom schema

Author: Weves

backend/ee/onyx/server/saml.py

@@ -28,6 +28,7 @@
 from onyx.db.auth import get_user_count
 from onyx.db.auth import get_user_db
 from onyx.db.engine import get_async_session
+from onyx.db.engine import get_async_session_context_manager
 from onyx.db.engine import get_session
 from onyx.db.models import User
 from onyx.utils.logger import setup_logger
@@ -39,13 +40,10 @@
 
 async def upsert_saml_user(email: str) -> User:
     logger.debug(f"Attempting to upsert SAML user with email: {email}")
-    get_async_session_context = contextlib.asynccontextmanager(
-        get_async_session
-    )  # type:ignore
     get_user_db_context = contextlib.asynccontextmanager(get_user_db)
     get_user_manager_context = contextlib.asynccontextmanager(get_user_manager)
 
-    async with get_async_session_context() as session:
+    async with get_async_session_context_manager() as session:
         async with get_user_db_context(session) as user_db:
             async with get_user_manager_context(user_db) as user_manager:
                 try:

backend/onyx/auth/users.py

@@ -92,7 +92,7 @@
 from onyx.db.auth import get_user_db
 from onyx.db.auth import SQLAlchemyUserAdminDB
 from onyx.db.engine import get_async_session
-from onyx.db.engine import get_async_session_with_tenant
+from onyx.db.engine import get_async_session_context_manager
 from onyx.db.engine import get_session_with_tenant
 from onyx.db.models import AccessToken
 from onyx.db.models import OAuthAccount
@@ -252,7 +252,7 @@ async def get_by_email(self, user_email: str) -> User:
         tenant_id = fetch_ee_implementation_or_noop(
             "onyx.server.tenants.user_mapping", "get_tenant_id_for_email", None
         )(user_email)
-        async with get_async_session_with_tenant(tenant_id) as db_session:
+        async with get_async_session_context_manager(tenant_id) as db_session:
             if MULTI_TENANT:
                 tenant_user_db = SQLAlchemyUserAdminDB[User, uuid.UUID](
                     db_session, User, OAuthAccount
@@ -295,7 +295,7 @@ async def create(
         )
         user: User
 
-        async with get_async_session_with_tenant(tenant_id) as db_session:
+        async with get_async_session_context_manager(tenant_id) as db_session:
             token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)
             verify_email_is_invited(user_create.email)
             verify_email_domain(user_create.email)
@@ -395,7 +395,7 @@ async def oauth_callback(
 
         # Proceed with the tenant context
         token = None
-        async with get_async_session_with_tenant(tenant_id) as db_session:
+        async with get_async_session_context_manager(tenant_id) as db_session:
             token = CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)
 
             verify_email_in_whitelist(account_email, tenant_id)
@@ -634,7 +634,7 @@ async def authenticate(
             return None
 
         # Create a tenant-specific session
-        async with get_async_session_with_tenant(tenant_id) as tenant_session:
+        async with get_async_session_context_manager(tenant_id) as tenant_session:
             tenant_user_db: SQLAlchemyUserDatabase = SQLAlchemyUserDatabase(
                 tenant_session, User
             )

backend/onyx/db/auth.py

@@ -17,7 +17,7 @@
 from onyx.auth.schemas import UserRole
 from onyx.db.api_key import get_api_key_email_pattern
 from onyx.db.engine import get_async_session
-from onyx.db.engine import get_async_session_with_tenant
+from onyx.db.engine import get_async_session_context_manager
 from onyx.db.models import AccessToken
 from onyx.db.models import OAuthAccount
 from onyx.db.models import User
@@ -55,8 +55,9 @@ def get_total_users_count(db_session: Session) -> int:
 
 
 async def get_user_count(only_admin_users: bool = False) -> int:
-    async with get_async_session_with_tenant() as session:
-        stmt = select(func.count(User.id))
+    async with get_async_session_context_manager() as session:
+        count_stmt = func.count(User.id)  # type: ignore
+        stmt = select(count_stmt)
         if only_admin_users:
             stmt = stmt.where(User.role == UserRole.ADMIN)
         result = await session.execute(stmt)

backend/onyx/db/engine.py

@@ -10,7 +10,7 @@
 from contextlib import contextmanager
 from datetime import datetime
 from typing import Any
-from typing import ContextManager
+from typing import AsyncContextManager
 
 import asyncpg  # type: ignore
 import boto3
@@ -46,6 +46,7 @@
 from onyx.utils.logger import setup_logger
 from shared_configs.configs import MULTI_TENANT
 from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
+from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE
 from shared_configs.configs import TENANT_ID_PREFIX
 from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR
 from shared_configs.contextvars import get_current_tenant_id
@@ -352,18 +353,6 @@ def provide_iam_token_async(
     return _ASYNC_ENGINE
 
 
-# Listen for events on the synchronous Session class
-@event.listens_for(Session, "after_begin")
-def _set_search_path(
-    session: Session, transaction: Any, connection: Any, *args: Any, **kwargs: Any
-) -> None:
-    """Every time a new transaction is started,
-    set the search_path from the session's info."""
-    tenant_id = session.info.get("tenant_id")
-    if tenant_id:
-        connection.exec_driver_sql(f'SET search_path = "{tenant_id}"')
-
-
 engine = get_sqlalchemy_async_engine()
 AsyncSessionLocal = sessionmaker(  # type: ignore
     bind=engine,
@@ -372,33 +361,6 @@ def _set_search_path(
 )
 
 
-@asynccontextmanager
-async def get_async_session_with_tenant(
-    tenant_id: str | None = None,
-) -> AsyncGenerator[AsyncSession, None]:
-    if tenant_id is None:
-        tenant_id = get_current_tenant_id()
-
-    if not is_valid_schema_name(tenant_id):
-        logger.error(f"Invalid tenant ID: {tenant_id}")
-        raise ValueError("Invalid tenant ID")
-
-    async with AsyncSessionLocal() as session:
-        session.sync_session.info["tenant_id"] = tenant_id
-
-        if POSTGRES_IDLE_SESSIONS_TIMEOUT:
-            await session.execute(
-                text(
-                    f"SET idle_in_transaction_session_timeout = {POSTGRES_IDLE_SESSIONS_TIMEOUT}"
-                )
-            )
-
-        try:
-            yield session
-        finally:
-            pass
-
-
 @contextmanager
 def get_session_with_current_tenant() -> Generator[Session, None, None]:
     tenant_id = get_current_tenant_id()
@@ -416,17 +378,24 @@ def get_session_with_shared_schema() -> Generator[Session, None, None]:
     CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
 
 
+def _set_search_path_on_checkout__listener(
+    dbapi_conn: Any, connection_record: Any, connection_proxy: Any
+) -> None:
+    """Listener to make sure we ALWAYS set the search path on checkout."""
+    tenant_id = get_current_tenant_id()
+    if tenant_id and is_valid_schema_name(tenant_id):
+        with dbapi_conn.cursor() as cursor:
+            cursor.execute(f'SET search_path TO "{tenant_id}"')
+
+
 @contextmanager
 def get_session_with_tenant(*, tenant_id: str) -> Generator[Session, None, None]:
     """
     Generate a database session for a specific tenant.
     """
-    if tenant_id is None:
-        tenant_id = POSTGRES_DEFAULT_SCHEMA
-
     engine = get_sqlalchemy_engine()
 
-    event.listen(engine, "checkout", set_search_path_on_checkout)
+    event.listen(engine, "checkout", _set_search_path_on_checkout__listener)
 
     if not is_valid_schema_name(tenant_id):
         raise HTTPException(status_code=400, detail="Invalid tenant ID")
@@ -457,57 +426,84 @@ def get_session_with_tenant(*, tenant_id: str) -> Generator[Session, None, None]
                         cursor.close()
 
 
-def set_search_path_on_checkout(
-    dbapi_conn: Any, connection_record: Any, connection_proxy: Any
-) -> None:
+def get_session() -> Generator[Session, None, None]:
+    """For use w/ Depends for FastAPI endpoints.
+
+    Has some additional validation, and likely should be merged
+    with get_session_context_manager in the future."""
     tenant_id = get_current_tenant_id()
-    if tenant_id and is_valid_schema_name(tenant_id):
-        with dbapi_conn.cursor() as cursor:
-            cursor.execute(f'SET search_path TO "{tenant_id}"')
+    if tenant_id == POSTGRES_DEFAULT_SCHEMA and MULTI_TENANT:
+        raise BasicAuthenticationError(detail="User must authenticate")
+
+    if not is_valid_schema_name(tenant_id):
+        raise HTTPException(status_code=400, detail="Invalid tenant ID")
 
+    with get_session_context_manager() as db_session:
+        yield db_session
 
-def get_session_generator_with_tenant() -> Generator[Session, None, None]:
+
+@contextlib.contextmanager
+def get_session_context_manager() -> Generator[Session, None, None]:
+    """Context manager for database sessions."""
     tenant_id = get_current_tenant_id()
     with get_session_with_tenant(tenant_id=tenant_id) as session:
         yield session
 
 
-def get_session() -> Generator[Session, None, None]:
-    tenant_id = get_current_tenant_id()
-    if tenant_id == POSTGRES_DEFAULT_SCHEMA and MULTI_TENANT:
-        raise BasicAuthenticationError(detail="User must authenticate")
+def _set_search_path_on_transaction__listener(
+    session: Session, transaction: Any, connection: Any, *args: Any, **kwargs: Any
+) -> None:
+    """Every time a new transaction is started,
+    set the search_path from the session's info."""
+    tenant_id = session.info.get("tenant_id")
+    if tenant_id:
+        connection.exec_driver_sql(f'SET search_path = "{tenant_id}"')
 
-    engine = get_sqlalchemy_engine()
 
-    with Session(engine, expire_on_commit=False) as session:
-        if MULTI_TENANT:
-            if not is_valid_schema_name(tenant_id):
-                raise HTTPException(status_code=400, detail="Invalid tenant ID")
-            session.execute(text(f'SET search_path = "{tenant_id}"'))
-        yield session
+async def get_async_session(
+    tenant_id: str | None = None,
+) -> AsyncGenerator[AsyncSession, None]:
+    """For use w/ Depends for *async* FastAPI endpoints.
 
+    For standard `async with ... as ...` use, use get_async_session_context_manager.
+    """
+
+    if tenant_id is None:
+        tenant_id = get_current_tenant_id()
 
-async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
-    tenant_id = get_current_tenant_id()
     engine = get_sqlalchemy_async_engine()
+
     async with AsyncSession(engine, expire_on_commit=False) as async_session:
-        if MULTI_TENANT:
-            if not is_valid_schema_name(tenant_id):
-                raise HTTPException(status_code=400, detail="Invalid tenant ID")
-            await async_session.execute(text(f'SET search_path = "{tenant_id}"'))
-        yield async_session
+        # set the search path on sync session as well to be extra safe
+        event.listen(
+            async_session.sync_session,
+            "after_begin",
+            _set_search_path_on_transaction__listener,
+        )
 
+        if POSTGRES_IDLE_SESSIONS_TIMEOUT:
+            await async_session.execute(
+                text(
+                    f"SET idle_in_transaction_session_timeout = {POSTGRES_IDLE_SESSIONS_TIMEOUT}"
+                )
+            )
 
-def get_session_context_manager() -> ContextManager[Session]:
-    """Context manager for database sessions."""
-    return contextlib.contextmanager(get_session_generator_with_tenant)()
+        if not is_valid_schema_name(tenant_id):
+            raise HTTPException(status_code=400, detail="Invalid tenant ID")
 
+        # don't need to set the search path for self-hosted + default schema
+        # this is also true for sync sessions, but just not adding it there for
+        # now to simplify / not change too much
+        if MULTI_TENANT or tenant_id != POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE:
+            await async_session.execute(text(f'SET search_path = "{tenant_id}"'))
+
+        yield async_session
 
-def get_session_factory() -> sessionmaker[Session]:
-    global SessionFactory
-    if SessionFactory is None:
-        SessionFactory = sessionmaker(bind=get_sqlalchemy_engine())
-    return SessionFactory
+
+def get_async_session_context_manager(
+    tenant_id: str | None = None,
+) -> AsyncContextManager[AsyncSession]:
+    return asynccontextmanager(get_async_session)(tenant_id)
 
 
 async def warm_up_connections(