k

backend/alembic/env.py

@@ -1,48 +1,38 @@
-from typing import Any, Literal
-from onyx.db.engine import get_iam_auth_token
-from onyx.configs.app_configs import USE_IAM_AUTH
-from onyx.configs.app_configs import POSTGRES_HOST
-from onyx.configs.app_configs import POSTGRES_PORT
-from onyx.configs.app_configs import POSTGRES_USER
-from onyx.configs.app_configs import AWS_REGION
-from onyx.db.engine import build_connection_string
-from onyx.db.engine import get_all_tenant_ids
-from sqlalchemy import event
-from sqlalchemy import pool
-from sqlalchemy import text
 from sqlalchemy.engine.base import Connection
-import os
-import ssl
+from typing import Literal
 import asyncio
-import logging
 from logging.config import fileConfig
+import logging
 
 from alembic import context
+from sqlalchemy import pool
 from sqlalchemy.ext.asyncio import create_async_engine
+from sqlalchemy.sql import text
 from sqlalchemy.sql.schema import SchemaItem
-from onyx.configs.constants import SSL_CERT_FILE
-from shared_configs.configs import MULTI_TENANT, POSTGRES_DEFAULT_SCHEMA
+
+from shared_configs.configs import MULTI_TENANT
+from onyx.db.engine import build_connection_string
 from onyx.db.models import Base
 from celery.backends.database.session import ResultModelBase  # type: ignore
+from onyx.db.engine import get_all_tenant_ids
+from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
 
 # Alembic Config object
 config = context.config
 
+# Interpret the config file for Python logging.
 if config.config_file_name is not None and config.attributes.get(
     "configure_logger", True
 ):
     fileConfig(config.config_file_name)
 
+# Add your model's MetaData object here for 'autogenerate' support
 target_metadata = [Base.metadata, ResultModelBase.metadata]
 
 EXCLUDE_TABLES = {"kombu_queue", "kombu_message"}
-logger = logging.getLogger(__name__)
 
