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
@@ -49,13 +50,10 @@ async def upsert_saml_user(email: str) -> User:
     Identity Provider, but we need a valid password to satisfy system requirements.
     """
     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
@@ -253,7 +253,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
@@ -296,7 +296,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)
@@ -402,7 +402,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)
@@ -642,7 +642,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,7 +55,7 @@ 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:
+    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:

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
@@ -189,45 +190,6 @@ class SqlEngine:
     _lock: threading.Lock = threading.Lock()
     _app_name: str = POSTGRES_UNKNOWN_APP_NAME
 
-    # NOTE(rkuo) - this appears to be unused, clean it up?
-    # @classmethod
-    # def _init_engine(cls, **engine_kwargs: Any) -> Engine:
-    #     connection_string = build_connection_string(
-    #         db_api=SYNC_DB_API, app_name=cls._app_name + "_sync", use_iam=USE_IAM_AUTH
-    #     )
-
-    #     # Start with base kwargs that are valid for all pool types
-    #     final_engine_kwargs: dict[str, Any] = {}
-
-    #     if POSTGRES_USE_NULL_POOL:
-    #         # if null pool is specified, then we need to make sure that
-    #         # we remove any passed in kwargs related to pool size that would
-    #         # cause the initialization to fail
-    #         final_engine_kwargs.update(engine_kwargs)
-
-    #         final_engine_kwargs["poolclass"] = pool.NullPool
-    #         if "pool_size" in final_engine_kwargs:
-    #             del final_engine_kwargs["pool_size"]
-    #         if "max_overflow" in final_engine_kwargs:
-    #             del final_engine_kwargs["max_overflow"]
-    #     else:
-    #         final_engine_kwargs["pool_size"] = 20
-    #         final_engine_kwargs["max_overflow"] = 5
-    #         final_engine_kwargs["pool_pre_ping"] = POSTGRES_POOL_PRE_PING
-    #         final_engine_kwargs["pool_recycle"] = POSTGRES_POOL_RECYCLE
-
-    #         # any passed in kwargs override the defaults
-    #         final_engine_kwargs.update(engine_kwargs)
-
-    #     logger.info(f"Creating engine with kwargs: {final_engine_kwargs}")
-    #     # echo=True here for inspecting all emitted db queries
-    #     engine = create_engine(connection_string, **final_engine_kwargs)
-
-    #     if USE_IAM_AUTH:
-    #         event.listen(engine, "do_connect", provide_iam_token)
-
-    #     return engine
-
     @classmethod
     def init_engine(
         cls,
@@ -410,18 +372,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,
@@ -430,33 +380,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()
@@ -474,17 +397,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")
@@ -519,57 +449,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(

backend/shared_configs/configs.py

@@ -142,7 +142,10 @@ def validate_cors_origin(origin: str) -> None:
 # Multi-tenancy configuration
 MULTI_TENANT = os.environ.get("MULTI_TENANT", "").lower() == "true"
 
-POSTGRES_DEFAULT_SCHEMA = os.environ.get("POSTGRES_DEFAULT_SCHEMA") or "public"
+POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE = "public"
+POSTGRES_DEFAULT_SCHEMA = (
+    os.environ.get("POSTGRES_DEFAULT_SCHEMA") or POSTGRES_DEFAULT_SCHEMA_STANDARD_VALUE
+)
 DEFAULT_REDIS_PREFIX = os.environ.get("DEFAULT_REDIS_PREFIX") or "default"