-ssl_context: ssl.SSLContext | None = None
-if USE_IAM_AUTH:
-    if not os.path.exists(SSL_CERT_FILE):
-        raise FileNotFoundError(f"Expected {SSL_CERT_FILE} when USE_IAM_AUTH is true.")
-    ssl_context = ssl.create_default_context(cafile=SSL_CERT_FILE)
+# Set up logging
+logger = logging.getLogger(__name__)
 
 
 def include_object(
@@ -59,12 +49,20 @@ def include_object(
     reflected: bool,
     compare_to: SchemaItem | None,
 ) -> bool:
+    """
+    Determines whether a database object should be included in migrations.
+    Excludes specified tables from migrations.
+    """
     if type_ == "table" and name in EXCLUDE_TABLES:
         return False
     return True
 
 
 def get_schema_options() -> tuple[str, bool, bool]:
+    """
+    Parses command-line options passed via '-x' in Alembic commands.
+    Recognizes 'schema', 'create_schema', and 'upgrade_all_tenants' options.
+    """
     x_args_raw = context.get_x_argument()
     x_args = {}
     for arg in x_args_raw:
@@ -92,12 +90,16 @@ def get_schema_options() -> tuple[str, bool, bool]:
 def do_run_migrations(
     connection: Connection, schema_name: str, create_schema: bool
 ) -> None:
+    """
+    Executes migrations in the specified schema.
+    """
     logger.info(f"About to migrate schema: {schema_name}")
 
     if create_schema:
         connection.execute(text(f'CREATE SCHEMA IF NOT EXISTS "{schema_name}"'))
         connection.execute(text("COMMIT"))
 
+    # Set search_path to the target schema
     connection.execute(text(f'SET search_path TO "{schema_name}"'))
 
     context.configure(
@@ -115,25 +117,11 @@ def do_run_migrations(
         context.run_migrations()
 
 
-def provide_iam_token_for_alembic(
-    dialect: Any, conn_rec: Any, cargs: Any, cparams: Any
-) -> None:
-    if USE_IAM_AUTH:
-        # Database connection settings
-        region = AWS_REGION
-        host = POSTGRES_HOST
-        port = POSTGRES_PORT
-        user = POSTGRES_USER
-
-        # Get IAM authentication token
-        token = get_iam_auth_token(host, port, user, region)
-
-        # For Alembic / SQLAlchemy in this context, set SSL and password
-        cparams["password"] = token
-        cparams["ssl"] = ssl_context
-
-
 async def run_async_migrations() -> None:
+    """
+    Determines whether to run migrations for a single schema or all schemas,
+    and executes migrations accordingly.
+    """
     schema_name, create_schema, upgrade_all_tenants = get_schema_options()
 
     engine = create_async_engine(
@@ -141,16 +129,10 @@ async def run_async_migrations() -> None:
         poolclass=pool.NullPool,
     )
 
-    if USE_IAM_AUTH:
-
-        @event.listens_for(engine.sync_engine, "do_connect")
-        def event_provide_iam_token_for_alembic(
-            dialect: Any, conn_rec: Any, cargs: Any, cparams: Any
-        ) -> None:
-            provide_iam_token_for_alembic(dialect, conn_rec, cargs, cparams)
-
     if upgrade_all_tenants:
+        # Run migrations for all tenant schemas sequentially
         tenant_schemas = get_all_tenant_ids()
+
         for schema in tenant_schemas:
             try:
                 logger.info(f"Migrating schema: {schema}")
@@ -180,20 +162,15 @@ async def run_async_migrations() -> None:
 
 
 def run_migrations_offline() -> None:
+    """
+    Run migrations in 'offline' mode.
+    """
     schema_name, _, upgrade_all_tenants = get_schema_options()
     url = build_connection_string()
 
     if upgrade_all_tenants:
+        # Run offline migrations for all tenant schemas
         engine = create_async_engine(url)
-
-        if USE_IAM_AUTH:
-
-            @event.listens_for(engine.sync_engine, "do_connect")
-            def event_provide_iam_token_for_alembic_offline(
-                dialect: Any, conn_rec: Any, cargs: Any, cparams: Any
-            ) -> None:
-                provide_iam_token_for_alembic(dialect, conn_rec, cargs, cparams)
-
         tenant_schemas = get_all_tenant_ids()
         engine.sync_engine.dispose()
 
@@ -230,6 +207,9 @@ def run_migrations_offline() -> None:
 
 
 def run_migrations_online() -> None:
+    """
+    Runs migrations in 'online' mode using an asynchronous engine.
+    """
     asyncio.run(run_async_migrations())
 
 

backend/onyx/configs/app_configs.py

@@ -144,7 +144,6 @@ POSTGRES_PASSWORD = urllib.parse.quote_plus(
 POSTGRES_HOST = os.environ.get("POSTGRES_HOST") or "localhost"
 POSTGRES_PORT = os.environ.get("POSTGRES_PORT") or "5432"
 POSTGRES_DB = os.environ.get("POSTGRES_DB") or "postgres"
-AWS_REGION = os.environ.get("AWS_REGION") or "us-east-2"
 
 POSTGRES_API_SERVER_POOL_SIZE = int(
     os.environ.get("POSTGRES_API_SERVER_POOL_SIZE") or 40
@@ -175,9 +174,6 @@ try:
 except ValueError:
     POSTGRES_IDLE_SESSIONS_TIMEOUT = POSTGRES_IDLE_SESSIONS_TIMEOUT_DEFAULT
 
-USE_IAM_AUTH = os.getenv("USE_IAM_AUTH", "False").lower() == "true"
-
-
 REDIS_SSL = os.getenv("REDIS_SSL", "").lower() == "true"
 REDIS_HOST = os.environ.get("REDIS_HOST") or "localhost"
 REDIS_PORT = int(os.environ.get("REDIS_PORT", 6379))

backend/onyx/configs/constants.py

@@ -49,7 +49,6 @@ POSTGRES_CELERY_WORKER_INDEXING_CHILD_APP_NAME = "celery_worker_indexing_child"
 POSTGRES_PERMISSIONS_APP_NAME = "permissions"
 POSTGRES_UNKNOWN_APP_NAME = "unknown"
 
-SSL_CERT_FILE = "bundle.pem"
 # API Keys
 DANSWER_API_KEY_PREFIX = "API_KEY__"
 DANSWER_API_KEY_DUMMY_EMAIL_DOMAIN = "onyxapikey.ai"

backend/onyx/db/engine.py

@@ -1,7 +1,5 @@
 import contextlib
-import os
 import re
-import ssl
 import threading
 import time
 from collections.abc import AsyncGenerator
@@ -12,8 +10,6 @@ from datetime import datetime
 from typing import Any
 from typing import ContextManager
 
-import asyncpg  # type: ignore
-import boto3
 import jwt
 from fastapi import HTTPException
 from fastapi import Request
@@ -27,7 +23,6 @@ from sqlalchemy.ext.asyncio import create_async_engine
 from sqlalchemy.orm import Session
 from sqlalchemy.orm import sessionmaker
 
-from onyx.configs.app_configs import AWS_REGION
 from onyx.configs.app_configs import LOG_POSTGRES_CONN_COUNTS
 from onyx.configs.app_configs import LOG_POSTGRES_LATENCY
 from onyx.configs.app_configs import POSTGRES_API_SERVER_POOL_OVERFLOW
@@ -42,7 +37,6 @@ from onyx.configs.app_configs import POSTGRES_PORT
 from onyx.configs.app_configs import POSTGRES_USER
 from onyx.configs.app_configs import USER_AUTH_SECRET
 from onyx.configs.constants import POSTGRES_UNKNOWN_APP_NAME
-from onyx.configs.constants import SSL_CERT_FILE
 from onyx.server.utils import BasicAuthenticationError
 from onyx.utils.logger import setup_logger
 from shared_configs.configs import MULTI_TENANT
@@ -55,87 +49,28 @@ logger = setup_logger()
 SYNC_DB_API = "psycopg2"
 ASYNC_DB_API = "asyncpg"
 
-USE_IAM_AUTH = os.getenv("USE_IAM_AUTH", "False").lower() == "true"
+# global so we don't create more than one engine per process
+# outside of being best practice, this is needed so we can properly pool
+# connections and not create a new pool on every request
 
-# Global so we don't create more than one engine per process
 _ASYNC_ENGINE: AsyncEngine | None = None
 SessionFactory: sessionmaker[Session] | None = None
 
-
-def create_ssl_context_if_iam() -> ssl.SSLContext | None:
-    """Create an SSL context if IAM authentication is enabled, else return None."""
-    if USE_IAM_AUTH:
-        return ssl.create_default_context(cafile=SSL_CERT_FILE)
-    return None
-
-
-ssl_context = create_ssl_context_if_iam()
-
-
-def get_iam_auth_token(
-    host: str, port: str, user: str, region: str = "us-east-2"
-) -> str:
-    """
-    Generate an IAM authentication token using boto3.
-    """
-    client = boto3.client("rds", region_name=region)
-    token = client.generate_db_auth_token(
-        DBHostname=host, Port=int(port), DBUsername=user
-    )
-    return token
-
-
-def configure_psycopg2_iam_auth(
-    cparams: dict[str, Any], host: str, port: str, user: str, region: str
-) -> None:
-    """
-    Configure cparams for psycopg2 with IAM token and SSL.
-    """
-    token = get_iam_auth_token(host, port, user, region)
-    cparams["password"] = token
-    cparams["sslmode"] = "require"
-    cparams["sslrootcert"] = SSL_CERT_FILE
-
-
-def build_connection_string(
-    *,
-    db_api: str = ASYNC_DB_API,
-    user: str = POSTGRES_USER,
-    password: str = POSTGRES_PASSWORD,
-    host: str = POSTGRES_HOST,
-    port: str = POSTGRES_PORT,
-    db: str = POSTGRES_DB,
-    app_name: str | None = None,
-    use_iam: bool = USE_IAM_AUTH,
-    region: str = "us-west-2",
-) -> str:
-    if use_iam:
-        base_conn_str = f"postgresql+{db_api}://{user}@{host}:{port}/{db}"
-    else:
-        base_conn_str = f"postgresql+{db_api}://{user}:{password}@{host}:{port}/{db}"
-
-    # For asyncpg, do not include application_name in the connection string
-    if app_name and db_api != "asyncpg":
-        if "?" in base_conn_str:
-            return f"{base_conn_str}&application_name={app_name}"
-        else:
-            return f"{base_conn_str}?application_name={app_name}"
-    return base_conn_str
-
-
 if LOG_POSTGRES_LATENCY:
-
+    # Function to log before query execution
     @event.listens_for(Engine, "before_cursor_execute")
     def before_cursor_execute(  # type: ignore
         conn, cursor, statement, parameters, context, executemany
     ):
         conn.info["query_start_time"] = time.time()
 
+    # Function to log after query execution
     @event.listens_for(Engine, "after_cursor_execute")
     def after_cursor_execute(  # type: ignore
         conn, cursor, statement, parameters, context, executemany
     ):
         total_time = time.time() - conn.info["query_start_time"]
+        # don't spam TOO hard
         if total_time > 0.1:
             logger.debug(
                 f"Query Complete: {statement}\n\nTotal Time: {total_time:.4f} seconds"
@@ -143,6 +78,7 @@ if LOG_POSTGRES_LATENCY:
 
 
 if LOG_POSTGRES_CONN_COUNTS:
+    # Global counter for connection checkouts and checkins
     checkout_count = 0
     checkin_count = 0
 
@@ -169,13 +105,21 @@ if LOG_POSTGRES_CONN_COUNTS:
         logger.debug(f"Total connection checkins: {checkin_count}")
 
 
+"""END DEBUGGING LOGGING"""
+
+
 def get_db_current_time(db_session: Session) -> datetime:
+    """Get the current time from Postgres representing the start of the transaction
+    Within the same transaction this value will not update
+    This datetime object returned should be timezone aware, default Postgres timezone is UTC
+    """
     result = db_session.execute(text("SELECT NOW()")).scalar()
     if result is None:
         raise ValueError("Database did not return a time")
     return result
 
 
+# Regular expression to validate schema names to prevent SQL injection
 SCHEMA_NAME_REGEX = re.compile(r"^[a-zA-Z0-9_-]+$")
 
 
@@ -184,9 +128,16 @@ def is_valid_schema_name(name: str) -> bool:
 
 
 class SqlEngine:
+    """Class to manage a global SQLAlchemy engine (needed for proper resource control).
+    Will eventually subsume most of the standalone functions in this file.
+    Sync only for now.
+    """
+
     _engine: Engine | None = None
     _lock: threading.Lock = threading.Lock()
     _app_name: str = POSTGRES_UNKNOWN_APP_NAME
+
+    # Default parameters for engine creation
     DEFAULT_ENGINE_KWARGS = {
         "pool_size": 20,
         "max_overflow": 5,
@@ -194,27 +145,33 @@ class SqlEngine:
         "pool_recycle": POSTGRES_POOL_RECYCLE,
     }
 
+    def __init__(self) -> None:
+        pass
+
     @classmethod
     def _init_engine(cls, **engine_kwargs: Any) -> Engine:
+        """Private helper method to create and return an Engine."""
         connection_string = build_connection_string(
-            db_api=SYNC_DB_API, app_name=cls._app_name + "_sync", use_iam=USE_IAM_AUTH
+            db_api=SYNC_DB_API, app_name=cls._app_name + "_sync"
         )
         merged_kwargs = {**cls.DEFAULT_ENGINE_KWARGS, **engine_kwargs}
-        engine = create_engine(connection_string, **merged_kwargs)
-
-        if USE_IAM_AUTH:
-            event.listen(engine, "do_connect", provide_iam_token)
-
-        return engine
+        return create_engine(connection_string, **merged_kwargs)
 
     @classmethod
     def init_engine(cls, **engine_kwargs: Any) -> None:
+        """Allow the caller to init the engine with extra params. Different clients
+        such as the API server and different Celery workers and tasks
+        need different settings.
+        """
         with cls._lock:
             if not cls._engine:
                 cls._engine = cls._init_engine(**engine_kwargs)
 
     @classmethod
     def get_engine(cls) -> Engine:
+        """Gets the SQLAlchemy engine. Will init a default engine if init hasn't
+        already been called. You probably want to init first!
+        """
         if not cls._engine:
             with cls._lock:
                 if not cls._engine:
@@ -223,10 +180,12 @@ class SqlEngine:
 
     @classmethod
     def set_app_name(cls, app_name: str) -> None:
+        """Class method to set the app name."""
         cls._app_name = app_name
 
     @classmethod
     def get_app_name(cls) -> str:
+        """Class method to get current app name."""
         if not cls._app_name:
             return ""
         return cls._app_name
@@ -258,71 +217,56 @@ def get_all_tenant_ids() -> list[str] | list[None]:
         for tenant in tenant_ids
         if tenant is None or tenant.startswith(TENANT_ID_PREFIX)
     ]
+
     return valid_tenants
 
 
+def build_connection_string(
+    *,
+    db_api: str = ASYNC_DB_API,
+    user: str = POSTGRES_USER,
+    password: str = POSTGRES_PASSWORD,
+    host: str = POSTGRES_HOST,
+    port: str = POSTGRES_PORT,
+    db: str = POSTGRES_DB,
+    app_name: str | None = None,
+) -> str:
+    if app_name:
+        return f"postgresql+{db_api}://{user}:{password}@{host}:{port}/{db}?application_name={app_name}"
+    return f"postgresql+{db_api}://{user}:{password}@{host}:{port}/{db}"
+
+
 def get_sqlalchemy_engine() -> Engine:
     return SqlEngine.get_engine()
 
 
-async def get_async_connection() -> Any:
-    """
-    Custom connection function for async engine when using IAM auth.
-    """
-    host = POSTGRES_HOST
-    port = POSTGRES_PORT
-    user = POSTGRES_USER
-    db = POSTGRES_DB
-    token = get_iam_auth_token(host, port, user, AWS_REGION)
-
-    # asyncpg requires 'ssl="require"' if SSL needed
-    return await asyncpg.connect(
-        user=user, password=token, host=host, port=int(port), database=db, ssl="require"
-    )
-
-
 def get_sqlalchemy_async_engine() -> AsyncEngine:
     global _ASYNC_ENGINE
     if _ASYNC_ENGINE is None:
-        app_name = SqlEngine.get_app_name() + "_async"
-        connection_string = build_connection_string(
-            db_api=ASYNC_DB_API,
-            use_iam=USE_IAM_AUTH,
-        )
-
-        connect_args: dict[str, Any] = {}
-        if app_name:
-            connect_args["server_settings"] = {"application_name": app_name}
-
-        connect_args["ssl"] = ssl_context
-
+        # Underlying asyncpg cannot accept application_name directly in the connection string
+        # https://github.com/MagicStack/asyncpg/issues/798
+        connection_string = build_connection_string()
         _ASYNC_ENGINE = create_async_engine(
             connection_string,
-            connect_args=connect_args,
+            connect_args={
+                "server_settings": {
+                    "application_name": SqlEngine.get_app_name() + "_async"
+                }
+            },
+            # async engine is only used by API server, so we can use those values
+            # here as well
             pool_size=POSTGRES_API_SERVER_POOL_SIZE,
             max_overflow=POSTGRES_API_SERVER_POOL_OVERFLOW,
             pool_pre_ping=POSTGRES_POOL_PRE_PING,
             pool_recycle=POSTGRES_POOL_RECYCLE,
         )
-
-        if USE_IAM_AUTH:
-
-            @event.listens_for(_ASYNC_ENGINE.sync_engine, "do_connect")
-            def provide_iam_token_async(
-                dialect: Any, conn_rec: Any, cargs: Any, cparams: Any
-            ) -> None:
-                # For async engine using asyncpg, we still need to set the IAM token here.
-                host = POSTGRES_HOST
-                port = POSTGRES_PORT
-                user = POSTGRES_USER
-                token = get_iam_auth_token(host, port, user, AWS_REGION)
-                cparams["password"] = token
-                cparams["ssl"] = ssl_context
-
     return _ASYNC_ENGINE
 
 
+# Dependency to get the current tenant ID
+# If no token is present, uses the default schema for this use case
 def get_current_tenant_id(request: Request) -> str:
+    """Dependency that extracts the tenant ID from the JWT token in the request and sets the context variable."""
     if not MULTI_TENANT:
         tenant_id = POSTGRES_DEFAULT_SCHEMA
         CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)
@@ -331,6 +275,7 @@ def get_current_tenant_id(request: Request) -> str:
     token = request.cookies.get("fastapiusersauth")
     if not token:
         current_value = CURRENT_TENANT_ID_CONTEXTVAR.get()
+        # If no token is present, use the default schema or handle accordingly
         return current_value
 
     try:
@@ -344,6 +289,7 @@ def get_current_tenant_id(request: Request) -> str:
         if not is_valid_schema_name(tenant_id):
             raise HTTPException(status_code=400, detail="Invalid tenant ID format")
         CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)
+
         return tenant_id
     except jwt.InvalidTokenError:
         return CURRENT_TENANT_ID_CONTEXTVAR.get()
@@ -370,6 +316,7 @@ async def get_async_session_with_tenant(
 
     async with async_session_factory() as session:
         try:
+            # Set the search_path to the tenant's schema
             await session.execute(text(f'SET search_path = "{tenant_id}"'))
             if POSTGRES_IDLE_SESSIONS_TIMEOUT:
                 await session.execute(
@@ -379,6 +326,8 @@ async def get_async_session_with_tenant(
                 )
         except Exception:
             logger.exception("Error setting search_path.")
+            # You can choose to re-raise the exception or handle it
+            # Here, we'll re-raise to prevent proceeding with an incorrect session
             raise
         else:
             yield session
@@ -386,6 +335,9 @@ async def get_async_session_with_tenant(
 
 @contextmanager
 def get_session_with_default_tenant() -> Generator[Session, None, None]:
+    """
+    Get a database session using the current tenant ID from the context variable.
+    """
     tenant_id = CURRENT_TENANT_ID_CONTEXTVAR.get()
     with get_session_with_tenant(tenant_id) as session:
         yield session
@@ -397,6 +349,7 @@ def get_session_with_tenant(
 ) -> Generator[Session, None, None]:
     """
     Generate a database session for a specific tenant.
+
     This function:
     1. Sets the database schema to the specified tenant's schema.
     2. Preserves the tenant ID across the session.
@@ -404,20 +357,27 @@ def get_session_with_tenant(
     4. Uses the default schema if no tenant ID is provided.
     """
     engine = get_sqlalchemy_engine()
+
+    # Store the previous tenant ID
     previous_tenant_id = CURRENT_TENANT_ID_CONTEXTVAR.get() or POSTGRES_DEFAULT_SCHEMA
 
     if tenant_id is None:
         tenant_id = POSTGRES_DEFAULT_SCHEMA
 
     CURRENT_TENANT_ID_CONTEXTVAR.set(tenant_id)
+
     event.listen(engine, "checkout", set_search_path_on_checkout)
 
     if not is_valid_schema_name(tenant_id):
         raise HTTPException(status_code=400, detail="Invalid tenant ID")
 
     try:
+        # Establish a raw connection
         with engine.connect() as connection:
+            # Access the raw DBAPI connection and set the search_path
             dbapi_connection = connection.connection
+
+            # Set the search_path outside of any transaction
             cursor = dbapi_connection.cursor()
             try:
                 cursor.execute(f'SET search_path = "{tenant_id}"')
@@ -430,17 +390,21 @@ def get_session_with_tenant(
             finally:
                 cursor.close()
 
+            # Bind the session to the connection
             with Session(bind=connection, expire_on_commit=False) as session:
                 try:
                     yield session
                 finally:
+                    # Reset search_path to default after the session is used
                     if MULTI_TENANT:
                         cursor = dbapi_connection.cursor()
                         try:
                             cursor.execute('SET search_path TO "$user", public')
                         finally:
                             cursor.close()
+
     finally:
+        # Restore the previous tenant ID
         CURRENT_TENANT_ID_CONTEXTVAR.set(previous_tenant_id)
 
 
@@ -460,9 +424,12 @@ def get_session_generator_with_tenant() -> Generator[Session, None, None]:
 
 
 def get_session() -> Generator[Session, None, None]:
+    """Generate a database session with the appropriate tenant schema set."""
     tenant_id = CURRENT_TENANT_ID_CONTEXTVAR.get()
     if tenant_id == POSTGRES_DEFAULT_SCHEMA and MULTI_TENANT:
-        raise BasicAuthenticationError(detail="User must authenticate")
+        raise BasicAuthenticationError(
+            detail="User must authenticate",
+        )
 
     engine = get_sqlalchemy_engine()
 
@@ -470,17 +437,20 @@ def get_session() -> Generator[Session, None, None]:
         if MULTI_TENANT:
             if not is_valid_schema_name(tenant_id):
                 raise HTTPException(status_code=400, detail="Invalid tenant ID")
+            # Set the search_path to the tenant's schema
             session.execute(text(f'SET search_path = "{tenant_id}"'))
         yield session
 
 
 async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
+    """Generate an async database session with the appropriate tenant schema set."""
     tenant_id = CURRENT_TENANT_ID_CONTEXTVAR.get()
     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")
+            # Set the search_path to the tenant's schema
             await async_session.execute(text(f'SET search_path = "{tenant_id}"'))
         yield async_session
 
@@ -491,6 +461,7 @@ def get_session_context_manager() -> ContextManager[Session]:
 
 
 def get_session_factory() -> sessionmaker[Session]:
+    """Get a session factory."""
     global SessionFactory
     if SessionFactory is None:
         SessionFactory = sessionmaker(bind=get_sqlalchemy_engine())
@@ -518,13 +489,3 @@ async def warm_up_connections(
         await async_conn.execute(text("SELECT 1"))
     for async_conn in async_connections:
         await async_conn.close()
-
-
-def provide_iam_token(dialect: Any, conn_rec: Any, cargs: Any, cparams: Any) -> None:
-    if USE_IAM_AUTH:
-        host = POSTGRES_HOST
-        port = POSTGRES_PORT
-        user = POSTGRES_USER
-        region = os.getenv("AWS_REGION", "us-east-2")
-        # Configure for psycopg2 with IAM token
-        configure_psycopg2_iam_auth(cparams, host, port, user, region)

deployment/cloud_kubernetes/workers/beat.yaml

@@ -14,7 +14,7 @@ spec:
     spec:
       containers:
         - name: celery-beat
-          image: onyxdotapp/onyx-backend-cloud:v0.14.0-cloud.beta.21
+          image: onyxdotapp/onyx-backend-cloud:v0.14.0-cloud.beta.20
           imagePullPolicy: IfNotPresent
           command:
             [

deployment/cloud_kubernetes/workers/heavy_worker.yaml

@@ -14,7 +14,7 @@ spec:
     spec:
       containers:
         - name: celery-worker-heavy
-          image: onyxdotapp/onyx-backend-cloud:v0.14.0-cloud.beta.21
+          image: onyxdotapp/onyx-backend-cloud:v0.14.0-cloud.beta.20
           imagePullPolicy: IfNotPresent
           command:
             [

deployment/cloud_kubernetes/workers/indexing_worker.yaml

@@ -14,7 +14,7 @@ spec:
     spec:
       containers:
         - name: celery-worker-indexing
-          image: onyxdotapp/onyx-backend-cloud:v0.14.0-cloud.beta.21
+          image: onyxdotapp/onyx-backend-cloud:v0.14.0-cloud.beta.20
           imagePullPolicy: IfNotPresent
           command:
             [

deployment/cloud_kubernetes/workers/light_worker.yaml

@@ -14,7 +14,7 @@ spec:
     spec:
       containers:
         - name: celery-worker-light
-          image: onyxdotapp/onyx-backend-cloud:v0.14.0-cloud.beta.21
+          image: onyxdotapp/onyx-backend-cloud:v0.14.0-cloud.beta.20
           imagePullPolicy: IfNotPresent
           command:
             [

deployment/cloud_kubernetes/workers/primary.yaml

@@ -14,7 +14,7 @@ spec:
     spec:
       containers:
         - name: celery-worker-primary
-          image: onyxdotapp/onyx-backend-cloud:v0.14.0-cloud.beta.21
+          image: onyxdotapp/onyx-backend-cloud:v0.14.0-cloud.beta.20
           imagePullPolicy: IfNotPresent
           command:
             [

deployment/docker_compose/docker-compose.dev.yml

@@ -103,13 +103,6 @@ services:
       - ENABLE_PAID_ENTERPRISE_EDITION_FEATURES=${ENABLE_PAID_ENTERPRISE_EDITION_FEATURES:-false}
       - API_KEY_HASH_ROUNDS=${API_KEY_HASH_ROUNDS:-}
       # Seeding configuration
-      - USE_IAM_AUTH=${USE_IAM_AUTH:-}
-      - AWS_REGION=${AWS_REGION-}
-      - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID-}
-      - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY-}
-    # Uncomment the line below to use if IAM_AUTH is true and you are using iam auth for postgres
-    # volumes:
-    #   - ./bundle.pem:/app/bundle.pem:ro
     extra_hosts:
       - "host.docker.internal:host-gateway"
     logging:
@@ -230,13 +223,6 @@ services:
 
       # Enterprise Edition stuff
       - ENABLE_PAID_ENTERPRISE_EDITION_FEATURES=${ENABLE_PAID_ENTERPRISE_EDITION_FEATURES:-false}
-      - USE_IAM_AUTH=${USE_IAM_AUTH:-}
-      - AWS_REGION=${AWS_REGION-}
-      - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID-}
-      - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY-}
-    # Uncomment the line below to use if IAM_AUTH is true and you are using iam auth for postgres
-    # volumes:
-    #   - ./bundle.pem:/app/bundle.pem:ro
     extra_hosts:
       - "host.docker.internal:host-gateway"
     logging:

deployment/docker_compose/docker-compose.gpu-dev.yml

@@ -91,13 +91,6 @@ services:
       # Enterprise Edition only
       - API_KEY_HASH_ROUNDS=${API_KEY_HASH_ROUNDS:-}
       - ENABLE_PAID_ENTERPRISE_EDITION_FEATURES=${ENABLE_PAID_ENTERPRISE_EDITION_FEATURES:-false}
-      - USE_IAM_AUTH=${USE_IAM_AUTH}
-      - AWS_REGION=${AWS_REGION-}
-      - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID-}
-      - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY-}
-    # Uncomment the line below to use if IAM_AUTH is true and you are using iam auth for postgres
-    # volumes:
-    #   - ./bundle.pem:/app/bundle.pem:ro
     extra_hosts:
       - "host.docker.internal:host-gateway"
     logging:
@@ -199,13 +192,6 @@ services:
       # Enterprise Edition only
       - API_KEY_HASH_ROUNDS=${API_KEY_HASH_ROUNDS:-}
       - ENABLE_PAID_ENTERPRISE_EDITION_FEATURES=${ENABLE_PAID_ENTERPRISE_EDITION_FEATURES:-false}
-      - USE_IAM_AUTH=${USE_IAM_AUTH}
-      - AWS_REGION=${AWS_REGION-}
-      - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID-}
-      - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY-}
-    # Uncomment the line below to use if IAM_AUTH is true and you are using iam auth for postgres
-    # volumes:
-    #   - ./bundle.pem:/app/bundle.pem:ro
     extra_hosts:
       - "host.docker.internal:host-gateway"
     logging:

deployment/docker_compose/docker-compose.prod-no-letsencrypt.yml

@@ -22,13 +22,6 @@ services:
       - VESPA_HOST=index
       - REDIS_HOST=cache
       - MODEL_SERVER_HOST=${MODEL_SERVER_HOST:-inference_model_server}
-      - USE_IAM_AUTH=${USE_IAM_AUTH}
-      - AWS_REGION=${AWS_REGION-}
-      - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID-}
-      - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY-}
-    # Uncomment the line below to use if IAM_AUTH is true and you are using iam auth for postgres
-    # volumes:
-    #   - ./bundle.pem:/app/bundle.pem:ro
     extra_hosts:
       - "host.docker.internal:host-gateway"
     logging:
@@ -59,13 +52,6 @@ services:
       - REDIS_HOST=cache
       - MODEL_SERVER_HOST=${MODEL_SERVER_HOST:-inference_model_server}
       - INDEXING_MODEL_SERVER_HOST=${INDEXING_MODEL_SERVER_HOST:-indexing_model_server}
-      - USE_IAM_AUTH=${USE_IAM_AUTH}
-      - AWS_REGION=${AWS_REGION-}
-      - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID-}
-      - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY-}
-    # Uncomment the line below to use if IAM_AUTH is true and you are using iam auth for postgres
-    # volumes:
-    #   - ./bundle.pem:/app/bundle.pem:ro
     extra_hosts:
       - "host.docker.internal:host-gateway"
     logging:

deployment/docker_compose/docker-compose.prod.yml

@@ -23,13 +23,6 @@ services:
       - VESPA_HOST=index
       - REDIS_HOST=cache
       - MODEL_SERVER_HOST=${MODEL_SERVER_HOST:-inference_model_server}
-      - USE_IAM_AUTH=${USE_IAM_AUTH}
-      - AWS_REGION=${AWS_REGION-}
-      - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID-}
-      - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY-}
-    # Uncomment the line below to use if IAM_AUTH is true and you are using iam auth for postgres
-    # volumes:
-    #   - ./bundle.pem:/app/bundle.pem:ro
     extra_hosts:
       - "host.docker.internal:host-gateway"
     logging:
@@ -64,13 +57,6 @@ services:
       - REDIS_HOST=cache
       - MODEL_SERVER_HOST=${MODEL_SERVER_HOST:-inference_model_server}
       - INDEXING_MODEL_SERVER_HOST=${INDEXING_MODEL_SERVER_HOST:-indexing_model_server}
-      - USE_IAM_AUTH=${USE_IAM_AUTH}
-      - AWS_REGION=${AWS_REGION-}
-      - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID-}
-      - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY-}
-    # Uncomment the line below to use if IAM_AUTH is true and you are using iam auth for postgres
-    # volumes:
-    #   - ./bundle.pem:/app/bundle.pem:ro
     extra_hosts:
       - "host.docker.internal:host-gateway"
     logging:
@@ -237,7 +223,7 @@ services:
     volumes:
       - ../data/certbot/conf:/etc/letsencrypt
       - ../data/certbot/www:/var/www/certbot
-    logging::wq
+    logging:
       driver: json-file
       options:
         max-size: "50m"
@@ -259,6 +245,3 @@ volumes:
   # Created by the container itself
   model_cache_huggingface:
   indexing_huggingface_model_cache:
-
-
-

deployment/kubernetes/api_server-service-deployment.yaml

@@ -60,12 +60,3 @@ spec:
           envFrom:
             - configMapRef:
                 name: env-configmap
-      # Uncomment if you are using IAM auth for Postgres
-      #     volumeMounts:
-      #       - name: bundle-pem
-      #         mountPath: "/app/certs"
-      #         readOnly: true
-      # volumes:
-      #   - name: bundle-pem
-      #     secret:
-      #       secretName: bundle-pem-secret

deployment/kubernetes/background-deployment.yaml

@@ -43,7 +43,6 @@ spec:
       #     - name: my-ca-cert-volume
       #       mountPath: /etc/ssl/certs/custom-ca.crt
       #       subPath: my-ca.crt
-
       # Optional volume for CA certificate
       # volumes:
       #   - name: my-cas-cert-volume
@@ -52,13 +51,3 @@ spec:
       #       items:
       #         - key: my-ca.crt
       #           path: my-ca.crt
-
-      # Uncomment if you are using IAM auth for Postgres
-      #     volumeMounts:
-      #       - name: bundle-pem
-      #         mountPath: "/app/certs"
-      #         readOnly: true
-      # volumes:
-      #   - name: bundle-pem
-      #     secret:
-      #       secretName: bundle-pem-secret

web/playwright.config.ts

@@ -1,38 +1,18 @@
 import { defineConfig, devices } from "@playwright/test";
+import path from "path";
 
 export default defineConfig({
-  workers: 1, // temporary change to see if single threaded testing stabilizes the tests
-  testDir: "./tests/e2e", // Folder for test files
+  workers: 1,
+  testDir: "./tests/e2e",
   reporter: "list",
-  // Configure paths for screenshots
-  // expect: {
-  //   toMatchSnapshot: {
-  //     threshold: 0.2, // Adjust the threshold for visual diffs
-  //   },
-  // },
-  // reporter: [["html", { outputFolder: "test-results/output/report" }]], // HTML report location
-  // outputDir: "test-results/output/screenshots", // Set output folder for test artifacts
+  globalSetup: path.join(__dirname, "tests/e2e/global-setup.ts"),
+  use: {
+    baseURL: "http://localhost:3000",
+    trace: "on-first-retry",
+  },
   projects: [
     {
-      // dependency for admin workflows
-      name: "admin_setup",
-      testMatch: /.*\admin_auth\.setup\.ts/,
-    },
-    {
-      // tests admin workflows
-      name: "chromium-admin",
-      grep: /@admin/,
-      use: {
-        ...devices["Desktop Chrome"],
-        // Use prepared auth state.
-        storageState: "admin_auth.json",
-      },
-      dependencies: ["admin_setup"],
-    },
-    {
-      // tests logged out / guest workflows
-      name: "chromium-guest",
-      grep: /@guest/,
+      name: "chromium",
       use: {
         ...devices["Desktop Chrome"],
       },

web/src/app/admin/embeddings/pages/EmbeddingFormPage.tsx

@@ -216,7 +216,8 @@ export default function EmbeddingForm() {
         <button
           className="enabled:cursor-pointer ml-auto disabled:bg-accent/50 disabled:cursor-not-allowed bg-accent flex mx-auto gap-x-1 items-center text-white py-2.5 px-3.5 text-sm font-regular rounded-sm"
           onClick={async () => {
-            updateSearch();
+            await updateSearch();
+            router.refresh();
             navigateToEmbeddingPage("search settings");
           }}
         >

web/tests/e2e/admin_auth.setup.ts

@@ -18,6 +18,9 @@ setup("authenticate", async ({ page }) => {
   // Click the login button
   await page.click('button[type="submit"]');
 
+  // Take a screenshot after successful login
+  await page.screenshot({ path: "admin_login_success.png" });
+
   await page.waitForURL("http://localhost:3000/chat");
 
   await page.context().storageState({ path: "admin_auth.json" });

web/tests/e2e/auth.setup.ts

@@ -0,0 +1,38 @@
+import { test as setup, expect } from "@playwright/test";
+import { loginAs } from "./utils/auth";
+
+setup("create auth states", async ({ browser }) => {
+  // Setup admin auth
+  const adminContext = await browser.newContext();
+  const adminPage = await adminContext.newPage();
+  await loginAs(adminPage, "admin");
+  await adminContext.storageState({ path: "admin_auth.json" });
+  await adminContext.close();
+
+  // Setup user auth
+  const userContext = await browser.newContext();
+  const userPage = await userContext.newPage();
+
+  // Login as admin to create user
+  await loginAs(userPage, "admin");
+  await userPage.goto("http://localhost:3000/admin/indexing/status");
+  await userPage.getByRole("button", { name: "Users" }).click();
+  await userPage.getByRole("button", { name: "Invite Users" }).click();
+  await userPage.locator("#emails").fill("user_user@test.com");
+  await userPage.getByRole("button", { name: "Add!" }).click();
+
+  // Logout admin
+  await userPage.getByText("A", { exact: true }).click();
+  await userPage.getByText("Log out").click();
+
+  // Create and login as new user
+  await userPage.goto("http://localhost:3000/auth/login");
+  await userPage.getByRole("link", { name: "Create an account" }).click();
+  await userPage.getByTestId("email").fill("user_user@test.com");
+  await userPage.getByTestId("password").fill("test");
+  await userPage.getByRole("button", { name: "Sign Up" }).click();
+  await userPage.waitForURL("http://localhost:3000/chat");
+
+  await userContext.storageState({ path: "user_auth.json" });
+  await userContext.close();
+});

web/tests/e2e/constants.js

@@ -2,3 +2,8 @@ export const TEST_CREDENTIALS = {
   email: "admin_user@test.com",
   password: "TestPassword123!",
 };
+
+export const TEST_USER_CREDENTIALS = {
+  email: "basic_user@test.com",
+  password: "test",
+};

web/tests/e2e/core_flows/search_settings_update.test.ts

@@ -0,0 +1,16 @@
+import { test, expect } from "../utils/auth-utils";
+
+test("update search settings", async ({ auth }) => {
+  const { page } = await auth("admin");
+
+  await page.goto("http://localhost:3000/admin/indexing/status");
+  await page.getByRole("button", { name: "Search Settings" }).click();
+  await page.getByRole("button", { name: "Update Search Settings" }).click();
+  await page.getByRole("button", { name: "Continue" }).click();
+  await page.getByRole("button", { name: "Self-hosted" }).click();
+  await page.getByText("MixedBread XSmallFastest,").click();
+  await page.getByRole("button", { name: "Update Search" }).click();
+  await expect(page.locator("body")).toContainText(
+    "mixedbread-ai/mxbai-rerank-xsmall-v1"
+  );
+});

web/tests/e2e/global-setup.ts

@@ -0,0 +1,44 @@
+import { chromium, FullConfig } from "@playwright/test";
+import { loginAs } from "./utils/auth";
+
+async function globalSetup(config: FullConfig) {
+  const browser = await chromium.launch();
+
+  // Setup admin auth
+  const adminContext = await browser.newContext();
+  const adminPage = await adminContext.newPage();
+  await loginAs(adminPage, "admin");
+  await adminContext.storageState({ path: "admin_auth.json" });
+  await adminContext.close();
+
+  // Setup user auth
+  const userContext = await browser.newContext();
+  const userPage = await userContext.newPage();
+
+  // Login as admin to create user
+  await loginAs(userPage, "admin");
+  await userPage.goto("/admin/indexing/status");
+  await userPage.getByRole("button", { name: "Users" }).click();
+  await userPage.getByRole("button", { name: "Invite Users" }).click();
+  await userPage.locator("#emails").fill("user_user@test.com");
+  await userPage.getByRole("button", { name: "Add!" }).click();
+
+  // Logout admin
+  await userPage.getByText("A", { exact: true }).click();
+  await userPage.getByText("Log out").click();
+
+  // Create and login as new user
+  await userPage.goto("/auth/login");
+  await userPage.getByRole("link", { name: "Create an account" }).click();
+  await userPage.getByTestId("email").fill("user_user@test.com");
+  await userPage.getByTestId("password").fill("test");
+  await userPage.getByRole("button", { name: "Sign Up" }).click();
+  await userPage.waitForURL("/chat");
+
+  await userContext.storageState({ path: "user_auth.json" });
+  await userContext.close();
+
+  await browser.close();
+}
+
+export default globalSetup;

web/tests/e2e/user_auth.setup.ts

@@ -0,0 +1,40 @@
+import { test as setup, expect } from "@playwright/test";
+import { loginAs } from "./utils/auth";
+
+setup("create test user", async ({ page, context }) => {
+  // Login as admin
+  await loginAs(page, "admin");
+
+  // Create test user
+  await page.goto("http://localhost:3000/admin/indexing/status");
+  await page.getByRole("button", { name: "Users" }).click();
+  await page.getByRole("button", { name: "Invite Users" }).click();
+  await page.locator("#emails").click();
+  await page.locator("#emails").fill("user_user@test.com");
+  await page.getByRole("button", { name: "Add!" }).click();
+
+  // Logout
+  await page.getByText("A", { exact: true }).click();
+  await page.getByText("Log out").click();
+
+  // Create account for the invited user
+  await page.goto("http://localhost:3000/auth/login");
+  await page.getByRole("link", { name: "Create an account" }).click();
+  await page.getByTestId("email").click();
+  await page.waitForTimeout(500);
+  await page.getByTestId("email").fill("user_user@test.com");
+  await page.waitForTimeout(500);
+  await page.getByTestId("password").click();
+  await page.waitForTimeout(500);
+  await page.getByTestId("password").fill("test");
+  await page.waitForTimeout(500);
+  await page.getByRole("button", { name: "Sign Up" }).click();
+  await page.waitForTimeout(2000);
+
+  // Verify successful account creation
+  await page.waitForURL("http://localhost:3000/chat");
+  await expect(page).toHaveURL("http://localhost:3000/chat");
+
+  // Save authentication state
+  await context.storageState({ path: "user_auth.json" });
+});

web/tests/e2e/utils/auth-utils.ts

@@ -0,0 +1,22 @@
+import { test as base, expect, Page, BrowserContext } from "@playwright/test";
+
+type AuthFixture = {
+  auth: (
+    role: "admin" | "user"
+  ) => Promise<{ context: BrowserContext; page: Page }>;
+};
+
+export const test = base.extend<AuthFixture>({
+  auth: async ({ browser, contextOptions }, use) => {
+    await use(async (role: "admin" | "user") => {
+      const context = await browser.newContext({
+        ...contextOptions, // Inherit all context options, including 'headless'
+        storageState: role === "admin" ? "admin_auth.json" : "user_auth.json",
+      });
+      const page = await context.newPage();
+      return { context, page };
+    });
+  },
+});
+
+export { expect };

web/tests/e2e/utils/auth.ts

@@ -0,0 +1,22 @@
+import { Page } from "@playwright/test";
+import { TEST_CREDENTIALS, TEST_USER_CREDENTIALS } from "../constants";
+
+export async function loginAs(page: Page, userType: "admin" | "user") {
+  const { email, password } =
+    userType === "admin" ? TEST_CREDENTIALS : TEST_USER_CREDENTIALS;
+
+  await page.goto("http://localhost:3000/chat");
+
+  await page.waitForURL("http://localhost:3000/auth/login?next=%2Fchat");
+
+  await page.fill("#email", email);
+  await page.fill("#password", password);
+
+  // Click the login button
+  await page.click('button[type="submit"]');
+
+  // Log the entire page contents after login
+  console.log("Page contents after login:", await page.content());
+
+  await page.waitForURL("http://localhost:3000/chat");
+}

web/user_auth.json

@@ -0,0 +1,15 @@
+{
+  "cookies": [
+    {
+      "name": "fastapiusersauth",
+      "value": "lbgX6spEvhdzvCE6uNugvMQE3AqtSW1JVEZb1hXF79A",
+      "domain": "localhost",
+      "path": "/",
+      "expires": 1734908599.292503,
+      "httpOnly": true,
+      "secure": false,
+      "sameSite": "Lax"
+    }
+  ],
+  "origins": []
+}