remove oauth.py again

# Conflicts:
#	backend/ee/onyx/server/oauth.py

.github/CODEOWNERS

@@ -0,0 +1 @@
+* @onyx-dot-app/onyx-core-team

.github/workflows/nightly-scan-licenses.yml

@@ -53,24 +53,90 @@ jobs:
           exclude: '(?i)^(pylint|aio[-_]*).*'
           
       - name: Print report
-        if: ${{ always() }}
+        if: always()
         run: echo "${{ steps.license_check_report.outputs.report }}"
       
       - name: Install npm dependencies
         working-directory: ./web
         run: npm ci
-        
-      - name: Run Trivy vulnerability scanner in repo mode
-        uses: aquasecurity/trivy-action@0.28.0
-        with:
-          scan-type: fs
-          scanners: license
-          format: table
-#           format: sarif
-#           output: trivy-results.sarif
-          severity: HIGH,CRITICAL
 
-#       - name: Upload Trivy scan results to GitHub Security tab
-#         uses: github/codeql-action/upload-sarif@v3
+        # be careful enabling the sarif and upload as it may spam the security tab
+        # with a huge amount of items. Work out the issues before enabling upload.       
+#       - name: Run Trivy vulnerability scanner in repo mode
+#         if: always()
+#         uses: aquasecurity/trivy-action@0.29.0
 #         with:
-#           sarif_file: trivy-results.sarif
+#           scan-type: fs
+#           scan-ref: .
+#           scanners: license
+#           format: table
+#           severity: HIGH,CRITICAL
+# #           format: sarif
+# #           output: trivy-results.sarif
+# 
+# #       - name: Upload Trivy scan results to GitHub Security tab
+# #         uses: github/codeql-action/upload-sarif@v3
+# #         with:
+# #           sarif_file: trivy-results.sarif
+
+  scan-trivy:
+    # See https://runs-on.com/runners/linux/
+    runs-on: [runs-on,runner=2cpu-linux-x64,"run-id=${{ github.run_id }}"]
+      
+    steps:
+    - name: Set up Docker Buildx
+      uses: docker/setup-buildx-action@v3
+
+    - name: Login to Docker Hub
+      uses: docker/login-action@v3
+      with:
+        username: ${{ secrets.DOCKER_USERNAME }}
+        password: ${{ secrets.DOCKER_TOKEN }}
+
+    # Backend
+    - name: Pull backend docker image
+      run: docker pull onyxdotapp/onyx-backend:latest
+
+    - name: Run Trivy vulnerability scanner on backend
+      uses: aquasecurity/trivy-action@0.29.0
+      env:
+        TRIVY_DB_REPOSITORY: 'public.ecr.aws/aquasecurity/trivy-db:2'
+        TRIVY_JAVA_DB_REPOSITORY: 'public.ecr.aws/aquasecurity/trivy-java-db:1'
+      with:
+        image-ref: onyxdotapp/onyx-backend:latest
+        scanners: license
+        severity: HIGH,CRITICAL
+        vuln-type: library
+        exit-code: 0  # Set to 1 if we want a failed scan to fail the workflow
+
+    # Web server
+    - name: Pull web server docker image
+      run: docker pull onyxdotapp/onyx-web-server:latest
+          
+    - name: Run Trivy vulnerability scanner on web server
+      uses: aquasecurity/trivy-action@0.29.0
+      env:
+        TRIVY_DB_REPOSITORY: 'public.ecr.aws/aquasecurity/trivy-db:2'
+        TRIVY_JAVA_DB_REPOSITORY: 'public.ecr.aws/aquasecurity/trivy-java-db:1'
+      with:
+        image-ref: onyxdotapp/onyx-web-server:latest
+        scanners: license
+        severity: HIGH,CRITICAL
+        vuln-type: library
+        exit-code: 0
+
+    # Model server
+    - name: Pull model server docker image
+      run: docker pull onyxdotapp/onyx-model-server:latest
+
+    - name: Run Trivy vulnerability scanner
+      uses: aquasecurity/trivy-action@0.29.0
+      env:
+        TRIVY_DB_REPOSITORY: 'public.ecr.aws/aquasecurity/trivy-db:2'
+        TRIVY_JAVA_DB_REPOSITORY: 'public.ecr.aws/aquasecurity/trivy-java-db:1'
+      with:
+        image-ref: onyxdotapp/onyx-model-server:latest
+        scanners: license
+        severity: HIGH,CRITICAL
+        vuln-type: library
+        exit-code: 0

backend/alembic/env.py

@@ -6,7 +6,7 @@ from onyx.configs.app_configs import POSTGRES_PORT
 from onyx.configs.app_configs import POSTGRES_USER
 from onyx.configs.app_configs import AWS_REGION_NAME
 from onyx.db.engine import build_connection_string
-from onyx.db.engine import get_all_tenant_ids
+from onyx.db.tenant import get_all_tenant_ids
 from sqlalchemy import event
 from sqlalchemy import pool
 from sqlalchemy import text

backend/alembic/versions/3bd4c84fe72f_improved_index.py

@@ -0,0 +1,84 @@
+"""improved index
+
+Revision ID: 3bd4c84fe72f
+Revises: 8f43500ee275
+Create Date: 2025-02-26 13:07:56.217791
+
+"""
+from alembic import op
+
+
+# revision identifiers, used by Alembic.
+revision = "3bd4c84fe72f"
+down_revision = "8f43500ee275"
+branch_labels = None
+depends_on = None
+
+
+# NOTE:
+# This migration addresses issues with the previous migration (8f43500ee275) which caused
+# an outage by creating an index without using CONCURRENTLY. This migration:
+#
+# 1. Creates more efficient full-text search capabilities using tsvector columns and GIN indexes
+# 2. Uses CONCURRENTLY for all index creation to prevent table locking
+# 3. Explicitly manages transactions with COMMIT statements to allow CONCURRENTLY to work
+# (see: https://www.postgresql.org/docs/9.4/sql-createindex.html#SQL-CREATEINDEX-CONCURRENTLY)
+# (see: https://github.com/sqlalchemy/alembic/issues/277)
+# 4. Adds indexes to both chat_message and chat_session tables for comprehensive search
+
+
+def upgrade() -> None:
+    # Create a GIN index for full-text search on chat_message.message
+    op.execute(
+        """
+        ALTER TABLE chat_message
+        ADD COLUMN message_tsv tsvector
+        GENERATED ALWAYS AS (to_tsvector('english', message)) STORED;
+        """
+    )
+
+    # Commit the current transaction before creating concurrent indexes
+    op.execute("COMMIT")
+
+    op.execute(
+        """
+        CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_chat_message_tsv
+        ON chat_message
+        USING GIN (message_tsv)
+        """
+    )
+
+    # Also add a stored tsvector column for chat_session.description
+    op.execute(
+        """
+        ALTER TABLE chat_session
+        ADD COLUMN description_tsv tsvector
+        GENERATED ALWAYS AS (to_tsvector('english', coalesce(description, ''))) STORED;
+        """
+    )
+
+    # Commit again before creating the second concurrent index
+    op.execute("COMMIT")
+
+    op.execute(
+        """
+        CREATE INDEX CONCURRENTLY IF NOT EXISTS idx_chat_session_desc_tsv
+        ON chat_session
+        USING GIN (description_tsv)
+        """
+    )
+
+
+def downgrade() -> None:
+    # Drop the indexes first (use CONCURRENTLY for dropping too)
+    op.execute("COMMIT")
+    op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_chat_message_tsv;")
+
+    op.execute("COMMIT")
+    op.execute("DROP INDEX CONCURRENTLY IF EXISTS idx_chat_session_desc_tsv;")
+
+    # Then drop the columns
+    op.execute("ALTER TABLE chat_message DROP COLUMN IF EXISTS message_tsv;")
+    op.execute("ALTER TABLE chat_session DROP COLUMN IF EXISTS description_tsv;")
+
+    op.execute("DROP INDEX IF EXISTS idx_chat_message_message_lower;")

backend/alembic/versions/b7c2b63c4a03_add_background_reindex_enabled_field.py

@@ -0,0 +1,55 @@
+"""add background_reindex_enabled field
+
+Revision ID: b7c2b63c4a03
+Revises: f11b408e39d3
+Create Date: 2024-03-26 12:34:56.789012
+
+"""
+from alembic import op
+import sqlalchemy as sa
+
+from onyx.db.enums import EmbeddingPrecision
+
+
+# revision identifiers, used by Alembic.
+revision = "b7c2b63c4a03"
+down_revision = "f11b408e39d3"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # Add background_reindex_enabled column with default value of True
+    op.add_column(
+        "search_settings",
+        sa.Column(
+            "background_reindex_enabled",
+            sa.Boolean(),
+            nullable=False,
+            server_default="true",
+        ),
+    )
+
+    # Add embedding_precision column with default value of FLOAT
+    op.add_column(
+        "search_settings",
+        sa.Column(
+            "embedding_precision",
+            sa.Enum(EmbeddingPrecision, native_enum=False),
+            nullable=False,
+            server_default=EmbeddingPrecision.FLOAT.name,
+        ),
+    )
+
+    # Add reduced_dimension column with default value of None
+    op.add_column(
+        "search_settings",
+        sa.Column("reduced_dimension", sa.Integer(), nullable=True),
+    )
+
+
+def downgrade() -> None:
+    # Remove the background_reindex_enabled column
+    op.drop_column("search_settings", "background_reindex_enabled")
+    op.drop_column("search_settings", "embedding_precision")
+    op.drop_column("search_settings", "reduced_dimension")

backend/alembic/versions/f11b408e39d3_force_lowercase_all_users.py

@@ -0,0 +1,36 @@
+"""force lowercase all users
+
+Revision ID: f11b408e39d3
+Revises: 3bd4c84fe72f
+Create Date: 2025-02-26 17:04:55.683500
+
+"""
+
+
+# revision identifiers, used by Alembic.
+revision = "f11b408e39d3"
+down_revision = "3bd4c84fe72f"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # 1) Convert all existing user emails to lowercase
+    from alembic import op
+
+    op.execute(
+        """
+        UPDATE "user"
+        SET email = LOWER(email)
+        """
+    )
+
+    # 2) Add a check constraint to ensure emails are always lowercase
+    op.create_check_constraint("ensure_lowercase_email", "user", "email = LOWER(email)")
+
+
+def downgrade() -> None:
+    # Drop the check constraint
+    from alembic import op
+
+    op.drop_constraint("ensure_lowercase_email", "user", type_="check")

backend/alembic_tenants/versions/34e3630c7f32_lowercase_multi_tenant_user_auth.py

@@ -0,0 +1,42 @@
+"""lowercase multi-tenant user auth
+
+Revision ID: 34e3630c7f32
+Revises: a4f6ee863c47
+Create Date: 2025-02-26 15:03:01.211894
+
+"""
+from alembic import op
+
+
+# revision identifiers, used by Alembic.
+revision = "34e3630c7f32"
+down_revision = "a4f6ee863c47"
+branch_labels = None
+depends_on = None
+
+
+def upgrade() -> None:
+    # 1) Convert all existing rows to lowercase
+    op.execute(
+        """
+        UPDATE user_tenant_mapping
+        SET email = LOWER(email)
+        """
+    )
+    # 2) Add a check constraint so that emails cannot be written in uppercase
+    op.create_check_constraint(
+        "ensure_lowercase_email",
+        "user_tenant_mapping",
+        "email = LOWER(email)",
+        schema="public",
+    )
+
+
+def downgrade() -> None:
+    # Drop the check constraint
+    op.drop_constraint(
+        "ensure_lowercase_email",
+        "user_tenant_mapping",
+        schema="public",
+        type_="check",
+    )

backend/ee/onyx/background/celery/apps/primary.py

@@ -5,7 +5,7 @@ from onyx.background.celery.apps.primary import celery_app
 from onyx.background.task_utils import build_celery_task_wrapper
 from onyx.configs.app_configs import JOB_TIMEOUT
 from onyx.db.chat import delete_chat_sessions_older_than
-from onyx.db.engine import get_session_with_current_tenant
+from onyx.db.session import get_session_with_current_tenant
 from onyx.server.settings.store import load_settings
 from onyx.utils.logger import setup_logger
 

backend/ee/onyx/configs/app_configs.py

@@ -59,10 +59,14 @@ SUPER_CLOUD_API_KEY = os.environ.get("SUPER_CLOUD_API_KEY", "api_key")
 
 OAUTH_SLACK_CLIENT_ID = os.environ.get("OAUTH_SLACK_CLIENT_ID", "")
 OAUTH_SLACK_CLIENT_SECRET = os.environ.get("OAUTH_SLACK_CLIENT_SECRET", "")
-OAUTH_CONFLUENCE_CLIENT_ID = os.environ.get("OAUTH_CONFLUENCE_CLIENT_ID", "")
-OAUTH_CONFLUENCE_CLIENT_SECRET = os.environ.get("OAUTH_CONFLUENCE_CLIENT_SECRET", "")
-OAUTH_JIRA_CLIENT_ID = os.environ.get("OAUTH_JIRA_CLIENT_ID", "")
-OAUTH_JIRA_CLIENT_SECRET = os.environ.get("OAUTH_JIRA_CLIENT_SECRET", "")
+OAUTH_CONFLUENCE_CLOUD_CLIENT_ID = os.environ.get(
+    "OAUTH_CONFLUENCE_CLOUD_CLIENT_ID", ""
+)
+OAUTH_CONFLUENCE_CLOUD_CLIENT_SECRET = os.environ.get(
+    "OAUTH_CONFLUENCE_CLOUD_CLIENT_SECRET", ""
+)
+OAUTH_JIRA_CLOUD_CLIENT_ID = os.environ.get("OAUTH_JIRA_CLOUD_CLIENT_ID", "")
+OAUTH_JIRA_CLOUD_CLIENT_SECRET = os.environ.get("OAUTH_JIRA_CLOUD_CLIENT_SECRET", "")
 OAUTH_GOOGLE_DRIVE_CLIENT_ID = os.environ.get("OAUTH_GOOGLE_DRIVE_CLIENT_ID", "")
 OAUTH_GOOGLE_DRIVE_CLIENT_SECRET = os.environ.get(
     "OAUTH_GOOGLE_DRIVE_CLIENT_SECRET", ""

backend/ee/onyx/db/user_group.py

@@ -424,7 +424,7 @@ def _validate_curator_status__no_commit(
         )
 
         # if the user is a curator in any of their groups, set their role to CURATOR
-        # otherwise, set their role to BASIC
+        # otherwise, set their role to BASIC only if they were previously a CURATOR
         if curator_relationships:
             user.role = UserRole.CURATOR
         elif user.role == UserRole.CURATOR:
@@ -631,7 +631,16 @@ def update_user_group(
     removed_users = db_session.scalars(
         select(User).where(User.id.in_(removed_user_ids))  # type: ignore
     ).unique()
-    _validate_curator_status__no_commit(db_session, list(removed_users))
+
+    # Filter out admin and global curator users before validating curator status
+    users_to_validate = [
+        user
+        for user in removed_users
+        if user.role not in [UserRole.ADMIN, UserRole.GLOBAL_CURATOR]
+    ]
+
+    if users_to_validate:
+        _validate_curator_status__no_commit(db_session, users_to_validate)
 
     # update "time_updated" to now
     db_user_group.time_last_modified_by_user = func.now()

backend/ee/onyx/external_permissions/confluence/doc_sync.py

@@ -9,12 +9,16 @@ from ee.onyx.external_permissions.confluence.constants import ALL_CONF_EMAILS_GR
 from onyx.access.models import DocExternalAccess
 from onyx.access.models import ExternalAccess
 from onyx.connectors.confluence.connector import ConfluenceConnector
+from onyx.connectors.confluence.onyx_confluence import (
+    get_user_email_from_username__server,
+)
 from onyx.connectors.confluence.onyx_confluence import OnyxConfluence
-from onyx.connectors.confluence.utils import get_user_email_from_username__server
+from onyx.connectors.credentials_provider import OnyxDBCredentialsProvider
 from onyx.connectors.models import SlimDocument
 from onyx.db.models import ConnectorCredentialPair
 from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface
 from onyx.utils.logger import setup_logger
+from shared_configs.contextvars import get_current_tenant_id
 
 logger = setup_logger()
 
@@ -342,7 +346,8 @@ def _fetch_all_page_restrictions(
 
 
 def confluence_doc_sync(
-    cc_pair: ConnectorCredentialPair, callback: IndexingHeartbeatInterface | None
+    cc_pair: ConnectorCredentialPair,
+    callback: IndexingHeartbeatInterface | None,
 ) -> list[DocExternalAccess]:
     """
     Adds the external permissions to the documents in postgres
@@ -354,7 +359,11 @@ def confluence_doc_sync(
     confluence_connector = ConfluenceConnector(
         **cc_pair.connector.connector_specific_config
     )
-    confluence_connector.load_credentials(cc_pair.credential.credential_json)
+
+    provider = OnyxDBCredentialsProvider(
+        get_current_tenant_id(), "confluence", cc_pair.credential_id
+    )
+    confluence_connector.set_credentials_provider(provider)
 
     is_cloud = cc_pair.connector.connector_specific_config.get("is_cloud", False)
 

backend/ee/onyx/external_permissions/confluence/group_sync.py

@@ -1,9 +1,11 @@
 from ee.onyx.db.external_perm import ExternalUserGroup
 from ee.onyx.external_permissions.confluence.constants import ALL_CONF_EMAILS_GROUP_NAME
 from onyx.background.error_logging import emit_background_error
-from onyx.connectors.confluence.onyx_confluence import build_confluence_client
+from onyx.connectors.confluence.onyx_confluence import (
+    get_user_email_from_username__server,
+)
 from onyx.connectors.confluence.onyx_confluence import OnyxConfluence
-from onyx.connectors.confluence.utils import get_user_email_from_username__server
+from onyx.connectors.credentials_provider import OnyxDBCredentialsProvider
 from onyx.db.models import ConnectorCredentialPair
 from onyx.utils.logger import setup_logger
 
@@ -61,13 +63,27 @@ def _build_group_member_email_map(
 
 
 def confluence_group_sync(
+    tenant_id: str,
     cc_pair: ConnectorCredentialPair,
 ) -> list[ExternalUserGroup]:
-    confluence_client = build_confluence_client(
-        credentials=cc_pair.credential.credential_json,
-        is_cloud=cc_pair.connector.connector_specific_config.get("is_cloud", False),
-        wiki_base=cc_pair.connector.connector_specific_config["wiki_base"],
-    )
+    provider = OnyxDBCredentialsProvider(tenant_id, "confluence", cc_pair.credential_id)
+    is_cloud = cc_pair.connector.connector_specific_config.get("is_cloud", False)
+    wiki_base: str = cc_pair.connector.connector_specific_config["wiki_base"]
+    url = wiki_base.rstrip("/")
+
+    probe_kwargs = {
+        "max_backoff_retries": 6,
+        "max_backoff_seconds": 10,
+    }
+
+    final_kwargs = {
+        "max_backoff_retries": 10,
+        "max_backoff_seconds": 60,
+    }
+
+    confluence_client = OnyxConfluence(is_cloud, url, provider)
+    confluence_client._probe_connection(**probe_kwargs)
+    confluence_client._initialize_connection(**final_kwargs)
 
     group_member_email_map = _build_group_member_email_map(
         confluence_client=confluence_client,

backend/ee/onyx/external_permissions/gmail/doc_sync.py

@@ -32,7 +32,8 @@ def _get_slim_doc_generator(
 
 
 def gmail_doc_sync(
-    cc_pair: ConnectorCredentialPair, callback: IndexingHeartbeatInterface | None
+    cc_pair: ConnectorCredentialPair,
+    callback: IndexingHeartbeatInterface | None,
 ) -> list[DocExternalAccess]:
     """
     Adds the external permissions to the documents in postgres

backend/ee/onyx/external_permissions/google_drive/doc_sync.py

@@ -145,7 +145,8 @@ def _get_permissions_from_slim_doc(
 
 
 def gdrive_doc_sync(
-    cc_pair: ConnectorCredentialPair, callback: IndexingHeartbeatInterface | None
+    cc_pair: ConnectorCredentialPair,
+    callback: IndexingHeartbeatInterface | None,
 ) -> list[DocExternalAccess]:
     """
     Adds the external permissions to the documents in postgres

backend/ee/onyx/external_permissions/google_drive/group_sync.py

@@ -119,6 +119,7 @@ def _build_onyx_groups(
 
 
 def gdrive_group_sync(
+    tenant_id: str,
     cc_pair: ConnectorCredentialPair,
 ) -> list[ExternalUserGroup]:
     # Initialize connector and build credential/service objects

backend/ee/onyx/external_permissions/post_query_censoring.py

@@ -6,8 +6,8 @@ from ee.onyx.external_permissions.salesforce.postprocessing import (
 )
 from onyx.configs.constants import DocumentSource
 from onyx.context.search.pipeline import InferenceChunk
-from onyx.db.engine import get_session_context_manager
 from onyx.db.models import User
+from onyx.db.session import get_session_context_manager
 from onyx.utils.logger import setup_logger
 
 logger = setup_logger()

backend/ee/onyx/external_permissions/salesforce/postprocessing.py

@@ -10,7 +10,7 @@ from ee.onyx.external_permissions.salesforce.utils import (
 )
 from onyx.configs.app_configs import BLURB_SIZE
 from onyx.context.search.models import InferenceChunk
-from onyx.db.engine import get_session_context_manager
+from onyx.db.session import get_session_context_manager
 from onyx.utils.logger import setup_logger
 
 logger = setup_logger()

backend/ee/onyx/external_permissions/slack/doc_sync.py

@@ -123,7 +123,8 @@ def _fetch_channel_permissions(
 
 
 def slack_doc_sync(
-    cc_pair: ConnectorCredentialPair, callback: IndexingHeartbeatInterface | None
+    cc_pair: ConnectorCredentialPair,
+    callback: IndexingHeartbeatInterface | None,
 ) -> list[DocExternalAccess]:
     """
     Adds the external permissions to the documents in postgres

backend/ee/onyx/external_permissions/sync_params.py

@@ -28,6 +28,7 @@ DocSyncFuncType = Callable[
 
 GroupSyncFuncType = Callable[
     [
+        str,
         ConnectorCredentialPair,
     ],
     list[ExternalUserGroup],

backend/ee/onyx/main.py

@@ -15,7 +15,7 @@ from ee.onyx.server.enterprise_settings.api import (
 )
 from ee.onyx.server.manage.standard_answer import router as standard_answer_router
 from ee.onyx.server.middleware.tenant_tracking import add_tenant_id_middleware
-from ee.onyx.server.oauth import router as oauth_router
+from ee.onyx.server.oauth.api import router as oauth_router
 from ee.onyx.server.query_and_chat.chat_backend import (
     router as chat_router,
 )
@@ -152,4 +152,8 @@ def get_application() -> FastAPI:
     # environment variable. Used to automate deployment for multiple environments.
     seed_db()
 
+    # for debugging discovered routes
+    # for route in application.router.routes:
+    #     print(f"Path: {route.path}, Methods: {route.methods}")
+
     return application

backend/ee/onyx/onyxbot/slack/handlers/handle_standard_answers.py

@@ -22,7 +22,7 @@ from onyx.onyxbot.slack.blocks import get_restate_blocks
 from onyx.onyxbot.slack.constants import GENERATE_ANSWER_BUTTON_ACTION_ID
 from onyx.onyxbot.slack.handlers.utils import send_team_member_message
 from onyx.onyxbot.slack.models import SlackMessageInfo
-from onyx.onyxbot.slack.utils import respond_in_thread
+from onyx.onyxbot.slack.utils import respond_in_thread_or_channel
 from onyx.onyxbot.slack.utils import update_emote_react
 from onyx.utils.logger import OnyxLoggingAdapter
 from onyx.utils.logger import setup_logger
@@ -216,7 +216,7 @@ def _handle_standard_answers(
         all_blocks = restate_question_blocks + answer_blocks
 
         try:
-            respond_in_thread(
+            respond_in_thread_or_channel(
                 client=client,
                 channel=message_info.channel_to_respond,
                 receiver_ids=receiver_ids,
@@ -231,6 +231,7 @@ def _handle_standard_answers(
                     client=client,
                     channel=message_info.channel_to_respond,
                     thread_ts=slack_thread_id,
+                    receiver_ids=receiver_ids,
                 )
 
             return True

backend/ee/onyx/server/analytics/api.py

@@ -19,8 +19,8 @@ from ee.onyx.db.analytics import fetch_query_analytics
 from ee.onyx.db.analytics import user_can_view_assistant_stats
 from onyx.auth.users import current_admin_user
 from onyx.auth.users import current_user
-from onyx.db.engine import get_session
 from onyx.db.models import User
+from onyx.db.session import get_session
 
 router = APIRouter(prefix="/analytics")
 

backend/ee/onyx/server/enterprise_settings/api.py

@@ -26,8 +26,8 @@ from onyx.auth.users import current_admin_user
 from onyx.auth.users import current_user_with_expired_token
 from onyx.auth.users import get_user_manager
 from onyx.auth.users import UserManager
-from onyx.db.engine import get_session
 from onyx.db.models import User
+from onyx.db.session import get_session
 from onyx.file_store.file_store import get_default_file_store
 from onyx.utils.logger import setup_logger
 

backend/ee/onyx/server/manage/standard_answer.py

@@ -17,8 +17,8 @@ from ee.onyx.server.manage.models import StandardAnswerCategory
 from ee.onyx.server.manage.models import StandardAnswerCategoryCreationRequest
 from ee.onyx.server.manage.models import StandardAnswerCreationRequest
 from onyx.auth.users import current_admin_user
-from onyx.db.engine import get_session
 from onyx.db.models import User
+from onyx.db.session import get_session
 
 router = APIRouter(prefix="/manage")
 

backend/ee/onyx/server/middleware/tenant_tracking.py

@@ -11,7 +11,7 @@ from ee.onyx.auth.users import decode_anonymous_user_jwt_token
 from ee.onyx.configs.app_configs import ANONYMOUS_USER_COOKIE_NAME
 from onyx.auth.api_key import extract_tenant_from_api_key_header
 from onyx.configs.constants import TENANT_ID_COOKIE_NAME
-from onyx.db.engine import is_valid_schema_name
+from onyx.db.utils import is_valid_schema_name
 from onyx.redis.redis_pool import retrieve_auth_token_data_from_redis
 from shared_configs.configs import MULTI_TENANT
 from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA

backend/ee/onyx/server/oauth.py

@@ -1,629 +0,0 @@
-import base64
-import json
-import uuid
-from typing import Any
-from typing import cast
-
-import requests
-from fastapi import APIRouter
-from fastapi import Depends
-from fastapi import HTTPException
-from fastapi.responses import JSONResponse
-from pydantic import BaseModel
-from sqlalchemy.orm import Session
-
-from ee.onyx.configs.app_configs import OAUTH_CONFLUENCE_CLIENT_ID
-from ee.onyx.configs.app_configs import OAUTH_CONFLUENCE_CLIENT_SECRET
-from ee.onyx.configs.app_configs import OAUTH_GOOGLE_DRIVE_CLIENT_ID
-from ee.onyx.configs.app_configs import OAUTH_GOOGLE_DRIVE_CLIENT_SECRET
-from ee.onyx.configs.app_configs import OAUTH_SLACK_CLIENT_ID
-from ee.onyx.configs.app_configs import OAUTH_SLACK_CLIENT_SECRET
-from onyx.auth.users import current_user
-from onyx.configs.app_configs import WEB_DOMAIN
-from onyx.configs.constants import DocumentSource
-from onyx.connectors.google_utils.google_auth import get_google_oauth_creds
-from onyx.connectors.google_utils.google_auth import sanitize_oauth_credentials
-from onyx.connectors.google_utils.shared_constants import (
-    DB_CREDENTIALS_AUTHENTICATION_METHOD,
-)
-from onyx.connectors.google_utils.shared_constants import (
-    DB_CREDENTIALS_DICT_TOKEN_KEY,
-)
-from onyx.connectors.google_utils.shared_constants import (
-    DB_CREDENTIALS_PRIMARY_ADMIN_KEY,
-)
-from onyx.connectors.google_utils.shared_constants import (
-    GoogleOAuthAuthenticationMethod,
-)
-from onyx.db.credentials import create_credential
-from onyx.db.engine import get_session
-from onyx.db.models import User
-from onyx.redis.redis_pool import get_redis_client
-from onyx.server.documents.models import CredentialBase
-from onyx.utils.logger import setup_logger
-from shared_configs.contextvars import get_current_tenant_id
-
-
-logger = setup_logger()
-
-router = APIRouter(prefix="/oauth")
-
-
-class SlackOAuth:
-    # https://knock.app/blog/how-to-authenticate-users-in-slack-using-oauth
-    # Example: https://api.slack.com/authentication/oauth-v2#exchanging
-
-    class OAuthSession(BaseModel):
-        """Stored in redis to be looked up on callback"""
-
-        email: str
-        redirect_on_success: str | None  # Where to send the user if OAuth flow succeeds
-
-    CLIENT_ID = OAUTH_SLACK_CLIENT_ID
-    CLIENT_SECRET = OAUTH_SLACK_CLIENT_SECRET
-
-    TOKEN_URL = "https://slack.com/api/oauth.v2.access"
-
-    # SCOPE is per https://docs.onyx.app/connectors/slack
-    BOT_SCOPE = (
-        "channels:history,"
-        "channels:read,"
-        "groups:history,"
-        "groups:read,"
-        "channels:join,"
-        "im:history,"
-        "users:read,"
-        "users:read.email,"
-        "usergroups:read"
-    )
-
-    REDIRECT_URI = f"{WEB_DOMAIN}/admin/connectors/slack/oauth/callback"
-    DEV_REDIRECT_URI = f"https://redirectmeto.com/{REDIRECT_URI}"
-
-    @classmethod
-    def generate_oauth_url(cls, state: str) -> str:
-        return cls._generate_oauth_url_helper(cls.REDIRECT_URI, state)
-
-    @classmethod
-    def generate_dev_oauth_url(cls, state: str) -> str:
-        """dev mode workaround for localhost testing
-        - https://www.nango.dev/blog/oauth-redirects-on-localhost-with-https
-        """
-
-        return cls._generate_oauth_url_helper(cls.DEV_REDIRECT_URI, state)
-
-    @classmethod
-    def _generate_oauth_url_helper(cls, redirect_uri: str, state: str) -> str:
-        url = (
-            f"https://slack.com/oauth/v2/authorize"
-            f"?client_id={cls.CLIENT_ID}"
-            f"&redirect_uri={redirect_uri}"
-            f"&scope={cls.BOT_SCOPE}"
-            f"&state={state}"
-        )
-        return url
-
-    @classmethod
-    def session_dump_json(cls, email: str, redirect_on_success: str | None) -> str:
-        """Temporary state to store in redis. to be looked up on auth response.
-        Returns a json string.
-        """
-        session = SlackOAuth.OAuthSession(
-            email=email, redirect_on_success=redirect_on_success
-        )
-        return session.model_dump_json()
-
-    @classmethod
-    def parse_session(cls, session_json: str) -> OAuthSession:
-        session = SlackOAuth.OAuthSession.model_validate_json(session_json)
-        return session
-
-
-class ConfluenceCloudOAuth:
-    """work in progress"""
-
-    # https://developer.atlassian.com/cloud/confluence/oauth-2-3lo-apps/
-
-    class OAuthSession(BaseModel):
-        """Stored in redis to be looked up on callback"""
-
-        email: str
-        redirect_on_success: str | None  # Where to send the user if OAuth flow succeeds
-
-    CLIENT_ID = OAUTH_CONFLUENCE_CLIENT_ID
-    CLIENT_SECRET = OAUTH_CONFLUENCE_CLIENT_SECRET
-    TOKEN_URL = "https://auth.atlassian.com/oauth/token"
-
-    # All read scopes per https://developer.atlassian.com/cloud/confluence/scopes-for-oauth-2-3LO-and-forge-apps/
-    CONFLUENCE_OAUTH_SCOPE = (
-        "read:confluence-props%20"
-        "read:confluence-content.all%20"
-        "read:confluence-content.summary%20"
-        "read:confluence-content.permission%20"
-        "read:confluence-user%20"
-        "read:confluence-groups%20"
-        "readonly:content.attachment:confluence"
-    )
-
-    REDIRECT_URI = f"{WEB_DOMAIN}/admin/connectors/confluence/oauth/callback"
-    DEV_REDIRECT_URI = f"https://redirectmeto.com/{REDIRECT_URI}"
-
-    # eventually for Confluence Data Center
-    # oauth_url = (
-    #     f"http://localhost:8090/rest/oauth/v2/authorize?client_id={CONFLUENCE_OAUTH_CLIENT_ID}"
-    #     f"&scope={CONFLUENCE_OAUTH_SCOPE_2}"
-    #     f"&redirect_uri={redirectme_uri}"
-    # )
-
-    @classmethod
-    def generate_oauth_url(cls, state: str) -> str:
-        return cls._generate_oauth_url_helper(cls.REDIRECT_URI, state)
-
-    @classmethod
-    def generate_dev_oauth_url(cls, state: str) -> str:
-        """dev mode workaround for localhost testing
-        - https://www.nango.dev/blog/oauth-redirects-on-localhost-with-https
-        """
-        return cls._generate_oauth_url_helper(cls.DEV_REDIRECT_URI, state)
-
-    @classmethod
-    def _generate_oauth_url_helper(cls, redirect_uri: str, state: str) -> str:
-        url = (
-            "https://auth.atlassian.com/authorize"
-            f"?audience=api.atlassian.com"
-            f"&client_id={cls.CLIENT_ID}"
-            f"&redirect_uri={redirect_uri}"
-            f"&scope={cls.CONFLUENCE_OAUTH_SCOPE}"
-            f"&state={state}"
-            "&response_type=code"
-            "&prompt=consent"
-        )
-        return url
-
-    @classmethod
-    def session_dump_json(cls, email: str, redirect_on_success: str | None) -> str:
-        """Temporary state to store in redis. to be looked up on auth response.
-        Returns a json string.
-        """
-        session = ConfluenceCloudOAuth.OAuthSession(
-            email=email, redirect_on_success=redirect_on_success
-        )
-        return session.model_dump_json()
-
-    @classmethod
-    def parse_session(cls, session_json: str) -> SlackOAuth.OAuthSession:
-        session = SlackOAuth.OAuthSession.model_validate_json(session_json)
-        return session
-
-
-class GoogleDriveOAuth:
-    # https://developers.google.com/identity/protocols/oauth2
-    # https://developers.google.com/identity/protocols/oauth2/web-server
-
-    class OAuthSession(BaseModel):
-        """Stored in redis to be looked up on callback"""
-
-        email: str
-        redirect_on_success: str | None  # Where to send the user if OAuth flow succeeds
-
-    CLIENT_ID = OAUTH_GOOGLE_DRIVE_CLIENT_ID
-    CLIENT_SECRET = OAUTH_GOOGLE_DRIVE_CLIENT_SECRET
-
-    TOKEN_URL = "https://oauth2.googleapis.com/token"
-
-    # SCOPE is per https://docs.onyx.app/connectors/google-drive
-    # TODO: Merge with or use google_utils.GOOGLE_SCOPES
-    SCOPE = (
-        "https://www.googleapis.com/auth/drive.readonly%20"
-        "https://www.googleapis.com/auth/drive.metadata.readonly%20"
-        "https://www.googleapis.com/auth/admin.directory.user.readonly%20"
-        "https://www.googleapis.com/auth/admin.directory.group.readonly"
-    )
-
-    REDIRECT_URI = f"{WEB_DOMAIN}/admin/connectors/google-drive/oauth/callback"
-    DEV_REDIRECT_URI = f"https://redirectmeto.com/{REDIRECT_URI}"
-
-    @classmethod
-    def generate_oauth_url(cls, state: str) -> str:
-        return cls._generate_oauth_url_helper(cls.REDIRECT_URI, state)
-
-    @classmethod
-    def generate_dev_oauth_url(cls, state: str) -> str:
-        """dev mode workaround for localhost testing
-        - https://www.nango.dev/blog/oauth-redirects-on-localhost-with-https
-        """
-
-        return cls._generate_oauth_url_helper(cls.DEV_REDIRECT_URI, state)
-
-    @classmethod
-    def _generate_oauth_url_helper(cls, redirect_uri: str, state: str) -> str:
-        # without prompt=consent, a refresh token is only issued the first time the user approves
-        url = (
-            f"https://accounts.google.com/o/oauth2/v2/auth"
-            f"?client_id={cls.CLIENT_ID}"
-            f"&redirect_uri={redirect_uri}"
-            "&response_type=code"
-            f"&scope={cls.SCOPE}"
-            "&access_type=offline"
-            f"&state={state}"
-            "&prompt=consent"
-        )
-        return url
-
-    @classmethod
-    def session_dump_json(cls, email: str, redirect_on_success: str | None) -> str:
-        """Temporary state to store in redis. to be looked up on auth response.
-        Returns a json string.
-        """
-        session = GoogleDriveOAuth.OAuthSession(
-            email=email, redirect_on_success=redirect_on_success
-        )
-        return session.model_dump_json()
-
-    @classmethod
-    def parse_session(cls, session_json: str) -> OAuthSession:
-        session = GoogleDriveOAuth.OAuthSession.model_validate_json(session_json)
-        return session
-
-
-@router.post("/prepare-authorization-request")
-def prepare_authorization_request(
-    connector: DocumentSource,
-    redirect_on_success: str | None,
-    user: User = Depends(current_user),
-) -> JSONResponse:
-    """Used by the frontend to generate the url for the user's browser during auth request.
-
-    Example: https://www.oauth.com/oauth2-servers/authorization/the-authorization-request/
-    """
-    tenant_id = get_current_tenant_id()
-
-    # create random oauth state param for security and to retrieve user data later
-    oauth_uuid = uuid.uuid4()
-    oauth_uuid_str = str(oauth_uuid)
-
-    # urlsafe b64 encode the uuid for the oauth url
-    oauth_state = (
-        base64.urlsafe_b64encode(oauth_uuid.bytes).rstrip(b"=").decode("utf-8")
-    )
-    session: str
-
-    if connector == DocumentSource.SLACK:
-        oauth_url = SlackOAuth.generate_oauth_url(oauth_state)
-        session = SlackOAuth.session_dump_json(
-            email=user.email, redirect_on_success=redirect_on_success
-        )
-    elif connector == DocumentSource.GOOGLE_DRIVE:
-        oauth_url = GoogleDriveOAuth.generate_oauth_url(oauth_state)
-        session = GoogleDriveOAuth.session_dump_json(
-            email=user.email, redirect_on_success=redirect_on_success
-        )
-    # elif connector == DocumentSource.CONFLUENCE:
-    #     oauth_url = ConfluenceCloudOAuth.generate_oauth_url(oauth_state)
-    #     session = ConfluenceCloudOAuth.session_dump_json(
-    #         email=user.email, redirect_on_success=redirect_on_success
-    #     )
-    # elif connector == DocumentSource.JIRA:
-    #     oauth_url = JiraCloudOAuth.generate_dev_oauth_url(oauth_state)
-    else:
-        oauth_url = None
-
-    if not oauth_url:
-        raise HTTPException(
-            status_code=404,
-            detail=f"The document source type {connector} does not have OAuth implemented",
-        )
-
-    r = get_redis_client(tenant_id=tenant_id)
-
-    # store important session state to retrieve when the user is redirected back
-    # 10 min is the max we want an oauth flow to be valid
-    r.set(f"da_oauth:{oauth_uuid_str}", session, ex=600)
-
-    return JSONResponse(content={"url": oauth_url})
-
-
-@router.post("/connector/slack/callback")
-def handle_slack_oauth_callback(
-    code: str,
-    state: str,
-    user: User = Depends(current_user),
-    db_session: Session = Depends(get_session),
-) -> JSONResponse:
-    if not SlackOAuth.CLIENT_ID or not SlackOAuth.CLIENT_SECRET:
-        raise HTTPException(
-            status_code=500,
-            detail="Slack client ID or client secret is not configured.",
-        )
-
-    r = get_redis_client()
-
-    # recover the state
-    padded_state = state + "=" * (
-        -len(state) % 4
-    )  # Add padding back (Base64 decoding requires padding)
-    uuid_bytes = base64.urlsafe_b64decode(
-        padded_state
-    )  # Decode the Base64 string back to bytes
-
-    # Convert bytes back to a UUID
-    oauth_uuid = uuid.UUID(bytes=uuid_bytes)
-    oauth_uuid_str = str(oauth_uuid)
-
-    r_key = f"da_oauth:{oauth_uuid_str}"
-
-    session_json_bytes = cast(bytes, r.get(r_key))
-    if not session_json_bytes:
-        raise HTTPException(
-            status_code=400,
-            detail=f"Slack OAuth failed - OAuth state key not found: key={r_key}",
-        )
-
-    session_json = session_json_bytes.decode("utf-8")
-    try:
-        session = SlackOAuth.parse_session(session_json)
-
-        # Exchange the authorization code for an access token
-        response = requests.post(
-            SlackOAuth.TOKEN_URL,
-            headers={"Content-Type": "application/x-www-form-urlencoded"},
-            data={
-                "client_id": SlackOAuth.CLIENT_ID,
-                "client_secret": SlackOAuth.CLIENT_SECRET,
-                "code": code,
-                "redirect_uri": SlackOAuth.REDIRECT_URI,
-            },
-        )
-
-        response_data = response.json()
-
-        if not response_data.get("ok"):
-            raise HTTPException(
-                status_code=400,
-                detail=f"Slack OAuth failed: {response_data.get('error')}",
-            )
-
-        # Extract token and team information
-        access_token: str = response_data.get("access_token")
-        team_id: str = response_data.get("team", {}).get("id")
-        authed_user_id: str = response_data.get("authed_user", {}).get("id")
-
-        credential_info = CredentialBase(
-            credential_json={"slack_bot_token": access_token},
-            admin_public=True,
-            source=DocumentSource.SLACK,
-            name="Slack OAuth",
-        )
-
-        create_credential(credential_info, user, db_session)
-    except Exception as e:
-        return JSONResponse(
-            status_code=500,
-            content={
-                "success": False,
-                "message": f"An error occurred during Slack OAuth: {str(e)}",
-            },
-        )
-    finally:
-        r.delete(r_key)
-
-    # return the result
-    return JSONResponse(
-        content={
-            "success": True,
-            "message": "Slack OAuth completed successfully.",
-            "team_id": team_id,
-            "authed_user_id": authed_user_id,
-            "redirect_on_success": session.redirect_on_success,
-        }
-    )
-
-
-# Work in progress
-# @router.post("/connector/confluence/callback")
-# def handle_confluence_oauth_callback(
-#     code: str,
-#     state: str,
-#     user: User = Depends(current_user),
-#     db_session: Session = Depends(get_session),
-#     tenant_id: str | None = Depends(get_current_tenant_id),
-# ) -> JSONResponse:
-#     if not ConfluenceCloudOAuth.CLIENT_ID or not ConfluenceCloudOAuth.CLIENT_SECRET:
-#         raise HTTPException(
-#             status_code=500,
-#             detail="Confluence client ID or client secret is not configured."
-#         )
-
-#     r = get_redis_client(tenant_id=tenant_id)
-
-#     # recover the state
-#     padded_state = state + '=' * (-len(state) % 4)  # Add padding back (Base64 decoding requires padding)
-#     uuid_bytes = base64.urlsafe_b64decode(padded_state)  # Decode the Base64 string back to bytes
-
-#     # Convert bytes back to a UUID
-#     oauth_uuid = uuid.UUID(bytes=uuid_bytes)
-#     oauth_uuid_str = str(oauth_uuid)
-
-#     r_key = f"da_oauth:{oauth_uuid_str}"
-
-#     result = r.get(r_key)
-#     if not result:
-#         raise HTTPException(
-#             status_code=400,
-#             detail=f"Confluence OAuth failed - OAuth state key not found: key={r_key}"
-#         )
-
-#     try:
-#         session = ConfluenceCloudOAuth.parse_session(result)
-
-#         # Exchange the authorization code for an access token
-#         response = requests.post(
-#             ConfluenceCloudOAuth.TOKEN_URL,
-#             headers={"Content-Type": "application/x-www-form-urlencoded"},
-#             data={
-#                 "client_id": ConfluenceCloudOAuth.CLIENT_ID,
-#                 "client_secret": ConfluenceCloudOAuth.CLIENT_SECRET,
-#                 "code": code,
-#                 "redirect_uri": ConfluenceCloudOAuth.DEV_REDIRECT_URI,
-#             },
-#         )
-
-#         response_data = response.json()
-
-#         if not response_data.get("ok"):
-#             raise HTTPException(
-#                 status_code=400,
-#                 detail=f"ConfluenceCloudOAuth OAuth failed: {response_data.get('error')}"
-#             )
-
-#         # Extract token and team information
-#         access_token: str = response_data.get("access_token")
-#         team_id: str = response_data.get("team", {}).get("id")
-#         authed_user_id: str = response_data.get("authed_user", {}).get("id")
-
-#         credential_info = CredentialBase(
-#             credential_json={"slack_bot_token": access_token},
-#             admin_public=True,
-#             source=DocumentSource.CONFLUENCE,
-#             name="Confluence OAuth",
-#         )
-
-#         logger.info(f"Slack access token: {access_token}")
-
-#         credential = create_credential(credential_info, user, db_session)
-
-#         logger.info(f"new_credential_id={credential.id}")
-#     except Exception as e:
-#         return JSONResponse(
-#             status_code=500,
-#             content={
-#                 "success": False,
-#                 "message": f"An error occurred during Slack OAuth: {str(e)}",
-#             },
-#         )
-#     finally:
-#         r.delete(r_key)
-
-#     # return the result
-#     return JSONResponse(
-#         content={
-#             "success": True,
-#             "message": "Slack OAuth completed successfully.",
-#             "team_id": team_id,
-#             "authed_user_id": authed_user_id,
-#             "redirect_on_success": session.redirect_on_success,
-#         }
-#     )
-
-
-@router.post("/connector/google-drive/callback")
-def handle_google_drive_oauth_callback(
-    code: str,
-    state: str,
-    user: User = Depends(current_user),
-    db_session: Session = Depends(get_session),
-) -> JSONResponse:
-    if not GoogleDriveOAuth.CLIENT_ID or not GoogleDriveOAuth.CLIENT_SECRET:
-        raise HTTPException(
-            status_code=500,
-            detail="Google Drive client ID or client secret is not configured.",
-        )
-
-    r = get_redis_client()
-
-    # recover the state
-    padded_state = state + "=" * (
-        -len(state) % 4
-    )  # Add padding back (Base64 decoding requires padding)
-    uuid_bytes = base64.urlsafe_b64decode(
-        padded_state
-    )  # Decode the Base64 string back to bytes
-
-    # Convert bytes back to a UUID
-    oauth_uuid = uuid.UUID(bytes=uuid_bytes)
-    oauth_uuid_str = str(oauth_uuid)
-
-    r_key = f"da_oauth:{oauth_uuid_str}"
-
-    session_json_bytes = cast(bytes, r.get(r_key))
-    if not session_json_bytes:
-        raise HTTPException(
-            status_code=400,
-            detail=f"Google Drive OAuth failed - OAuth state key not found: key={r_key}",
-        )
-
-    session_json = session_json_bytes.decode("utf-8")
-    session: GoogleDriveOAuth.OAuthSession
-    try:
-        session = GoogleDriveOAuth.parse_session(session_json)
-
-        # Exchange the authorization code for an access token
-        response = requests.post(
-            GoogleDriveOAuth.TOKEN_URL,
-            headers={"Content-Type": "application/x-www-form-urlencoded"},
-            data={
-                "client_id": GoogleDriveOAuth.CLIENT_ID,
-                "client_secret": GoogleDriveOAuth.CLIENT_SECRET,
-                "code": code,
-                "redirect_uri": GoogleDriveOAuth.REDIRECT_URI,
-                "grant_type": "authorization_code",
-            },
-        )
-
-        response.raise_for_status()
-
-        authorization_response: dict[str, Any] = response.json()
-
-        # the connector wants us to store the json in its authorized_user_info format
-        # returned from OAuthCredentials.get_authorized_user_info().
-        # So refresh immediately via get_google_oauth_creds with the params filled in
-        # from fields in authorization_response to get the json we need
-        authorized_user_info = {}
-        authorized_user_info["client_id"] = OAUTH_GOOGLE_DRIVE_CLIENT_ID
-        authorized_user_info["client_secret"] = OAUTH_GOOGLE_DRIVE_CLIENT_SECRET
-        authorized_user_info["refresh_token"] = authorization_response["refresh_token"]
-
-        token_json_str = json.dumps(authorized_user_info)
-        oauth_creds = get_google_oauth_creds(
-            token_json_str=token_json_str, source=DocumentSource.GOOGLE_DRIVE
-        )
-        if not oauth_creds:
-            raise RuntimeError("get_google_oauth_creds returned None.")
-
-        # save off the credentials
-        oauth_creds_sanitized_json_str = sanitize_oauth_credentials(oauth_creds)
-
-        credential_dict: dict[str, str] = {}
-        credential_dict[DB_CREDENTIALS_DICT_TOKEN_KEY] = oauth_creds_sanitized_json_str
-        credential_dict[DB_CREDENTIALS_PRIMARY_ADMIN_KEY] = session.email
-        credential_dict[
-            DB_CREDENTIALS_AUTHENTICATION_METHOD
-        ] = GoogleOAuthAuthenticationMethod.OAUTH_INTERACTIVE.value
-
-        credential_info = CredentialBase(
-            credential_json=credential_dict,
-            admin_public=True,
-            source=DocumentSource.GOOGLE_DRIVE,
-            name="OAuth (interactive)",
-        )
-
-        create_credential(credential_info, user, db_session)
-    except Exception as e:
-        return JSONResponse(
-            status_code=500,
-            content={
-                "success": False,
-                "message": f"An error occurred during Google Drive OAuth: {str(e)}",
-            },
-        )
-    finally:
-        r.delete(r_key)
-
-    # return the result
-    return JSONResponse(
-        content={
-            "success": True,
-            "message": "Google Drive OAuth completed successfully.",
-            "redirect_on_success": session.redirect_on_success,
-        }
-    )

backend/ee/onyx/server/oauth/api.py

@@ -0,0 +1,91 @@
+import base64
+import uuid
+
+from fastapi import Depends
+from fastapi import HTTPException
+from fastapi.responses import JSONResponse
+
+from ee.onyx.server.oauth.api_router import router
+from ee.onyx.server.oauth.confluence_cloud import ConfluenceCloudOAuth
+from ee.onyx.server.oauth.google_drive import GoogleDriveOAuth
+from ee.onyx.server.oauth.slack import SlackOAuth
+from onyx.auth.users import current_admin_user
+from onyx.configs.app_configs import DEV_MODE
+from onyx.configs.constants import DocumentSource
+from onyx.db.models import User
+from onyx.redis.redis_pool import get_redis_client
+from onyx.utils.logger import setup_logger
+from shared_configs.contextvars import get_current_tenant_id
+
+logger = setup_logger()
+
+
+@router.post("/prepare-authorization-request")
+def prepare_authorization_request(
+    connector: DocumentSource,
+    redirect_on_success: str | None,
+    user: User = Depends(current_admin_user),
+    tenant_id: str | None = Depends(get_current_tenant_id),
+) -> JSONResponse:
+    """Used by the frontend to generate the url for the user's browser during auth request.
+
+    Example: https://www.oauth.com/oauth2-servers/authorization/the-authorization-request/
+    """
+
+    # create random oauth state param for security and to retrieve user data later
+    oauth_uuid = uuid.uuid4()
+    oauth_uuid_str = str(oauth_uuid)
+
+    # urlsafe b64 encode the uuid for the oauth url
+    oauth_state = (
+        base64.urlsafe_b64encode(oauth_uuid.bytes).rstrip(b"=").decode("utf-8")
+    )
+
+    session: str | None = None
+    if connector == DocumentSource.SLACK:
+        if not DEV_MODE:
+            oauth_url = SlackOAuth.generate_oauth_url(oauth_state)
+        else:
+            oauth_url = SlackOAuth.generate_dev_oauth_url(oauth_state)
+
+        session = SlackOAuth.session_dump_json(
+            email=user.email, redirect_on_success=redirect_on_success
+        )
+    elif connector == DocumentSource.CONFLUENCE:
+        if not DEV_MODE:
+            oauth_url = ConfluenceCloudOAuth.generate_oauth_url(oauth_state)
+        else:
+            oauth_url = ConfluenceCloudOAuth.generate_dev_oauth_url(oauth_state)
+        session = ConfluenceCloudOAuth.session_dump_json(
+            email=user.email, redirect_on_success=redirect_on_success
+        )
+    elif connector == DocumentSource.GOOGLE_DRIVE:
+        if not DEV_MODE:
+            oauth_url = GoogleDriveOAuth.generate_oauth_url(oauth_state)
+        else:
+            oauth_url = GoogleDriveOAuth.generate_dev_oauth_url(oauth_state)
+        session = GoogleDriveOAuth.session_dump_json(
+            email=user.email, redirect_on_success=redirect_on_success
+        )
+    else:
+        oauth_url = None
+
+    if not oauth_url:
+        raise HTTPException(
+            status_code=404,
+            detail=f"The document source type {connector} does not have OAuth implemented",
+        )
+
+    if not session:
+        raise HTTPException(
+            status_code=500,
+            detail=f"The document source type {connector} failed to generate an OAuth session.",
+        )
+
+    r = get_redis_client(tenant_id=tenant_id)
+
+    # store important session state to retrieve when the user is redirected back
+    # 10 min is the max we want an oauth flow to be valid
+    r.set(f"da_oauth:{oauth_uuid_str}", session, ex=600)
+
+    return JSONResponse(content={"url": oauth_url})

backend/ee/onyx/server/oauth/api_router.py

@@ -0,0 +1,3 @@
+from fastapi import APIRouter
+
+router: APIRouter = APIRouter(prefix="/oauth")

backend/ee/onyx/server/oauth/confluence_cloud.py

@@ -0,0 +1,361 @@
+import base64
+import uuid
+from datetime import datetime
+from datetime import timedelta
+from datetime import timezone
+from typing import Any
+from typing import cast
+
+import requests
+from fastapi import Depends
+from fastapi import HTTPException
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel
+from pydantic import ValidationError
+from sqlalchemy.orm import Session
+
+from ee.onyx.configs.app_configs import OAUTH_CONFLUENCE_CLOUD_CLIENT_ID
+from ee.onyx.configs.app_configs import OAUTH_CONFLUENCE_CLOUD_CLIENT_SECRET
+from ee.onyx.server.oauth.api_router import router
+from onyx.auth.users import current_admin_user
+from onyx.configs.app_configs import DEV_MODE
+from onyx.configs.app_configs import WEB_DOMAIN
+from onyx.configs.constants import DocumentSource
+from onyx.connectors.confluence.utils import CONFLUENCE_OAUTH_TOKEN_URL
+from onyx.db.credentials import create_credential
+from onyx.db.credentials import fetch_credential_by_id_for_user
+from onyx.db.credentials import update_credential_json
+from onyx.db.models import User
+from onyx.db.session import get_session
+from onyx.redis.redis_pool import get_redis_client
+from onyx.server.documents.models import CredentialBase
+from onyx.utils.logger import setup_logger
+from shared_configs.contextvars import get_current_tenant_id
+
+logger = setup_logger()
+
+
+class ConfluenceCloudOAuth:
+    # https://developer.atlassian.com/cloud/confluence/oauth-2-3lo-apps/
+
+    class OAuthSession(BaseModel):
+        """Stored in redis to be looked up on callback"""
+
+        email: str
+        redirect_on_success: str | None  # Where to send the user if OAuth flow succeeds
+
+    class TokenResponse(BaseModel):
+        access_token: str
+        expires_in: int
+        token_type: str
+        refresh_token: str
+        scope: str
+
+    class AccessibleResources(BaseModel):
+        id: str
+        name: str
+        url: str
+        scopes: list[str]
+        avatarUrl: str
+
+    CLIENT_ID = OAUTH_CONFLUENCE_CLOUD_CLIENT_ID
+    CLIENT_SECRET = OAUTH_CONFLUENCE_CLOUD_CLIENT_SECRET
+    TOKEN_URL = CONFLUENCE_OAUTH_TOKEN_URL
+
+    ACCESSIBLE_RESOURCE_URL = (
+        "https://api.atlassian.com/oauth/token/accessible-resources"
+    )
+
+    # All read scopes per https://developer.atlassian.com/cloud/confluence/scopes-for-oauth-2-3LO-and-forge-apps/
+    CONFLUENCE_OAUTH_SCOPE = (
+        # classic scope
+        "read:confluence-space.summary%20"
+        "read:confluence-props%20"
+        "read:confluence-content.all%20"
+        "read:confluence-content.summary%20"
+        "read:confluence-content.permission%20"
+        "read:confluence-user%20"
+        "read:confluence-groups%20"
+        "readonly:content.attachment:confluence%20"
+        "search:confluence%20"
+        # granular scope
+        "read:attachment:confluence%20"  # possibly unneeded unless calling v2 attachments api
+        "offline_access"
+    )
+
+    REDIRECT_URI = f"{WEB_DOMAIN}/admin/connectors/confluence/oauth/callback"
+    DEV_REDIRECT_URI = f"https://redirectmeto.com/{REDIRECT_URI}"
+
+    # eventually for Confluence Data Center
+    # oauth_url = (
+    #     f"http://localhost:8090/rest/oauth/v2/authorize?client_id={CONFLUENCE_OAUTH_CLIENT_ID}"
+    #     f"&scope={CONFLUENCE_OAUTH_SCOPE_2}"
+    #     f"&redirect_uri={redirectme_uri}"
+    # )
+
+    @classmethod
+    def generate_oauth_url(cls, state: str) -> str:
+        return cls._generate_oauth_url_helper(cls.REDIRECT_URI, state)
+
+    @classmethod
+    def generate_dev_oauth_url(cls, state: str) -> str:
+        """dev mode workaround for localhost testing
+        - https://www.nango.dev/blog/oauth-redirects-on-localhost-with-https
+        """
+        return cls._generate_oauth_url_helper(cls.DEV_REDIRECT_URI, state)
+
+    @classmethod
+    def _generate_oauth_url_helper(cls, redirect_uri: str, state: str) -> str:
+        # https://developer.atlassian.com/cloud/jira/platform/oauth-2-3lo-apps/#1--direct-the-user-to-the-authorization-url-to-get-an-authorization-code
+
+        url = (
+            "https://auth.atlassian.com/authorize"
+            f"?audience=api.atlassian.com"
+            f"&client_id={cls.CLIENT_ID}"
+            f"&scope={cls.CONFLUENCE_OAUTH_SCOPE}"
+            f"&redirect_uri={redirect_uri}"
+            f"&state={state}"
+            "&response_type=code"
+            "&prompt=consent"
+        )
+        return url
+
+    @classmethod
+    def session_dump_json(cls, email: str, redirect_on_success: str | None) -> str:
+        """Temporary state to store in redis. to be looked up on auth response.
+        Returns a json string.
+        """
+        session = ConfluenceCloudOAuth.OAuthSession(
+            email=email, redirect_on_success=redirect_on_success
+        )
+        return session.model_dump_json()
+
+    @classmethod
+    def parse_session(cls, session_json: str) -> OAuthSession:
+        session = ConfluenceCloudOAuth.OAuthSession.model_validate_json(session_json)
+        return session
+
+    @classmethod
+    def generate_finalize_url(cls, credential_id: int) -> str:
+        return f"{WEB_DOMAIN}/admin/connectors/confluence/oauth/finalize?credential={credential_id}"
+
+
+@router.post("/connector/confluence/callback")
+def confluence_oauth_callback(
+    code: str,
+    state: str,
+    user: User = Depends(current_admin_user),
+    db_session: Session = Depends(get_session),
+    tenant_id: str | None = Depends(get_current_tenant_id),
+) -> JSONResponse:
+    """Handles the backend logic for the frontend page that the user is redirected to
+    after visiting the oauth authorization url."""
+
+    if not ConfluenceCloudOAuth.CLIENT_ID or not ConfluenceCloudOAuth.CLIENT_SECRET:
+        raise HTTPException(
+            status_code=500,
+            detail="Confluence Cloud client ID or client secret is not configured.",
+        )
+
+    r = get_redis_client(tenant_id=tenant_id)
+
+    # recover the state
+    padded_state = state + "=" * (
+        -len(state) % 4
+    )  # Add padding back (Base64 decoding requires padding)
+    uuid_bytes = base64.urlsafe_b64decode(
+        padded_state
+    )  # Decode the Base64 string back to bytes
+
+    # Convert bytes back to a UUID
+    oauth_uuid = uuid.UUID(bytes=uuid_bytes)
+    oauth_uuid_str = str(oauth_uuid)
+
+    r_key = f"da_oauth:{oauth_uuid_str}"
+
+    session_json_bytes = cast(bytes, r.get(r_key))
+    if not session_json_bytes:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Confluence Cloud OAuth failed - OAuth state key not found: key={r_key}",
+        )
+
+    session_json = session_json_bytes.decode("utf-8")
+    try:
+        session = ConfluenceCloudOAuth.parse_session(session_json)
+
+        if not DEV_MODE:
+            redirect_uri = ConfluenceCloudOAuth.REDIRECT_URI
+        else:
+            redirect_uri = ConfluenceCloudOAuth.DEV_REDIRECT_URI
+
+        # Exchange the authorization code for an access token
+        response = requests.post(
+            ConfluenceCloudOAuth.TOKEN_URL,
+            headers={"Content-Type": "application/x-www-form-urlencoded"},
+            data={
+                "client_id": ConfluenceCloudOAuth.CLIENT_ID,
+                "client_secret": ConfluenceCloudOAuth.CLIENT_SECRET,
+                "code": code,
+                "redirect_uri": redirect_uri,
+                "grant_type": "authorization_code",
+            },
+        )
+
+        token_response: ConfluenceCloudOAuth.TokenResponse | None = None
+
+        try:
+            token_response = ConfluenceCloudOAuth.TokenResponse.model_validate_json(
+                response.text
+            )
+        except Exception:
+            raise RuntimeError(
+                "Confluence Cloud OAuth failed during code/token exchange."
+            )
+
+        now = datetime.now(timezone.utc)
+        expires_at = now + timedelta(seconds=token_response.expires_in)
+
+        credential_info = CredentialBase(
+            credential_json={
+                "confluence_access_token": token_response.access_token,
+                "confluence_refresh_token": token_response.refresh_token,
+                "created_at": now.isoformat(),
+                "expires_at": expires_at.isoformat(),
+                "expires_in": token_response.expires_in,
+                "scope": token_response.scope,
+            },
+            admin_public=True,
+            source=DocumentSource.CONFLUENCE,
+            name="Confluence Cloud OAuth",
+        )
+
+        credential = create_credential(credential_info, user, db_session)
+    except Exception as e:
+        return JSONResponse(
+            status_code=500,
+            content={
+                "success": False,
+                "message": f"An error occurred during Confluence Cloud OAuth: {str(e)}",
+            },
+        )
+    finally:
+        r.delete(r_key)
+
+    # return the result
+    return JSONResponse(
+        content={
+            "success": True,
+            "message": "Confluence Cloud OAuth completed successfully.",
+            "finalize_url": ConfluenceCloudOAuth.generate_finalize_url(credential.id),
+            "redirect_on_success": session.redirect_on_success,
+        }
+    )
+
+
+@router.get("/connector/confluence/accessible-resources")
+def confluence_oauth_accessible_resources(
+    credential_id: int,
+    user: User = Depends(current_admin_user),
+    db_session: Session = Depends(get_session),
+    tenant_id: str | None = Depends(get_current_tenant_id),
+) -> JSONResponse:
+    """Atlassian's API is weird and does not supply us with enough info to be in a
+    usable state after authorizing.  All API's require a cloud id. We have to list
+    the accessible resources/sites and let the user choose which site to use."""
+
+    credential = fetch_credential_by_id_for_user(credential_id, user, db_session)
+    if not credential:
+        raise HTTPException(400, f"Credential {credential_id} not found.")
+
+    credential_dict = credential.credential_json
+    access_token = credential_dict["confluence_access_token"]
+
+    try:
+        # Exchange the authorization code for an access token
+        response = requests.get(
+            ConfluenceCloudOAuth.ACCESSIBLE_RESOURCE_URL,
+            headers={
+                "Authorization": f"Bearer {access_token}",
+                "Accept": "application/json",
+            },
+        )
+
+        response.raise_for_status()
+        accessible_resources_data = response.json()
+
+        # Validate the list of AccessibleResources
+        try:
+            accessible_resources = [
+                ConfluenceCloudOAuth.AccessibleResources(**resource)
+                for resource in accessible_resources_data
+            ]
+        except ValidationError as e:
+            raise RuntimeError(f"Failed to parse accessible resources: {e}")
+    except Exception as e:
+        return JSONResponse(
+            status_code=500,
+            content={
+                "success": False,
+                "message": f"An error occurred retrieving Confluence Cloud accessible resources: {str(e)}",
+            },
+        )
+
+    # return the result
+    return JSONResponse(
+        content={
+            "success": True,
+            "message": "Confluence Cloud get accessible resources completed successfully.",
+            "accessible_resources": [
+                resource.model_dump() for resource in accessible_resources
+            ],
+        }
+    )
+
+
+@router.post("/connector/confluence/finalize")
+def confluence_oauth_finalize(
+    credential_id: int,
+    cloud_id: str,
+    cloud_name: str,
+    cloud_url: str,
+    user: User = Depends(current_admin_user),
+    db_session: Session = Depends(get_session),
+    tenant_id: str | None = Depends(get_current_tenant_id),
+) -> JSONResponse:
+    """Saves the info for the selected cloud site to the credential.
+    This is the final step in the confluence oauth flow where after the traditional
+    OAuth process, the user has to select a site to associate with the credentials.
+    After this, the credential is usable."""
+
+    credential = fetch_credential_by_id_for_user(credential_id, user, db_session)
+    if not credential:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Confluence Cloud OAuth failed - credential {credential_id} not found.",
+        )
+
+    new_credential_json: dict[str, Any] = dict(credential.credential_json)
+    new_credential_json["cloud_id"] = cloud_id
+    new_credential_json["cloud_name"] = cloud_name
+    new_credential_json["wiki_base"] = cloud_url
+
+    try:
+        update_credential_json(credential_id, new_credential_json, user, db_session)
+    except Exception as e:
+        return JSONResponse(
+            status_code=500,
+            content={
+                "success": False,
+                "message": f"An error occurred during Confluence Cloud OAuth: {str(e)}",
+            },
+        )
+
+    # return the result
+    return JSONResponse(
+        content={
+            "success": True,
+            "message": "Confluence Cloud OAuth finalized successfully.",
+            "redirect_url": f"{WEB_DOMAIN}/admin/connectors/confluence",
+        }
+    )

backend/ee/onyx/server/oauth/google_drive.py

@@ -0,0 +1,229 @@
+import base64
+import json
+import uuid
+from typing import Any
+from typing import cast
+
+import requests
+from fastapi import Depends
+from fastapi import HTTPException
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel
+from sqlalchemy.orm import Session
+
+from ee.onyx.configs.app_configs import OAUTH_GOOGLE_DRIVE_CLIENT_ID
+from ee.onyx.configs.app_configs import OAUTH_GOOGLE_DRIVE_CLIENT_SECRET
+from ee.onyx.server.oauth.api_router import router
+from onyx.auth.users import current_admin_user
+from onyx.configs.app_configs import DEV_MODE
+from onyx.configs.app_configs import WEB_DOMAIN
+from onyx.configs.constants import DocumentSource
+from onyx.connectors.google_utils.google_auth import get_google_oauth_creds
+from onyx.connectors.google_utils.google_auth import sanitize_oauth_credentials
+from onyx.connectors.google_utils.shared_constants import (
+    DB_CREDENTIALS_AUTHENTICATION_METHOD,
+)
+from onyx.connectors.google_utils.shared_constants import (
+    DB_CREDENTIALS_DICT_TOKEN_KEY,
+)
+from onyx.connectors.google_utils.shared_constants import (
+    DB_CREDENTIALS_PRIMARY_ADMIN_KEY,
+)
+from onyx.connectors.google_utils.shared_constants import (
+    GoogleOAuthAuthenticationMethod,
+)
+from onyx.db.credentials import create_credential
+from onyx.db.models import User
+from onyx.db.session import get_session
+from onyx.redis.redis_pool import get_redis_client
+from onyx.server.documents.models import CredentialBase
+from shared_configs.contextvars import get_current_tenant_id
+
+
+class GoogleDriveOAuth:
+    # https://developers.google.com/identity/protocols/oauth2
+    # https://developers.google.com/identity/protocols/oauth2/web-server
+
+    class OAuthSession(BaseModel):
+        """Stored in redis to be looked up on callback"""
+
+        email: str
+        redirect_on_success: str | None  # Where to send the user if OAuth flow succeeds
+
+    CLIENT_ID = OAUTH_GOOGLE_DRIVE_CLIENT_ID
+    CLIENT_SECRET = OAUTH_GOOGLE_DRIVE_CLIENT_SECRET
+
+    TOKEN_URL = "https://oauth2.googleapis.com/token"
+
+    # SCOPE is per https://docs.danswer.dev/connectors/google-drive
+    # TODO: Merge with or use google_utils.GOOGLE_SCOPES
+    SCOPE = (
+        "https://www.googleapis.com/auth/drive.readonly%20"
+        "https://www.googleapis.com/auth/drive.metadata.readonly%20"
+        "https://www.googleapis.com/auth/admin.directory.user.readonly%20"
+        "https://www.googleapis.com/auth/admin.directory.group.readonly"
+    )
+
+    REDIRECT_URI = f"{WEB_DOMAIN}/admin/connectors/google-drive/oauth/callback"
+    DEV_REDIRECT_URI = f"https://redirectmeto.com/{REDIRECT_URI}"
+
+    @classmethod
+    def generate_oauth_url(cls, state: str) -> str:
+        return cls._generate_oauth_url_helper(cls.REDIRECT_URI, state)
+
+    @classmethod
+    def generate_dev_oauth_url(cls, state: str) -> str:
+        """dev mode workaround for localhost testing
+        - https://www.nango.dev/blog/oauth-redirects-on-localhost-with-https
+        """
+
+        return cls._generate_oauth_url_helper(cls.DEV_REDIRECT_URI, state)
+
+    @classmethod
+    def _generate_oauth_url_helper(cls, redirect_uri: str, state: str) -> str:
+        # without prompt=consent, a refresh token is only issued the first time the user approves
+        url = (
+            f"https://accounts.google.com/o/oauth2/v2/auth"
+            f"?client_id={cls.CLIENT_ID}"
+            f"&redirect_uri={redirect_uri}"
+            "&response_type=code"
+            f"&scope={cls.SCOPE}"
+            "&access_type=offline"
+            f"&state={state}"
+            "&prompt=consent"
+        )
+        return url
+
+    @classmethod
+    def session_dump_json(cls, email: str, redirect_on_success: str | None) -> str:
+        """Temporary state to store in redis. to be looked up on auth response.
+        Returns a json string.
+        """
+        session = GoogleDriveOAuth.OAuthSession(
+            email=email, redirect_on_success=redirect_on_success
+        )
+        return session.model_dump_json()
+
+    @classmethod
+    def parse_session(cls, session_json: str) -> OAuthSession:
+        session = GoogleDriveOAuth.OAuthSession.model_validate_json(session_json)
+        return session
+
+
+@router.post("/connector/google-drive/callback")
+def handle_google_drive_oauth_callback(
+    code: str,
+    state: str,
+    user: User = Depends(current_admin_user),
+    db_session: Session = Depends(get_session),
+    tenant_id: str | None = Depends(get_current_tenant_id),
+) -> JSONResponse:
+    if not GoogleDriveOAuth.CLIENT_ID or not GoogleDriveOAuth.CLIENT_SECRET:
+        raise HTTPException(
+            status_code=500,
+            detail="Google Drive client ID or client secret is not configured.",
+        )
+
+    r = get_redis_client(tenant_id=tenant_id)
+
+    # recover the state
+    padded_state = state + "=" * (
+        -len(state) % 4
+    )  # Add padding back (Base64 decoding requires padding)
+    uuid_bytes = base64.urlsafe_b64decode(
+        padded_state
+    )  # Decode the Base64 string back to bytes
+
+    # Convert bytes back to a UUID
+    oauth_uuid = uuid.UUID(bytes=uuid_bytes)
+    oauth_uuid_str = str(oauth_uuid)
+
+    r_key = f"da_oauth:{oauth_uuid_str}"
+
+    session_json_bytes = cast(bytes, r.get(r_key))
+    if not session_json_bytes:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Google Drive OAuth failed - OAuth state key not found: key={r_key}",
+        )
+
+    session_json = session_json_bytes.decode("utf-8")
+    try:
+        session = GoogleDriveOAuth.parse_session(session_json)
+
+        if not DEV_MODE:
+            redirect_uri = GoogleDriveOAuth.REDIRECT_URI
+        else:
+            redirect_uri = GoogleDriveOAuth.DEV_REDIRECT_URI
+
+        # Exchange the authorization code for an access token
+        response = requests.post(
+            GoogleDriveOAuth.TOKEN_URL,
+            headers={"Content-Type": "application/x-www-form-urlencoded"},
+            data={
+                "client_id": GoogleDriveOAuth.CLIENT_ID,
+                "client_secret": GoogleDriveOAuth.CLIENT_SECRET,
+                "code": code,
+                "redirect_uri": redirect_uri,
+                "grant_type": "authorization_code",
+            },
+        )
+
+        response.raise_for_status()
+
+        authorization_response: dict[str, Any] = response.json()
+
+        # the connector wants us to store the json in its authorized_user_info format
+        # returned from OAuthCredentials.get_authorized_user_info().
+        # So refresh immediately via get_google_oauth_creds with the params filled in
+        # from fields in authorization_response to get the json we need
+        authorized_user_info = {}
+        authorized_user_info["client_id"] = OAUTH_GOOGLE_DRIVE_CLIENT_ID
+        authorized_user_info["client_secret"] = OAUTH_GOOGLE_DRIVE_CLIENT_SECRET
+        authorized_user_info["refresh_token"] = authorization_response["refresh_token"]
+
+        token_json_str = json.dumps(authorized_user_info)
+        oauth_creds = get_google_oauth_creds(
+            token_json_str=token_json_str, source=DocumentSource.GOOGLE_DRIVE
+        )
+        if not oauth_creds:
+            raise RuntimeError("get_google_oauth_creds returned None.")
+
+        # save off the credentials
+        oauth_creds_sanitized_json_str = sanitize_oauth_credentials(oauth_creds)
+
+        credential_dict: dict[str, str] = {}
+        credential_dict[DB_CREDENTIALS_DICT_TOKEN_KEY] = oauth_creds_sanitized_json_str
+        credential_dict[DB_CREDENTIALS_PRIMARY_ADMIN_KEY] = session.email
+        credential_dict[
+            DB_CREDENTIALS_AUTHENTICATION_METHOD
+        ] = GoogleOAuthAuthenticationMethod.OAUTH_INTERACTIVE.value
+
+        credential_info = CredentialBase(
+            credential_json=credential_dict,
+            admin_public=True,
+            source=DocumentSource.GOOGLE_DRIVE,
+            name="OAuth (interactive)",
+        )
+
+        create_credential(credential_info, user, db_session)
+    except Exception as e:
+        return JSONResponse(
+            status_code=500,
+            content={
+                "success": False,
+                "message": f"An error occurred during Google Drive OAuth: {str(e)}",
+            },
+        )
+    finally:
+        r.delete(r_key)
+
+    # return the result
+    return JSONResponse(
+        content={
+            "success": True,
+            "message": "Google Drive OAuth completed successfully.",
+            "finalize_url": None,
+            "redirect_on_success": session.redirect_on_success,
+        }
+    )

backend/ee/onyx/server/oauth/slack.py

@@ -0,0 +1,197 @@
+import base64
+import uuid
+from typing import cast
+
+import requests
+from fastapi import Depends
+from fastapi import HTTPException
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel
+from sqlalchemy.orm import Session
+
+from ee.onyx.configs.app_configs import OAUTH_SLACK_CLIENT_ID
+from ee.onyx.configs.app_configs import OAUTH_SLACK_CLIENT_SECRET
+from ee.onyx.server.oauth.api_router import router
+from onyx.auth.users import current_admin_user
+from onyx.configs.app_configs import DEV_MODE
+from onyx.configs.app_configs import WEB_DOMAIN
+from onyx.configs.constants import DocumentSource
+from onyx.db.credentials import create_credential
+from onyx.db.models import User
+from onyx.db.session import get_session
+from onyx.redis.redis_pool import get_redis_client
+from onyx.server.documents.models import CredentialBase
+from shared_configs.contextvars import get_current_tenant_id
+
+
+class SlackOAuth:
+    # https://knock.app/blog/how-to-authenticate-users-in-slack-using-oauth
+    # Example: https://api.slack.com/authentication/oauth-v2#exchanging
+
+    class OAuthSession(BaseModel):
+        """Stored in redis to be looked up on callback"""
+
+        email: str
+        redirect_on_success: str | None  # Where to send the user if OAuth flow succeeds
+
+    CLIENT_ID = OAUTH_SLACK_CLIENT_ID
+    CLIENT_SECRET = OAUTH_SLACK_CLIENT_SECRET
+
+    TOKEN_URL = "https://slack.com/api/oauth.v2.access"
+
+    # SCOPE is per https://docs.danswer.dev/connectors/slack
+    BOT_SCOPE = (
+        "channels:history,"
+        "channels:read,"
+        "groups:history,"
+        "groups:read,"
+        "channels:join,"
+        "im:history,"
+        "users:read,"
+        "users:read.email,"
+        "usergroups:read"
+    )
+
+    REDIRECT_URI = f"{WEB_DOMAIN}/admin/connectors/slack/oauth/callback"
+    DEV_REDIRECT_URI = f"https://redirectmeto.com/{REDIRECT_URI}"
+
+    @classmethod
+    def generate_oauth_url(cls, state: str) -> str:
+        return cls._generate_oauth_url_helper(cls.REDIRECT_URI, state)
+
+    @classmethod
+    def generate_dev_oauth_url(cls, state: str) -> str:
+        """dev mode workaround for localhost testing
+        - https://www.nango.dev/blog/oauth-redirects-on-localhost-with-https
+        """
+
+        return cls._generate_oauth_url_helper(cls.DEV_REDIRECT_URI, state)
+
+    @classmethod
+    def _generate_oauth_url_helper(cls, redirect_uri: str, state: str) -> str:
+        url = (
+            f"https://slack.com/oauth/v2/authorize"
+            f"?client_id={cls.CLIENT_ID}"
+            f"&redirect_uri={redirect_uri}"
+            f"&scope={cls.BOT_SCOPE}"
+            f"&state={state}"
+        )
+        return url
+
+    @classmethod
+    def session_dump_json(cls, email: str, redirect_on_success: str | None) -> str:
+        """Temporary state to store in redis. to be looked up on auth response.
+        Returns a json string.
+        """
+        session = SlackOAuth.OAuthSession(
+            email=email, redirect_on_success=redirect_on_success
+        )
+        return session.model_dump_json()
+
+    @classmethod
+    def parse_session(cls, session_json: str) -> OAuthSession:
+        session = SlackOAuth.OAuthSession.model_validate_json(session_json)
+        return session
+
+
+@router.post("/connector/slack/callback")
+def handle_slack_oauth_callback(
+    code: str,
+    state: str,
+    user: User = Depends(current_admin_user),
+    db_session: Session = Depends(get_session),
+    tenant_id: str | None = Depends(get_current_tenant_id),
+) -> JSONResponse:
+    if not SlackOAuth.CLIENT_ID or not SlackOAuth.CLIENT_SECRET:
+        raise HTTPException(
+            status_code=500,
+            detail="Slack client ID or client secret is not configured.",
+        )
+
+    r = get_redis_client(tenant_id=tenant_id)
+
+    # recover the state
+    padded_state = state + "=" * (
+        -len(state) % 4
+    )  # Add padding back (Base64 decoding requires padding)
+    uuid_bytes = base64.urlsafe_b64decode(
+        padded_state
+    )  # Decode the Base64 string back to bytes
+
+    # Convert bytes back to a UUID
+    oauth_uuid = uuid.UUID(bytes=uuid_bytes)
+    oauth_uuid_str = str(oauth_uuid)
+
+    r_key = f"da_oauth:{oauth_uuid_str}"
+
+    session_json_bytes = cast(bytes, r.get(r_key))
+    if not session_json_bytes:
+        raise HTTPException(
+            status_code=400,
+            detail=f"Slack OAuth failed - OAuth state key not found: key={r_key}",
+        )
+
+    session_json = session_json_bytes.decode("utf-8")
+    try:
+        session = SlackOAuth.parse_session(session_json)
+
+        if not DEV_MODE:
+            redirect_uri = SlackOAuth.REDIRECT_URI
+        else:
+            redirect_uri = SlackOAuth.DEV_REDIRECT_URI
+
+        # Exchange the authorization code for an access token
+        response = requests.post(
+            SlackOAuth.TOKEN_URL,
+            headers={"Content-Type": "application/x-www-form-urlencoded"},
+            data={
+                "client_id": SlackOAuth.CLIENT_ID,
+                "client_secret": SlackOAuth.CLIENT_SECRET,
+                "code": code,
+                "redirect_uri": redirect_uri,
+            },
+        )
+
+        response_data = response.json()
+
+        if not response_data.get("ok"):
+            raise HTTPException(
+                status_code=400,
+                detail=f"Slack OAuth failed: {response_data.get('error')}",
+            )
+
+        # Extract token and team information
+        access_token: str = response_data.get("access_token")
+        team_id: str = response_data.get("team", {}).get("id")
+        authed_user_id: str = response_data.get("authed_user", {}).get("id")
+
+        credential_info = CredentialBase(
+            credential_json={"slack_bot_token": access_token},
+            admin_public=True,
+            source=DocumentSource.SLACK,
+            name="Slack OAuth",
+        )
+
+        create_credential(credential_info, user, db_session)
+    except Exception as e:
+        return JSONResponse(
+            status_code=500,
+            content={
+                "success": False,
+                "message": f"An error occurred during Slack OAuth: {str(e)}",
+            },
+        )
+    finally:
+        r.delete(r_key)
+
+    # return the result
+    return JSONResponse(
+        content={
+            "success": True,
+            "message": "Slack OAuth completed successfully.",
+            "finalize_url": None,
+            "redirect_on_success": session.redirect_on_success,
+            "team_id": team_id,
+            "authed_user_id": authed_user_id,
+        }
+    )

backend/ee/onyx/server/query_and_chat/chat_backend.py

@@ -31,8 +31,8 @@ from onyx.context.search.models import SavedSearchDoc
 from onyx.db.chat import create_chat_session
 from onyx.db.chat import create_new_chat_message
 from onyx.db.chat import get_or_create_root_message
-from onyx.db.engine import get_session
 from onyx.db.models import User
+from onyx.db.session import get_session
 from onyx.llm.factory import get_llms_for_persona
 from onyx.llm.utils import get_max_input_tokens
 from onyx.natural_language_processing.utils import get_tokenizer

backend/ee/onyx/server/query_and_chat/query_backend.py

@@ -31,10 +31,10 @@ from onyx.context.search.utils import dedupe_documents
 from onyx.context.search.utils import drop_llm_indices
 from onyx.context.search.utils import relevant_sections_to_indices
 from onyx.db.chat import get_prompt_by_id
-from onyx.db.engine import get_session
 from onyx.db.models import Persona
 from onyx.db.models import User
 from onyx.db.persona import get_persona_by_id
+from onyx.db.session import get_session
 from onyx.llm.factory import get_default_llms
 from onyx.llm.factory import get_llms_for_persona
 from onyx.llm.factory import get_main_llm_from_tuple

backend/ee/onyx/server/query_and_chat/token_limit.py

@@ -13,7 +13,6 @@ from sqlalchemy import select
 from sqlalchemy.orm import Session
 
 from onyx.db.api_key import is_api_key_email_address
-from onyx.db.engine import get_session_with_current_tenant
 from onyx.db.models import ChatMessage
 from onyx.db.models import ChatSession
 from onyx.db.models import TokenRateLimit
@@ -21,6 +20,7 @@ from onyx.db.models import TokenRateLimit__UserGroup
 from onyx.db.models import User
 from onyx.db.models import User__UserGroup
 from onyx.db.models import UserGroup
+from onyx.db.session import get_session_with_current_tenant
 from onyx.db.token_limit import fetch_all_user_token_rate_limits
 from onyx.server.query_and_chat.token_limit import _get_cutoff_time
 from onyx.server.query_and_chat.token_limit import _is_rate_limited

backend/ee/onyx/server/query_history/api.py

@@ -29,9 +29,9 @@ from onyx.configs.constants import QueryHistoryType
 from onyx.configs.constants import SessionType
 from onyx.db.chat import get_chat_session_by_id
 from onyx.db.chat import get_chat_sessions_by_user
-from onyx.db.engine import get_session
 from onyx.db.models import ChatSession
 from onyx.db.models import User
+from onyx.db.session import get_session
 from onyx.server.documents.models import PaginatedReturn
 from onyx.server.query_and_chat.models import ChatSessionDetails
 from onyx.server.query_and_chat.models import ChatSessionsResponse
@@ -138,6 +138,7 @@ def get_user_chat_sessions(
                 name=chat.description,
                 persona_id=chat.persona_id,
                 time_created=chat.time_created.isoformat(),
+                time_updated=chat.time_updated.isoformat(),
                 shared_status=chat.shared_status,
                 folder_id=chat.folder_id,
                 current_alternate_model=chat.current_alternate_model,

backend/ee/onyx/server/reporting/usage_export_api.py

@@ -14,8 +14,8 @@ from ee.onyx.db.usage_export import get_usage_report_data
 from ee.onyx.db.usage_export import UsageReportMetadata
 from ee.onyx.server.reporting.usage_export_generation import create_new_usage_report
 from onyx.auth.users import current_admin_user
-from onyx.db.engine import get_session
 from onyx.db.models import User
+from onyx.db.session import get_session
 from onyx.file_store.constants import STANDARD_CHUNK_SIZE
 
 router = APIRouter()

backend/ee/onyx/server/saml.py

@@ -27,9 +27,9 @@ from onyx.auth.users import get_user_manager
 from onyx.configs.app_configs import SESSION_EXPIRE_TIME_SECONDS
 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_session
 from onyx.db.models import User
+from onyx.db.session import get_async_session
+from onyx.db.session import get_session
 from onyx.utils.logger import setup_logger
 
 

backend/ee/onyx/server/seeding.py

@@ -19,11 +19,11 @@ from ee.onyx.server.enterprise_settings.store import (
 )
 from ee.onyx.server.enterprise_settings.store import upload_logo
 from onyx.context.search.enums import RecencyBiasSetting
-from onyx.db.engine import get_session_context_manager
 from onyx.db.llm import update_default_provider
 from onyx.db.llm import upsert_llm_provider
 from onyx.db.models import Tool
 from onyx.db.persona import upsert_persona
+from onyx.db.session import get_session_context_manager
 from onyx.server.features.persona.models import PersonaUpsertRequest
 from onyx.server.manage.llm.models import LLMProviderUpsertRequest
 from onyx.server.settings.models import Settings

backend/ee/onyx/server/tenants/api.py

@@ -41,9 +41,9 @@ from onyx.auth.users import User
 from onyx.configs.app_configs import WEB_DOMAIN
 from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME
 from onyx.db.auth import get_user_count
-from onyx.db.engine import get_session
-from onyx.db.engine import get_session_with_shared_schema
-from onyx.db.engine import get_session_with_tenant
+from onyx.db.session import get_session
+from onyx.db.session import get_session_with_shared_schema
+from onyx.db.session import get_session_with_tenant
 from onyx.db.users import delete_user_from_db
 from onyx.db.users import get_user_by_email
 from onyx.server.manage.models import UserByEmail

backend/ee/onyx/server/tenants/billing.py

@@ -7,6 +7,7 @@ from ee.onyx.configs.app_configs import STRIPE_PRICE_ID
 from ee.onyx.configs.app_configs import STRIPE_SECRET_KEY
 from ee.onyx.server.tenants.access import generate_data_plane_token
 from ee.onyx.server.tenants.models import BillingInformation
+from ee.onyx.server.tenants.models import SubscriptionStatusResponse
 from onyx.configs.app_configs import CONTROL_PLANE_API_BASE_URL
 from onyx.utils.logger import setup_logger
 
@@ -41,7 +42,9 @@ def fetch_tenant_stripe_information(tenant_id: str) -> dict:
     return response.json()
 
 
-def fetch_billing_information(tenant_id: str) -> BillingInformation:
+def fetch_billing_information(
+    tenant_id: str,
+) -> BillingInformation | SubscriptionStatusResponse:
     logger.info("Fetching billing information")
     token = generate_data_plane_token()
     headers = {
@@ -52,8 +55,19 @@ def fetch_billing_information(tenant_id: str) -> BillingInformation:
     params = {"tenant_id": tenant_id}
     response = requests.get(url, headers=headers, params=params)
     response.raise_for_status()
-    billing_info = BillingInformation(**response.json())
-    return billing_info
+
+    response_data = response.json()
+
+    # Check if the response indicates no subscription
+    if (
+        isinstance(response_data, dict)
+        and "subscribed" in response_data
+        and not response_data["subscribed"]
+    ):
+        return SubscriptionStatusResponse(**response_data)
+
+    # Otherwise, parse as BillingInformation
+    return BillingInformation(**response_data)
 
 
 def register_tenant_users(tenant_id: str, number_of_users: int) -> stripe.Subscription:

backend/ee/onyx/server/tenants/provisioning.py

@@ -26,7 +26,6 @@ from onyx.auth.users import exceptions
 from onyx.configs.app_configs import CONTROL_PLANE_API_BASE_URL
 from onyx.configs.app_configs import DEV_MODE
 from onyx.configs.constants import MilestoneRecordType
-from onyx.db.engine import get_session_with_tenant
 from onyx.db.engine import get_sqlalchemy_engine
 from onyx.db.llm import update_default_provider
 from onyx.db.llm import upsert_cloud_embedding_provider
@@ -34,6 +33,7 @@ from onyx.db.llm import upsert_llm_provider
 from onyx.db.models import IndexModelStatus
 from onyx.db.models import SearchSettings
 from onyx.db.models import UserTenantMapping
+from onyx.db.session import get_session_with_tenant
 from onyx.llm.llm_provider_options import ANTHROPIC_MODEL_NAMES
 from onyx.llm.llm_provider_options import ANTHROPIC_PROVIDER_NAME
 from onyx.llm.llm_provider_options import OPEN_AI_MODEL_NAMES
@@ -200,25 +200,6 @@ async def rollback_tenant_provisioning(tenant_id: str) -> None:
 
 
 def configure_default_api_keys(db_session: Session) -> None:
-    if OPENAI_DEFAULT_API_KEY:
-        open_provider = LLMProviderUpsertRequest(
-            name="OpenAI",
-            provider=OPENAI_PROVIDER_NAME,
-            api_key=OPENAI_DEFAULT_API_KEY,
-            default_model_name="gpt-4",
-            fast_default_model_name="gpt-4o-mini",
-            model_names=OPEN_AI_MODEL_NAMES,
-        )
-        try:
-            full_provider = upsert_llm_provider(open_provider, db_session)
-            update_default_provider(full_provider.id, db_session)
-        except Exception as e:
-            logger.error(f"Failed to configure OpenAI provider: {e}")
-    else:
-        logger.error(
-            "OPENAI_DEFAULT_API_KEY not set, skipping OpenAI provider configuration"
-        )
-
     if ANTHROPIC_DEFAULT_API_KEY:
         anthropic_provider = LLMProviderUpsertRequest(
             name="Anthropic",
@@ -227,6 +208,7 @@ def configure_default_api_keys(db_session: Session) -> None:
             default_model_name="claude-3-7-sonnet-20250219",
             fast_default_model_name="claude-3-5-sonnet-20241022",
             model_names=ANTHROPIC_MODEL_NAMES,
+            display_model_names=["claude-3-5-sonnet-20241022"],
         )
         try:
             full_provider = upsert_llm_provider(anthropic_provider, db_session)
@@ -238,6 +220,26 @@ def configure_default_api_keys(db_session: Session) -> None:
             "ANTHROPIC_DEFAULT_API_KEY not set, skipping Anthropic provider configuration"
         )
 
+    if OPENAI_DEFAULT_API_KEY:
+        open_provider = LLMProviderUpsertRequest(
+            name="OpenAI",
+            provider=OPENAI_PROVIDER_NAME,
+            api_key=OPENAI_DEFAULT_API_KEY,
+            default_model_name="gpt-4o",
+            fast_default_model_name="gpt-4o-mini",
+            model_names=OPEN_AI_MODEL_NAMES,
+            display_model_names=["o1", "o3-mini", "gpt-4o", "gpt-4o-mini"],
+        )
+        try:
+            full_provider = upsert_llm_provider(open_provider, db_session)
+            update_default_provider(full_provider.id, db_session)
+        except Exception as e:
+            logger.error(f"Failed to configure OpenAI provider: {e}")
+    else:
+        logger.error(
+            "OPENAI_DEFAULT_API_KEY not set, skipping OpenAI provider configuration"
+        )
+
     if COHERE_DEFAULT_API_KEY:
         cloud_embedding_provider = CloudEmbeddingProviderCreationRequest(
             provider_type=EmbeddingProvider.COHERE,

backend/ee/onyx/server/tenants/user_mapping.py

@@ -4,9 +4,9 @@ from fastapi_users import exceptions
 from sqlalchemy import select
 from sqlalchemy.orm import Session
 
-from onyx.db.engine import get_session_with_tenant
 from onyx.db.engine import get_sqlalchemy_engine
 from onyx.db.models import UserTenantMapping
+from onyx.db.session import get_session_with_tenant
 from shared_configs.configs import MULTI_TENANT
 from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
 

backend/ee/onyx/server/token_rate_limits/api.py

@@ -9,8 +9,8 @@ from ee.onyx.db.token_limit import fetch_user_group_token_rate_limits_for_user
 from ee.onyx.db.token_limit import insert_user_group_token_rate_limit
 from onyx.auth.users import current_admin_user
 from onyx.auth.users import current_curator_or_admin_user
-from onyx.db.engine import get_session
 from onyx.db.models import User
+from onyx.db.session import get_session
 from onyx.db.token_limit import fetch_all_user_token_rate_limits
 from onyx.db.token_limit import insert_user_token_rate_limit
 from onyx.server.query_and_chat.token_limit import any_rate_limit_exists

backend/ee/onyx/server/user_group/api.py

@@ -16,9 +16,9 @@ from ee.onyx.server.user_group.models import UserGroupCreate
 from ee.onyx.server.user_group.models import UserGroupUpdate
 from onyx.auth.users import current_admin_user
 from onyx.auth.users import current_curator_or_admin_user
-from onyx.db.engine import get_session
 from onyx.db.models import User
 from onyx.db.models import UserRole
+from onyx.db.session import get_session
 from onyx.utils.logger import setup_logger
 
 logger = setup_logger()

backend/model_server/encoders.py

@@ -78,7 +78,7 @@ class CloudEmbedding:
         self._closed = False
 
     async def _embed_openai(
-        self, texts: list[str], model: str | None
+        self, texts: list[str], model: str | None, reduced_dimension: int | None
     ) -> list[Embedding]:
         if not model:
             model = DEFAULT_OPENAI_MODEL
@@ -91,7 +91,11 @@ class CloudEmbedding:
         final_embeddings: list[Embedding] = []
         try:
             for text_batch in batch_list(texts, _OPENAI_MAX_INPUT_LEN):
-                response = await client.embeddings.create(input=text_batch, model=model)
+                response = await client.embeddings.create(
+                    input=text_batch,
+                    model=model,
+                    dimensions=reduced_dimension or openai.NOT_GIVEN,
+                )
                 final_embeddings.extend(
                     [embedding.embedding for embedding in response.data]
                 )
@@ -223,9 +227,10 @@ class CloudEmbedding:
         text_type: EmbedTextType,
         model_name: str | None = None,
         deployment_name: str | None = None,
+        reduced_dimension: int | None = None,
     ) -> list[Embedding]:
         if self.provider == EmbeddingProvider.OPENAI:
-            return await self._embed_openai(texts, model_name)
+            return await self._embed_openai(texts, model_name, reduced_dimension)
         elif self.provider == EmbeddingProvider.AZURE:
             return await self._embed_azure(texts, f"azure/{deployment_name}")
         elif self.provider == EmbeddingProvider.LITELLM:
@@ -326,6 +331,7 @@ async def embed_text(
     prefix: str | None,
     api_url: str | None,
     api_version: str | None,
+    reduced_dimension: int | None,
     gpu_type: str = "UNKNOWN",
 ) -> list[Embedding]:
     if not all(texts):
@@ -369,6 +375,7 @@ async def embed_text(
                 model_name=model_name,
                 deployment_name=deployment_name,
                 text_type=text_type,
+                reduced_dimension=reduced_dimension,
             )
 
         if any(embedding is None for embedding in embeddings):
@@ -508,6 +515,7 @@ async def process_embed_request(
             text_type=embed_request.text_type,
             api_url=embed_request.api_url,
             api_version=embed_request.api_version,
+            reduced_dimension=embed_request.reduced_dimension,
             prefix=prefix,
             gpu_type=gpu_type,
         )

backend/onyx/agents/agent_search/basic/graph_builder.py

@@ -78,7 +78,7 @@ def should_continue(state: BasicState) -> str:
 
 
 if __name__ == "__main__":
-    from onyx.db.engine import get_session_context_manager
+    from onyx.db.session import get_session_context_manager
     from onyx.context.search.models import SearchRequest
     from onyx.llm.factory import get_default_llms
     from onyx.agents.agent_search.shared_graph_utils.utils import get_test_config

backend/onyx/agents/agent_search/deep_search/initial/generate_individual_sub_answer/graph_builder.py

@@ -111,7 +111,7 @@ def answer_query_graph_builder() -> StateGraph:
 
 
 if __name__ == "__main__":
-    from onyx.db.engine import get_session_context_manager
+    from onyx.db.session import get_session_context_manager
     from onyx.llm.factory import get_default_llms
     from onyx.context.search.models import SearchRequest
 

backend/onyx/agents/agent_search/deep_search/main/graph_builder.py

@@ -238,7 +238,7 @@ def main_graph_builder(test_mode: bool = False) -> StateGraph:
 if __name__ == "__main__":
     pass
 
-    from onyx.db.engine import get_session_context_manager
+    from onyx.db.session import get_session_context_manager
     from onyx.llm.factory import get_default_llms
     from onyx.context.search.models import SearchRequest
 

backend/onyx/agents/agent_search/deep_search/refinement/consolidate_sub_answers/graph_builder.py

@@ -109,7 +109,7 @@ def answer_refined_query_graph_builder() -> StateGraph:
 
 
 if __name__ == "__main__":
-    from onyx.db.engine import get_session_context_manager
+    from onyx.db.session import get_session_context_manager
     from onyx.llm.factory import get_default_llms
     from onyx.context.search.models import SearchRequest
 

backend/onyx/agents/agent_search/deep_search/shared/expanded_retrieval/graph_builder.py

@@ -131,7 +131,7 @@ def expanded_retrieval_graph_builder() -> StateGraph:
 
 
 if __name__ == "__main__":
-    from onyx.db.engine import get_session_context_manager
+    from onyx.db.session import get_session_context_manager
     from onyx.llm.factory import get_default_llms
     from onyx.context.search.models import SearchRequest
 

backend/onyx/agents/agent_search/deep_search/shared/expanded_retrieval/nodes/rerank_documents.py

@@ -24,8 +24,8 @@ from onyx.context.search.models import InferenceSection
 from onyx.context.search.models import RerankingDetails
 from onyx.context.search.postprocessing.postprocessing import rerank_sections
 from onyx.context.search.postprocessing.postprocessing import should_rerank
-from onyx.db.engine import get_session_context_manager
 from onyx.db.search_settings import get_current_search_settings
+from onyx.db.session import get_session_context_manager
 from onyx.utils.timing import log_function_time
 
 

backend/onyx/agents/agent_search/deep_search/shared/expanded_retrieval/nodes/retrieve_documents.py

@@ -21,7 +21,7 @@ from onyx.agents.agent_search.shared_graph_utils.utils import (
 from onyx.configs.agent_configs import AGENT_MAX_QUERY_RETRIEVAL_RESULTS
 from onyx.configs.agent_configs import AGENT_RETRIEVAL_STATS
 from onyx.context.search.models import InferenceSection
-from onyx.db.engine import get_session_context_manager
+from onyx.db.session import get_session_context_manager
 from onyx.tools.models import SearchQueryInfo
 from onyx.tools.models import SearchToolOverrideKwargs
 from onyx.tools.tool_implementations.search.search_tool import (

backend/onyx/agents/agent_search/run_graph.py

@@ -29,7 +29,7 @@ from onyx.chat.models import ToolResponse
 from onyx.configs.agent_configs import ALLOW_REFINEMENT
 from onyx.configs.agent_configs import INITIAL_SEARCH_DECOMPOSITION_ENABLED
 from onyx.context.search.models import SearchRequest
-from onyx.db.engine import get_session_context_manager
+from onyx.db.session import get_session_context_manager
 from onyx.llm.factory import get_default_llms
 from onyx.tools.tool_runner import ToolCallKickoff
 from onyx.utils.logger import setup_logger

backend/onyx/agents/agent_search/shared_graph_utils/utils.py

@@ -55,9 +55,9 @@ from onyx.context.search.enums import LLMEvaluationType
 from onyx.context.search.models import InferenceSection
 from onyx.context.search.models import RetrievalDetails
 from onyx.context.search.models import SearchRequest
-from onyx.db.engine import get_session_context_manager
 from onyx.db.persona import get_persona_by_id
 from onyx.db.persona import Persona
+from onyx.db.session import get_session_context_manager
 from onyx.llm.chat_llm import LLMRateLimitError
 from onyx.llm.chat_llm import LLMTimeoutError
 from onyx.llm.interfaces import LLM

backend/onyx/auth/users.py

@@ -86,12 +86,12 @@ from onyx.db.auth import get_default_admin_user_emails
 from onyx.db.auth import get_user_count
 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_session_with_tenant
 from onyx.db.models import AccessToken
 from onyx.db.models import OAuthAccount
 from onyx.db.models import User
+from onyx.db.session import get_async_session
+from onyx.db.session import get_async_session_with_tenant
+from onyx.db.session import get_session_with_tenant
 from onyx.db.users import get_user_by_email
 from onyx.redis.redis_pool import get_async_redis_connection
 from onyx.redis.redis_pool import get_redis_client
@@ -411,7 +411,7 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
                 "refresh_token": refresh_token,
             }
 
-            user: User
+            user: User | None = None
 
             try:
                 # Attempt to get user by OAuth account
@@ -420,15 +420,20 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
             except exceptions.UserNotExists:
                 try:
                     # Attempt to get user by email
-                    user = cast(User, await self.user_db.get_by_email(account_email))
+                    user = await self.user_db.get_by_email(account_email)
                     if not associate_by_email:
                         raise exceptions.UserAlreadyExists()
 
-                    user = await self.user_db.add_oauth_account(
-                        user, oauth_account_dict
-                    )
+                    # Make sure user is not None before adding OAuth account
+                    if user is not None:
+                        user = await self.user_db.add_oauth_account(
+                            user, oauth_account_dict
+                        )
+                    else:
+                        # This shouldn't happen since get_by_email would raise UserNotExists
+                        # but adding as a safeguard
+                        raise exceptions.UserNotExists()
 
-                    # If user not found by OAuth account or email, create a new user
                 except exceptions.UserNotExists:
                     password = self.password_helper.generate()
                     user_dict = {
@@ -439,26 +444,36 @@ class UserManager(UUIDIDMixin, BaseUserManager[User, uuid.UUID]):
 
                     user = await self.user_db.create(user_dict)
 
-                    # Explicitly set the Postgres schema for this session to ensure
-                    # OAuth account creation happens in the correct tenant schema
-
-                    # Add OAuth account
-                    await self.user_db.add_oauth_account(user, oauth_account_dict)
-                    await self.on_after_register(user, request)
+                    # Add OAuth account only if user creation was successful
+                    if user is not None:
+                        await self.user_db.add_oauth_account(user, oauth_account_dict)
+                        await self.on_after_register(user, request)
+                    else:
+                        raise HTTPException(
+                            status_code=500, detail="Failed to create user account"
+                        )
 
             else:
-                for existing_oauth_account in user.oauth_accounts:
-                    if (
-                        existing_oauth_account.account_id == account_id
-                        and existing_oauth_account.oauth_name == oauth_name
-                    ):
-                        user = await self.user_db.update_oauth_account(
-                            user,
-                            # NOTE: OAuthAccount DOES implement the OAuthAccountProtocol
-                            # but the type checker doesn't know that :(
-                            existing_oauth_account,  # type: ignore
-                            oauth_account_dict,
-                        )
+                # User exists, update OAuth account if needed
+                if user is not None:  # Add explicit check
+                    for existing_oauth_account in user.oauth_accounts:
+                        if (
+                            existing_oauth_account.account_id == account_id
+                            and existing_oauth_account.oauth_name == oauth_name
+                        ):
+                            user = await self.user_db.update_oauth_account(
+                                user,
+                                # NOTE: OAuthAccount DOES implement the OAuthAccountProtocol
+                                # but the type checker doesn't know that :(
+                                existing_oauth_account,  # type: ignore
+                                oauth_account_dict,
+                            )
+
+            # Ensure user is not None before proceeding
+            if user is None:
+                raise HTTPException(
+                    status_code=500, detail="Failed to authenticate or create user"
+                )
 
             # NOTE: Most IdPs have very short expiry times, and we don't want to force the user to
             # re-authenticate that frequently, so by default this is disabled

backend/onyx/background/celery/apps/beat.py

@@ -13,8 +13,8 @@ from onyx.background.celery.tasks.beat_schedule import CLOUD_BEAT_MULTIPLIER_DEF
 from onyx.configs.constants import ONYX_CLOUD_REDIS_RUNTIME
 from onyx.configs.constants import ONYX_CLOUD_TENANT_ID
 from onyx.configs.constants import POSTGRES_CELERY_BEAT_APP_NAME
-from onyx.db.engine import get_all_tenant_ids
 from onyx.db.engine import SqlEngine
+from onyx.db.tenant import get_all_tenant_ids
 from onyx.redis.redis_pool import get_redis_replica_client
 from onyx.utils.variable_functionality import fetch_versioned_implementation
 from shared_configs.configs import IGNORED_SYNCING_TENANT_LIST

backend/onyx/background/celery/apps/primary.py

@@ -24,10 +24,10 @@ from onyx.configs.constants import CELERY_PRIMARY_WORKER_LOCK_TIMEOUT
 from onyx.configs.constants import OnyxRedisConstants
 from onyx.configs.constants import OnyxRedisLocks
 from onyx.configs.constants import POSTGRES_CELERY_WORKER_PRIMARY_APP_NAME
-from onyx.db.engine import get_session_with_current_tenant
 from onyx.db.engine import SqlEngine
 from onyx.db.index_attempt import get_index_attempt
 from onyx.db.index_attempt import mark_attempt_canceled
+from onyx.db.session import get_session_with_current_tenant
 from onyx.redis.redis_connector_credential_pair import (
     RedisGlobalConnectorCredentialPair,
 )

backend/onyx/background/celery/tasks/connector_deletion/tasks.py

@@ -32,12 +32,12 @@ from onyx.db.connector_credential_pair import get_connector_credential_pair_from
 from onyx.db.connector_credential_pair import get_connector_credential_pairs
 from onyx.db.document import get_document_ids_for_connector_credential_pair
 from onyx.db.document_set import delete_document_set_cc_pair_relationship__no_commit
-from onyx.db.engine import get_session_with_current_tenant
 from onyx.db.enums import ConnectorCredentialPairStatus
 from onyx.db.enums import SyncStatus
 from onyx.db.enums import SyncType
 from onyx.db.index_attempt import delete_index_attempts
 from onyx.db.search_settings import get_all_search_settings
+from onyx.db.session import get_session_with_current_tenant
 from onyx.db.sync_record import cleanup_sync_records
 from onyx.db.sync_record import insert_sync_record
 from onyx.db.sync_record import update_sync_record_status

backend/onyx/background/celery/tasks/doc_permission_syncing/tasks.py

@@ -48,12 +48,12 @@ from onyx.db.connector import mark_cc_pair_as_permissions_synced
 from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id
 from onyx.db.connector_credential_pair import update_connector_credential_pair
 from onyx.db.document import upsert_document_by_connector_credential_pair
-from onyx.db.engine import get_session_with_current_tenant
 from onyx.db.enums import AccessType
 from onyx.db.enums import ConnectorCredentialPairStatus
 from onyx.db.enums import SyncStatus
 from onyx.db.enums import SyncType
 from onyx.db.models import ConnectorCredentialPair
+from onyx.db.session import get_session_with_current_tenant
 from onyx.db.sync_record import insert_sync_record
 from onyx.db.sync_record import update_sync_record_status
 from onyx.db.users import batch_add_ext_perm_user_if_not_exists

backend/onyx/background/celery/tasks/external_group_syncing/tasks.py

@@ -42,12 +42,12 @@ from onyx.connectors.factory import validate_ccpair_for_user
 from onyx.db.connector import mark_cc_pair_as_external_group_synced
 from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id
 from onyx.db.connector_credential_pair import update_connector_credential_pair
-from onyx.db.engine import get_session_with_current_tenant
 from onyx.db.enums import AccessType
 from onyx.db.enums import ConnectorCredentialPairStatus
 from onyx.db.enums import SyncStatus
 from onyx.db.enums import SyncType
 from onyx.db.models import ConnectorCredentialPair
+from onyx.db.session import get_session_with_current_tenant
 from onyx.db.sync_record import insert_sync_record
 from onyx.db.sync_record import update_sync_record_status
 from onyx.redis.redis_connector import RedisConnector
@@ -423,7 +423,7 @@ def connector_external_group_sync_generator_task(
             )
             external_user_groups: list[ExternalUserGroup] = []
             try:
-                external_user_groups = ext_group_sync_func(cc_pair)
+                external_user_groups = ext_group_sync_func(tenant_id, cc_pair)
             except ConnectorValidationError as e:
                 msg = f"Error syncing external groups for {source_type} for cc_pair: {cc_pair_id} {e}"
                 update_connector_credential_pair(

backend/onyx/background/celery/tasks/indexing/tasks.py

@@ -23,9 +23,9 @@ from sqlalchemy.orm import Session
 
 from onyx.background.celery.apps.app_base import task_logger
 from onyx.background.celery.celery_utils import httpx_init_vespa_pool
-from onyx.background.celery.tasks.indexing.utils import _should_index
 from onyx.background.celery.tasks.indexing.utils import get_unfenced_index_attempt_ids
 from onyx.background.celery.tasks.indexing.utils import IndexingCallback
+from onyx.background.celery.tasks.indexing.utils import should_index
 from onyx.background.celery.tasks.indexing.utils import try_creating_indexing_task
 from onyx.background.celery.tasks.indexing.utils import validate_indexing_fences
 from onyx.background.indexing.checkpointing_utils import cleanup_checkpoint
@@ -52,7 +52,6 @@ from onyx.connectors.exceptions import ConnectorValidationError
 from onyx.db.connector import mark_ccpair_with_indexing_trigger
 from onyx.db.connector_credential_pair import fetch_connector_credential_pairs
 from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id
-from onyx.db.engine import get_session_with_current_tenant
 from onyx.db.enums import IndexingMode
 from onyx.db.enums import IndexingStatus
 from onyx.db.index_attempt import get_index_attempt
@@ -61,7 +60,8 @@ from onyx.db.index_attempt import mark_attempt_canceled
 from onyx.db.index_attempt import mark_attempt_failed
 from onyx.db.search_settings import get_active_search_settings_list
 from onyx.db.search_settings import get_current_search_settings
-from onyx.db.swap_index import check_index_swap
+from onyx.db.session import get_session_with_current_tenant
+from onyx.db.swap_index import check_and_perform_index_swap
 from onyx.natural_language_processing.search_nlp_models import EmbeddingModel
 from onyx.natural_language_processing.search_nlp_models import warm_up_bi_encoder
 from onyx.redis.redis_connector import RedisConnector
@@ -406,7 +406,7 @@ def check_for_indexing(self: Task, *, tenant_id: str) -> int | None:
 
         # check for search settings swap
         with get_session_with_current_tenant() as db_session:
-            old_search_settings = check_index_swap(db_session=db_session)
+            old_search_settings = check_and_perform_index_swap(db_session=db_session)
             current_search_settings = get_current_search_settings(db_session)
             # So that the first time users aren't surprised by really slow speed of first
             # batch of documents indexed
@@ -439,6 +439,15 @@ def check_for_indexing(self: Task, *, tenant_id: str) -> int | None:
             with get_session_with_current_tenant() as db_session:
                 search_settings_list = get_active_search_settings_list(db_session)
                 for search_settings_instance in search_settings_list:
+                    # skip non-live search settings that don't have background reindex enabled
+                    # those should just auto-change to live shortly after creation without
+                    # requiring any indexing till that point
+                    if (
+                        not search_settings_instance.status.is_current()
+                        and not search_settings_instance.background_reindex_enabled
+                    ):
+                        continue
+
                     redis_connector_index = redis_connector.new_index(
                         search_settings_instance.id
                     )
@@ -456,23 +465,18 @@ def check_for_indexing(self: Task, *, tenant_id: str) -> int | None:
                         cc_pair.id, search_settings_instance.id, db_session
                     )
 
-                    search_settings_primary = False
-                    if search_settings_instance.id == search_settings_list[0].id:
-                        search_settings_primary = True
-
-                    if not _should_index(
+                    if not should_index(
                         cc_pair=cc_pair,
                         last_index=last_attempt,
                         search_settings_instance=search_settings_instance,
-                        search_settings_primary=search_settings_primary,
                         secondary_index_building=len(search_settings_list) > 1,
                         db_session=db_session,
                     ):
                         continue
 
                     reindex = False
-                    if search_settings_instance.id == search_settings_list[0].id:
-                        # the indexing trigger is only checked and cleared with the primary search settings
+                    if search_settings_instance.status.is_current():
+                        # the indexing trigger is only checked and cleared with the current search settings
                         if cc_pair.indexing_trigger is not None:
                             if cc_pair.indexing_trigger == IndexingMode.REINDEX:
                                 reindex = True

backend/onyx/background/celery/tasks/indexing/utils.py

@@ -23,7 +23,6 @@ from onyx.configs.constants import OnyxCeleryQueues
 from onyx.configs.constants import OnyxCeleryTask
 from onyx.configs.constants import OnyxRedisConstants
 from onyx.db.engine import get_db_current_time
-from onyx.db.engine import get_session_with_current_tenant
 from onyx.db.enums import ConnectorCredentialPairStatus
 from onyx.db.enums import IndexingStatus
 from onyx.db.enums import IndexModelStatus
@@ -35,6 +34,7 @@ from onyx.db.index_attempt import mark_attempt_failed
 from onyx.db.models import ConnectorCredentialPair
 from onyx.db.models import IndexAttempt
 from onyx.db.models import SearchSettings
+from onyx.db.session import get_session_with_current_tenant
 from onyx.indexing.indexing_heartbeat import IndexingHeartbeatInterface
 from onyx.redis.redis_connector import RedisConnector
 from onyx.redis.redis_connector_index import RedisConnectorIndex
@@ -346,11 +346,10 @@ def validate_indexing_fences(
     return
 
 
-def _should_index(
+def should_index(
     cc_pair: ConnectorCredentialPair,
     last_index: IndexAttempt | None,
     search_settings_instance: SearchSettings,
-    search_settings_primary: bool,
     secondary_index_building: bool,
     db_session: Session,
 ) -> bool:
@@ -415,9 +414,9 @@ def _should_index(
     ):
         return False
 
-    if search_settings_primary:
+    if search_settings_instance.status.is_current():
         if cc_pair.indexing_trigger is not None:
-            # if a manual indexing trigger is on the cc pair, honor it for primary search settings
+            # if a manual indexing trigger is on the cc pair, honor it for live search settings
             return True
 
     # if no attempt has ever occurred, we should index regardless of refresh_freq

backend/onyx/background/celery/tasks/llm_model_update/tasks.py

@@ -8,8 +8,8 @@ from onyx.background.celery.apps.app_base import task_logger
 from onyx.configs.app_configs import JOB_TIMEOUT
 from onyx.configs.app_configs import LLM_MODEL_UPDATE_API_URL
 from onyx.configs.constants import OnyxCeleryTask
-from onyx.db.engine import get_session_with_current_tenant
 from onyx.db.models import LLMProvider
+from onyx.db.session import get_session_with_current_tenant
 
 
 def _process_model_list_response(model_list_json: Any) -> list[str]:

backend/onyx/background/celery/tasks/monitoring/tasks.py

@@ -24,10 +24,7 @@ from onyx.configs.constants import ONYX_CLOUD_TENANT_ID
 from onyx.configs.constants import OnyxCeleryQueues
 from onyx.configs.constants import OnyxCeleryTask
 from onyx.configs.constants import OnyxRedisLocks
-from onyx.db.engine import get_all_tenant_ids
 from onyx.db.engine import get_db_current_time
-from onyx.db.engine import get_session_with_current_tenant
-from onyx.db.engine import get_session_with_shared_schema
 from onyx.db.enums import IndexingStatus
 from onyx.db.enums import SyncStatus
 from onyx.db.enums import SyncType
@@ -37,6 +34,9 @@ from onyx.db.models import IndexAttempt
 from onyx.db.models import SyncRecord
 from onyx.db.models import UserGroup
 from onyx.db.search_settings import get_active_search_settings_list
+from onyx.db.session import get_session_with_current_tenant
+from onyx.db.session import get_session_with_shared_schema
+from onyx.db.tenant import get_all_tenant_ids
 from onyx.redis.redis_pool import get_redis_client
 from onyx.redis.redis_pool import redis_lock_dump
 from onyx.utils.telemetry import optional_telemetry

backend/onyx/background/celery/tasks/periodic/tasks.py

@@ -15,7 +15,7 @@ from onyx.background.celery.apps.app_base import task_logger
 from onyx.configs.app_configs import JOB_TIMEOUT
 from onyx.configs.constants import OnyxCeleryTask
 from onyx.configs.constants import PostgresAdvisoryLocks
-from onyx.db.engine import get_session_with_current_tenant
+from onyx.db.session import get_session_with_current_tenant
 
 
 @shared_task(

backend/onyx/background/celery/tasks/pruning/tasks.py

@@ -41,12 +41,12 @@ from onyx.db.connector_credential_pair import get_connector_credential_pair
 from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id
 from onyx.db.connector_credential_pair import get_connector_credential_pairs
 from onyx.db.document import get_documents_for_connector_credential_pair
-from onyx.db.engine import get_session_with_current_tenant
 from onyx.db.enums import ConnectorCredentialPairStatus
 from onyx.db.enums import SyncStatus
 from onyx.db.enums import SyncType
 from onyx.db.models import ConnectorCredentialPair
 from onyx.db.search_settings import get_current_search_settings
+from onyx.db.session import get_session_with_current_tenant
 from onyx.db.sync_record import insert_sync_record
 from onyx.db.sync_record import update_sync_record_status
 from onyx.redis.redis_connector import RedisConnector

backend/onyx/background/celery/tasks/shared/tasks.py

@@ -27,9 +27,9 @@ from onyx.db.document import get_document_connector_count
 from onyx.db.document import mark_document_as_modified
 from onyx.db.document import mark_document_as_synced
 from onyx.db.document_set import fetch_document_sets_for_document
-from onyx.db.engine import get_all_tenant_ids
-from onyx.db.engine import get_session_with_current_tenant
 from onyx.db.search_settings import get_active_search_settings
+from onyx.db.session import get_session_with_current_tenant
+from onyx.db.tenant import get_all_tenant_ids
 from onyx.document_index.factory import get_default_document_index
 from onyx.document_index.interfaces import VespaDocumentFields
 from onyx.httpx.httpx_pool import HttpxPool

backend/onyx/background/celery/tasks/vespa/tasks.py

@@ -35,12 +35,12 @@ from onyx.db.document_set import fetch_document_sets
 from onyx.db.document_set import fetch_document_sets_for_document
 from onyx.db.document_set import get_document_set_by_id
 from onyx.db.document_set import mark_document_set_as_synced
-from onyx.db.engine import get_session_with_current_tenant
 from onyx.db.enums import SyncStatus
 from onyx.db.enums import SyncType
 from onyx.db.models import DocumentSet
 from onyx.db.models import UserGroup
 from onyx.db.search_settings import get_active_search_settings
+from onyx.db.session import get_session_with_current_tenant
 from onyx.db.sync_record import cleanup_sync_records
 from onyx.db.sync_record import insert_sync_record
 from onyx.db.sync_record import update_sync_record_status

backend/onyx/background/error_logging.py

@@ -1,7 +1,7 @@
 from sqlalchemy.exc import IntegrityError
 
 from onyx.db.background_error import create_background_error
-from onyx.db.engine import get_session_with_current_tenant
+from onyx.db.session import get_session_with_current_tenant
 
 
 def emit_background_error(
@@ -11,10 +11,27 @@ def emit_background_error(
     """Currently just saves a row in the background_errors table.
 
     In the future, could create notifications based on the severity."""
-    with get_session_with_current_tenant() as db_session:
-        try:
+    error_message = ""
+
+    # try to write to the db, but handle IntegrityError specifically
+    try:
+        with get_session_with_current_tenant() as db_session:
             create_background_error(db_session, message, cc_pair_id)
-        except IntegrityError as e:
-            # Log an error if the cc_pair_id was deleted or any other exception occurs
-            error_message = f"Failed to create background error: {str(e)}. Original message: {message}"
+    except IntegrityError as e:
+        # Log an error if the cc_pair_id was deleted or any other exception occurs
+        error_message = (
+            f"Failed to create background error: {str(e)}. Original message: {message}"
+        )
+    except Exception:
+        pass
+
+    if not error_message:
+        return
+
+    # if we get here from an IntegrityError, try to write the error message to the db
+    # we need a new session because the first session is now invalid
+    try:
+        with get_session_with_current_tenant() as db_session:
             create_background_error(db_session, error_message, None)
+    except Exception:
+        pass

backend/onyx/background/indexing/job_client.py

@@ -16,7 +16,7 @@ from typing import Optional
 
 from onyx.configs.constants import POSTGRES_CELERY_WORKER_INDEXING_CHILD_APP_NAME
 from onyx.db.engine import SqlEngine
-from onyx.utils.logger import setup_logger
+from onyx.setup import setup_logger
 from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
 from shared_configs.configs import TENANT_ID_PREFIX
 from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR

backend/onyx/background/indexing/run_indexing.py

@@ -30,7 +30,6 @@ from onyx.connectors.models import IndexAttemptMetadata
 from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id
 from onyx.db.connector_credential_pair import get_last_successful_attempt_time
 from onyx.db.connector_credential_pair import update_connector_credential_pair
-from onyx.db.engine import get_session_with_current_tenant
 from onyx.db.enums import ConnectorCredentialPairStatus
 from onyx.db.index_attempt import create_index_attempt_error
 from onyx.db.index_attempt import get_index_attempt
@@ -46,6 +45,7 @@ from onyx.db.models import IndexAttempt
 from onyx.db.models import IndexAttemptError
 from onyx.db.models import IndexingStatus
 from onyx.db.models import IndexModelStatus
+from onyx.db.session import get_session_with_current_tenant
 from onyx.document_index.factory import get_default_document_index
 from onyx.httpx.httpx_pool import HttpxPool
 from onyx.indexing.embedder import DefaultIndexingEmbedder
@@ -93,10 +93,11 @@ def _get_connector_runner(
             runnable_connector.validate_connector_settings()
 
     except Exception as e:
-        logger.exception(f"Unable to instantiate connector due to {e}")
-
+        logger.exception("Unable to instantiate connector.")
         # since we failed to even instantiate the connector, we pause the CCPair since
-        # it will never succeed. Sometimes there are cases where the connector will
+        # it will never succeed
+
+        # Sometimes there are cases where the connector will
         # intermittently fail to initialize in which case we should pass in
         # leave_connector_active=True to allow it to continue.
         # For example, if there is nightly maintenance on a Confluence Server instance,

backend/onyx/chat/process_message.py

@@ -73,7 +73,6 @@ from onyx.db.chat import get_or_create_root_message
 from onyx.db.chat import reserve_message_id
 from onyx.db.chat import translate_db_message_to_chat_message_detail
 from onyx.db.chat import translate_db_search_doc_to_server_search_doc
-from onyx.db.engine import get_session_context_manager
 from onyx.db.milestone import check_multi_assistant_milestone
 from onyx.db.milestone import create_milestone_if_not_exists
 from onyx.db.milestone import update_user_assistant_milestone
@@ -82,6 +81,7 @@ from onyx.db.models import ToolCall
 from onyx.db.models import User
 from onyx.db.persona import get_persona_by_id
 from onyx.db.search_settings import get_current_search_settings
+from onyx.db.session import get_session_context_manager
 from onyx.document_index.factory import get_default_document_index
 from onyx.file_store.models import ChatFileType
 from onyx.file_store.models import FileDescriptor

backend/onyx/connectors/confluence/connector.py

@@ -11,17 +11,20 @@ from onyx.configs.app_configs import CONFLUENCE_TIMEZONE_OFFSET
 from onyx.configs.app_configs import CONTINUE_ON_CONNECTOR_FAILURE
 from onyx.configs.app_configs import INDEX_BATCH_SIZE
 from onyx.configs.constants import DocumentSource
-from onyx.connectors.confluence.onyx_confluence import build_confluence_client
+from onyx.connectors.confluence.onyx_confluence import attachment_to_content
+from onyx.connectors.confluence.onyx_confluence import (
+    extract_text_from_confluence_html,
+)
 from onyx.connectors.confluence.onyx_confluence import OnyxConfluence
-from onyx.connectors.confluence.utils import attachment_to_content
 from onyx.connectors.confluence.utils import build_confluence_document_id
 from onyx.connectors.confluence.utils import datetime_from_string
-from onyx.connectors.confluence.utils import extract_text_from_confluence_html
 from onyx.connectors.confluence.utils import validate_attachment_filetype
 from onyx.connectors.exceptions import ConnectorValidationError
 from onyx.connectors.exceptions import CredentialExpiredError
 from onyx.connectors.exceptions import InsufficientPermissionsError
 from onyx.connectors.exceptions import UnexpectedError
+from onyx.connectors.interfaces import CredentialsConnector
+from onyx.connectors.interfaces import CredentialsProviderInterface
 from onyx.connectors.interfaces import GenerateDocumentsOutput
 from onyx.connectors.interfaces import GenerateSlimDocumentOutput
 from onyx.connectors.interfaces import LoadConnector
@@ -83,7 +86,9 @@ _FULL_EXTENSION_FILTER_STRING = "".join(
 )
 
 
-class ConfluenceConnector(LoadConnector, PollConnector, SlimConnector):
+class ConfluenceConnector(
+    LoadConnector, PollConnector, SlimConnector, CredentialsConnector
+):
     def __init__(
         self,
         wiki_base: str,
@@ -102,7 +107,6 @@ class ConfluenceConnector(LoadConnector, PollConnector, SlimConnector):
     ) -> None:
         self.batch_size = batch_size
         self.continue_on_failure = continue_on_failure
-        self._confluence_client: OnyxConfluence | None = None
         self.is_cloud = is_cloud
 
         # Remove trailing slash from wiki_base if present
@@ -137,6 +141,19 @@ class ConfluenceConnector(LoadConnector, PollConnector, SlimConnector):
             self.cql_label_filter = f" and label not in ({comma_separated_labels})"
 
         self.timezone: timezone = timezone(offset=timedelta(hours=timezone_offset))
+        self.credentials_provider: CredentialsProviderInterface | None = None
+
+        self.probe_kwargs = {
+            "max_backoff_retries": 6,
+            "max_backoff_seconds": 10,
+        }
+
+        self.final_kwargs = {
+            "max_backoff_retries": 10,
+            "max_backoff_seconds": 60,
+        }
+
+        self._confluence_client: OnyxConfluence | None = None
 
     @property
     def confluence_client(self) -> OnyxConfluence:
@@ -144,15 +161,22 @@ class ConfluenceConnector(LoadConnector, PollConnector, SlimConnector):
             raise ConnectorMissingCredentialError("Confluence")
         return self._confluence_client
 
-    def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:
-        # see https://github.com/atlassian-api/atlassian-python-api/blob/master/atlassian/rest_client.py
-        # for a list of other hidden constructor args
-        self._confluence_client = build_confluence_client(
-            credentials=credentials,
-            is_cloud=self.is_cloud,
-            wiki_base=self.wiki_base,
+    def set_credentials_provider(
+        self, credentials_provider: CredentialsProviderInterface
+    ) -> None:
+        self.credentials_provider = credentials_provider
+
+        # raises exception if there's a problem
+        confluence_client = OnyxConfluence(
+            self.is_cloud, self.wiki_base, credentials_provider
         )
-        return None
+        confluence_client._probe_connection(**self.probe_kwargs)
+        confluence_client._initialize_connection(**self.final_kwargs)
+
+        self._confluence_client = confluence_client
+
+    def load_credentials(self, credentials: dict[str, Any]) -> dict[str, Any] | None:
+        raise NotImplementedError("Use set_credentials_provider with this connector.")
 
     def _construct_page_query(
         self,
@@ -202,12 +226,17 @@ class ConfluenceConnector(LoadConnector, PollConnector, SlimConnector):
         return comment_string
 
     def _convert_object_to_document(
-        self, confluence_object: dict[str, Any]
+        self,
+        confluence_object: dict[str, Any],
+        parent_content_id: str | None = None,
     ) -> Document | None:
         """
         Takes in a confluence object, extracts all metadata, and converts it into a document.
         If its a page, it extracts the text, adds the comments for the document text.
         If its an attachment, it just downloads the attachment and converts that into a document.
+
+        parent_content_id: if the object is an attachment, specifies the content id that
+        the attachment is attached to
         """
         # The url and the id are the same
         object_url = build_confluence_document_id(
@@ -226,7 +255,9 @@ class ConfluenceConnector(LoadConnector, PollConnector, SlimConnector):
             object_text += self._get_comment_string_for_page_id(confluence_object["id"])
         elif confluence_object["type"] == "attachment":
             object_text = attachment_to_content(
-                confluence_client=self.confluence_client, attachment=confluence_object
+                confluence_client=self.confluence_client,
+                attachment=confluence_object,
+                parent_content_id=parent_content_id,
             )
 
         if object_text is None:
@@ -302,7 +333,7 @@ class ConfluenceConnector(LoadConnector, PollConnector, SlimConnector):
                 cql=attachment_query,
                 expand=",".join(_ATTACHMENT_EXPANSION_FIELDS),
             ):
-                doc = self._convert_object_to_document(attachment)
+                doc = self._convert_object_to_document(attachment, confluence_page_id)
                 if doc is not None:
                     doc_batch.append(doc)
                 if len(doc_batch) >= self.batch_size:

backend/onyx/connectors/confluence/onyx_confluence.py

@@ -1,19 +1,37 @@
-import math
+import io
+import json
 import time
 from collections.abc import Callable
 from collections.abc import Iterator
+from datetime import datetime
+from datetime import timedelta
+from datetime import timezone
 from typing import Any
 from typing import cast
 from typing import TypeVar
 from urllib.parse import quote
 
+import bs4
 from atlassian import Confluence  # type:ignore
 from pydantic import BaseModel
+from redis import Redis
 from requests import HTTPError
 
+from ee.onyx.configs.app_configs import OAUTH_CONFLUENCE_CLOUD_CLIENT_ID
+from ee.onyx.configs.app_configs import OAUTH_CONFLUENCE_CLOUD_CLIENT_SECRET
+from onyx.configs.app_configs import (
+    CONFLUENCE_CONNECTOR_ATTACHMENT_CHAR_COUNT_THRESHOLD,
+)
+from onyx.configs.app_configs import CONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD
+from onyx.connectors.confluence.utils import _handle_http_error
+from onyx.connectors.confluence.utils import confluence_refresh_tokens
 from onyx.connectors.confluence.utils import get_start_param_from_url
 from onyx.connectors.confluence.utils import update_param_in_path
-from onyx.connectors.exceptions import ConnectorValidationError
+from onyx.connectors.confluence.utils import validate_attachment_filetype
+from onyx.connectors.interfaces import CredentialsProviderInterface
+from onyx.file_processing.extract_file_text import extract_file_text
+from onyx.file_processing.html_utils import format_document_soup
+from onyx.redis.redis_pool import get_redis_client
 from onyx.utils.logger import setup_logger
 
 logger = setup_logger()
@@ -22,12 +40,14 @@ logger = setup_logger()
 F = TypeVar("F", bound=Callable[..., Any])
 
 
-RATE_LIMIT_MESSAGE_LOWERCASE = "Rate limit exceeded".lower()
-
 # https://jira.atlassian.com/browse/CONFCLOUD-76433
 _PROBLEMATIC_EXPANSIONS = "body.storage.value"
 _REPLACEMENT_EXPANSIONS = "body.view.value"
 
+_USER_NOT_FOUND = "Unknown Confluence User"
+_USER_ID_TO_DISPLAY_NAME_CACHE: dict[str, str | None] = {}
+_USER_EMAIL_CACHE: dict[str, str | None] = {}
+
 
 class ConfluenceRateLimitError(Exception):
     pass
@@ -43,124 +63,349 @@ class ConfluenceUser(BaseModel):
     type: str
 
 
-def _handle_http_error(e: HTTPError, attempt: int) -> int:
-    MIN_DELAY = 2
-    MAX_DELAY = 60
-    STARTING_DELAY = 5
-    BACKOFF = 2
-
-    # Check if the response or headers are None to avoid potential AttributeError
-    if e.response is None or e.response.headers is None:
-        logger.warning("HTTPError with `None` as response or as headers")
-        raise e
-
-    if (
-        e.response.status_code != 429
-        and RATE_LIMIT_MESSAGE_LOWERCASE not in e.response.text.lower()
-    ):
-        raise e
-
-    retry_after = None
-
-    retry_after_header = e.response.headers.get("Retry-After")
-    if retry_after_header is not None:
-        try:
-            retry_after = int(retry_after_header)
-            if retry_after > MAX_DELAY:
-                logger.warning(
-                    f"Clamping retry_after from {retry_after} to {MAX_DELAY} seconds..."
-                )
-                retry_after = MAX_DELAY
-            if retry_after < MIN_DELAY:
-                retry_after = MIN_DELAY
-        except ValueError:
-            pass
-
-    if retry_after is not None:
-        logger.warning(
-            f"Rate limiting with retry header. Retrying after {retry_after} seconds..."
-        )
-        delay = retry_after
-    else:
-        logger.warning(
-            "Rate limiting without retry header. Retrying with exponential backoff..."
-        )
-        delay = min(STARTING_DELAY * (BACKOFF**attempt), MAX_DELAY)
-
-    delay_until = math.ceil(time.monotonic() + delay)
-    return delay_until
-
-
-# https://developer.atlassian.com/cloud/confluence/rate-limiting/
-# this uses the native rate limiting option provided by the
-# confluence client and otherwise applies a simpler set of error handling
-def handle_confluence_rate_limit(confluence_call: F) -> F:
-    def wrapped_call(*args: list[Any], **kwargs: Any) -> Any:
-        MAX_RETRIES = 5
-
-        TIMEOUT = 600
-        timeout_at = time.monotonic() + TIMEOUT
-
-        for attempt in range(MAX_RETRIES):
-            if time.monotonic() > timeout_at:
-                raise TimeoutError(
-                    f"Confluence call attempts took longer than {TIMEOUT} seconds."
-                )
-
-            try:
-                # we're relying more on the client to rate limit itself
-                # and applying our own retries in a more specific set of circumstances
-                return confluence_call(*args, **kwargs)
-            except HTTPError as e:
-                delay_until = _handle_http_error(e, attempt)
-                logger.warning(
-                    f"HTTPError in confluence call. "
-                    f"Retrying in {delay_until} seconds..."
-                )
-                while time.monotonic() < delay_until:
-                    # in the future, check a signal here to exit
-                    time.sleep(1)
-            except AttributeError as e:
-                # Some error within the Confluence library, unclear why it fails.
-                # Users reported it to be intermittent, so just retry
-                if attempt == MAX_RETRIES - 1:
-                    raise e
-
-                logger.exception(
-                    "Confluence Client raised an AttributeError. Retrying..."
-                )
-                time.sleep(5)
-
-    return cast(F, wrapped_call)
-
-
 _DEFAULT_PAGINATION_LIMIT = 1000
 _MINIMUM_PAGINATION_LIMIT = 50
 
 
-class OnyxConfluence(Confluence):
+class OnyxConfluence:
     """
-    This is a custom Confluence class that overrides the default Confluence class to add a custom CQL method.
+    This is a custom Confluence class that:
+
+    A. overrides the default Confluence class to add a custom CQL method.
+    B.
     This is necessary because the default Confluence class does not properly support cql expansions.
     All methods are automatically wrapped with handle_confluence_rate_limit.
     """
 
-    def __init__(self, url: str, *args: Any, **kwargs: Any) -> None:
-        super(OnyxConfluence, self).__init__(url, *args, **kwargs)
-        self._wrap_methods()
+    CREDENTIAL_PREFIX = "connector:confluence:credential"
+    CREDENTIAL_TTL = 300  # 5 min
 
-    def _wrap_methods(self) -> None:
+    def __init__(
+        self,
+        is_cloud: bool,
+        url: str,
+        credentials_provider: CredentialsProviderInterface,
+    ) -> None:
+        self._is_cloud = is_cloud
+        self._url = url.rstrip("/")
+        self._credentials_provider = credentials_provider
+
+        self.redis_client: Redis | None = None
+        self.static_credentials: dict[str, Any] | None = None
+        if self._credentials_provider.is_dynamic():
+            self.redis_client = get_redis_client(
+                tenant_id=credentials_provider.get_tenant_id()
+            )
+        else:
+            self.static_credentials = self._credentials_provider.get_credentials()
+
+        self._confluence = Confluence(url)
+        self.credential_key: str = (
+            self.CREDENTIAL_PREFIX
+            + f":credential_{self._credentials_provider.get_provider_key()}"
+        )
+
+        self._kwargs: Any = None
+
+        self.shared_base_kwargs = {
+            "api_version": "cloud" if is_cloud else "latest",
+            "backoff_and_retry": True,
+            "cloud": is_cloud,
+        }
+
+    def _renew_credentials(self) -> tuple[dict[str, Any], bool]:
+        """credential_json - the current json credentials
+        Returns a tuple
+        1. The up to date credentials
+        2. True if the credentials were updated
+
+        This method is intended to be used within a distributed lock.
+        Lock, call this, update credentials if the tokens were refreshed, then release
         """
-        For each attribute that is callable (i.e., a method) and doesn't start with an underscore,
-        wrap it with handle_confluence_rate_limit.
-        """
-        for attr_name in dir(self):
-            if callable(getattr(self, attr_name)) and not attr_name.startswith("_"):
-                setattr(
-                    self,
-                    attr_name,
-                    handle_confluence_rate_limit(getattr(self, attr_name)),
+        # static credentials are preloaded, so no locking/redis required
+        if self.static_credentials:
+            return self.static_credentials, False
+
+        if not self.redis_client:
+            raise RuntimeError("self.redis_client is None")
+
+        # dynamic credentials need locking
+        # check redis first, then fallback to the DB
+        credential_raw = self.redis_client.get(self.credential_key)
+        if credential_raw is not None:
+            credential_bytes = cast(bytes, credential_raw)
+            credential_str = credential_bytes.decode("utf-8")
+            credential_json: dict[str, Any] = json.loads(credential_str)
+        else:
+            credential_json = self._credentials_provider.get_credentials()
+
+        if "confluence_refresh_token" not in credential_json:
+            # static credentials ... cache them permanently and return
+            self.static_credentials = credential_json
+            return credential_json, False
+
+        # check if we should refresh tokens. we're deciding to refresh halfway
+        # to expiration
+        now = datetime.now(timezone.utc)
+        created_at = datetime.fromisoformat(credential_json["created_at"])
+        expires_in: int = credential_json["expires_in"]
+        renew_at = created_at + timedelta(seconds=expires_in // 2)
+        if now <= renew_at:
+            # cached/current credentials are reasonably up to date
+            return credential_json, False
+
+        # we need to refresh
+        logger.info("Renewing Confluence Cloud credentials...")
+        new_credentials = confluence_refresh_tokens(
+            OAUTH_CONFLUENCE_CLOUD_CLIENT_ID,
+            OAUTH_CONFLUENCE_CLOUD_CLIENT_SECRET,
+            credential_json["cloud_id"],
+            credential_json["confluence_refresh_token"],
+        )
+
+        # store the new credentials to redis and to the db thru the provider
+        # redis: we use a 5 min TTL because we are given a 10 minute grace period
+        # when keys are rotated. it's easier to expire the cached credentials
+        # reasonably frequently rather than trying to handle strong synchronization
+        # between the db and redis everywhere the credentials might be updated
+        new_credential_str = json.dumps(new_credentials)
+        self.redis_client.set(
+            self.credential_key, new_credential_str, nx=True, ex=self.CREDENTIAL_TTL
+        )
+        self._credentials_provider.set_credentials(new_credentials)
+
+        return new_credentials, True
+
+    @staticmethod
+    def _make_oauth2_dict(credentials: dict[str, Any]) -> dict[str, Any]:
+        oauth2_dict: dict[str, Any] = {}
+        if "confluence_refresh_token" in credentials:
+            oauth2_dict["client_id"] = OAUTH_CONFLUENCE_CLOUD_CLIENT_ID
+            oauth2_dict["token"] = {}
+            oauth2_dict["token"]["access_token"] = credentials[
+                "confluence_access_token"
+            ]
+        return oauth2_dict
+
+    def _probe_connection(
+        self,
+        **kwargs: Any,
+    ) -> None:
+        merged_kwargs = {**self.shared_base_kwargs, **kwargs}
+
+        with self._credentials_provider:
+            credentials, _ = self._renew_credentials()
+
+            # probe connection with direct client, no retries
+            if "confluence_refresh_token" in credentials:
+                logger.info("Probing Confluence with OAuth Access Token.")
+
+                oauth2_dict: dict[str, Any] = OnyxConfluence._make_oauth2_dict(
+                    credentials
                 )
+                url = (
+                    f"https://api.atlassian.com/ex/confluence/{credentials['cloud_id']}"
+                )
+                confluence_client_with_minimal_retries = Confluence(
+                    url=url, oauth2=oauth2_dict, **merged_kwargs
+                )
+            else:
+                logger.info("Probing Confluence with Personal Access Token.")
+                url = self._url
+                if self._is_cloud:
+                    confluence_client_with_minimal_retries = Confluence(
+                        url=url,
+                        username=credentials["confluence_username"],
+                        password=credentials["confluence_access_token"],
+                        **merged_kwargs,
+                    )
+                else:
+                    confluence_client_with_minimal_retries = Confluence(
+                        url=url,
+                        token=credentials["confluence_access_token"],
+                        **merged_kwargs,
+                    )
+
+            spaces = confluence_client_with_minimal_retries.get_all_spaces(limit=1)
+
+            # uncomment the following for testing
+            # the following is an attempt to retrieve the user's timezone
+            # Unfornately, all data is returned in UTC regardless of the user's time zone
+            # even tho CQL parses incoming times based on the user's time zone
+            # space_key = spaces["results"][0]["key"]
+            # space_details = confluence_client_with_minimal_retries.cql(f"space.key={space_key}+AND+type=space")
+
+            if not spaces:
+                raise RuntimeError(
+                    f"No spaces found at {url}! "
+                    "Check your credentials and wiki_base and make sure "
+                    "is_cloud is set correctly."
+                )
+
+            logger.info("Confluence probe succeeded.")
+
+    def _initialize_connection(
+        self,
+        **kwargs: Any,
+    ) -> None:
+        """Called externally to init the connection in a thread safe manner."""
+        merged_kwargs = {**self.shared_base_kwargs, **kwargs}
+        with self._credentials_provider:
+            credentials, _ = self._renew_credentials()
+            self._confluence = self._initialize_connection_helper(
+                credentials, **merged_kwargs
+            )
+            self._kwargs = merged_kwargs
+
+    def _initialize_connection_helper(
+        self,
+        credentials: dict[str, Any],
+        **kwargs: Any,
+    ) -> Confluence:
+        """Called internally to init the connection. Distributed locking
+        to prevent multiple threads from modifying the credentials
+        must be handled around this function."""
+
+        confluence = None
+
+        # probe connection with direct client, no retries
+        if "confluence_refresh_token" in credentials:
+            logger.info("Connecting to Confluence Cloud with OAuth Access Token.")
+
+            oauth2_dict: dict[str, Any] = OnyxConfluence._make_oauth2_dict(credentials)
+            url = f"https://api.atlassian.com/ex/confluence/{credentials['cloud_id']}"
+            confluence = Confluence(url=url, oauth2=oauth2_dict, **kwargs)
+        else:
+            logger.info("Connecting to Confluence with Personal Access Token.")
+            if self._is_cloud:
+                confluence = Confluence(
+                    url=self._url,
+                    username=credentials["confluence_username"],
+                    password=credentials["confluence_access_token"],
+                    **kwargs,
+                )
+            else:
+                confluence = Confluence(
+                    url=self._url,
+                    token=credentials["confluence_access_token"],
+                    **kwargs,
+                )
+
+        return confluence
+
+    # https://developer.atlassian.com/cloud/confluence/rate-limiting/
+    # this uses the native rate limiting option provided by the
+    # confluence client and otherwise applies a simpler set of error handling
+    def _make_rate_limited_confluence_method(
+        self, name: str, credential_provider: CredentialsProviderInterface | None
+    ) -> Callable[..., Any]:
+        def wrapped_call(*args: list[Any], **kwargs: Any) -> Any:
+            MAX_RETRIES = 5
+
+            TIMEOUT = 600
+            timeout_at = time.monotonic() + TIMEOUT
+
+            for attempt in range(MAX_RETRIES):
+                if time.monotonic() > timeout_at:
+                    raise TimeoutError(
+                        f"Confluence call attempts took longer than {TIMEOUT} seconds."
+                    )
+
+                # we're relying more on the client to rate limit itself
+                # and applying our own retries in a more specific set of circumstances
+                try:
+                    if credential_provider:
+                        with credential_provider:
+                            credentials, renewed = self._renew_credentials()
+                            if renewed:
+                                self._confluence = self._initialize_connection_helper(
+                                    credentials, **self._kwargs
+                                )
+                            attr = getattr(self._confluence, name, None)
+                            if attr is None:
+                                # The underlying Confluence client doesn't have this attribute
+                                raise AttributeError(
+                                    f"'{type(self).__name__}' object has no attribute '{name}'"
+                                )
+
+                            return attr(*args, **kwargs)
+                    else:
+                        attr = getattr(self._confluence, name, None)
+                        if attr is None:
+                            # The underlying Confluence client doesn't have this attribute
+                            raise AttributeError(
+                                f"'{type(self).__name__}' object has no attribute '{name}'"
+                            )
+
+                        return attr(*args, **kwargs)
+
+                except HTTPError as e:
+                    delay_until = _handle_http_error(e, attempt)
+                    logger.warning(
+                        f"HTTPError in confluence call. "
+                        f"Retrying in {delay_until} seconds..."
+                    )
+                    while time.monotonic() < delay_until:
+                        # in the future, check a signal here to exit
+                        time.sleep(1)
+                except AttributeError as e:
+                    # Some error within the Confluence library, unclear why it fails.
+                    # Users reported it to be intermittent, so just retry
+                    if attempt == MAX_RETRIES - 1:
+                        raise e
+
+                    logger.exception(
+                        "Confluence Client raised an AttributeError. Retrying..."
+                    )
+                    time.sleep(5)
+
+        return wrapped_call
+
+    # def _wrap_methods(self) -> None:
+    #     """
+    #     For each attribute that is callable (i.e., a method) and doesn't start with an underscore,
+    #     wrap it with handle_confluence_rate_limit.
+    #     """
+    #     for attr_name in dir(self):
+    #         if callable(getattr(self, attr_name)) and not attr_name.startswith("_"):
+    #             setattr(
+    #                 self,
+    #                 attr_name,
+    #                 handle_confluence_rate_limit(getattr(self, attr_name)),
+    #             )
+
+    # def _ensure_token_valid(self) -> None:
+    #     if self._token_is_expired():
+    #         self._refresh_token()
+    #         # Re-init the Confluence client with the originally stored args
+    #         self._confluence = Confluence(self._url, *self._args, **self._kwargs)
+
+    def __getattr__(self, name: str) -> Any:
+        """Dynamically intercept attribute/method access."""
+        attr = getattr(self._confluence, name, None)
+        if attr is None:
+            # The underlying Confluence client doesn't have this attribute
+            raise AttributeError(
+                f"'{type(self).__name__}' object has no attribute '{name}'"
+            )
+
+        # If it's not a method, just return it after ensuring token validity
+        if not callable(attr):
+            return attr
+
+        # skip methods that start with "_"
+        if name.startswith("_"):
+            return attr
+
+        # wrap the method with our retry handler
+        rate_limited_method: Callable[
+            ..., Any
+        ] = self._make_rate_limited_confluence_method(name, self._credentials_provider)
+
+        def wrapped_method(*args: Any, **kwargs: Any) -> Any:
+            return rate_limited_method(*args, **kwargs)
+
+        return wrapped_method
 
     def _paginate_url(
         self, url_suffix: str, limit: int | None = None, auto_paginate: bool = False
@@ -507,63 +752,212 @@ class OnyxConfluence(Confluence):
         return response
 
 
-def _validate_connector_configuration(
-    credentials: dict[str, Any],
-    is_cloud: bool,
-    wiki_base: str,
-) -> None:
-    # test connection with direct client, no retries
-    confluence_client_with_minimal_retries = Confluence(
-        api_version="cloud" if is_cloud else "latest",
-        url=wiki_base.rstrip("/"),
-        username=credentials["confluence_username"] if is_cloud else None,
-        password=credentials["confluence_access_token"] if is_cloud else None,
-        token=credentials["confluence_access_token"] if not is_cloud else None,
-        backoff_and_retry=True,
-        max_backoff_retries=6,
-        max_backoff_seconds=10,
+def get_user_email_from_username__server(
+    confluence_client: OnyxConfluence, user_name: str
+) -> str | None:
+    global _USER_EMAIL_CACHE
+    if _USER_EMAIL_CACHE.get(user_name) is None:
+        try:
+            response = confluence_client.get_mobile_parameters(user_name)
+            email = response.get("email")
+        except Exception:
+            logger.warning(f"failed to get confluence email for {user_name}")
+            # For now, we'll just return None and log a warning. This means
+            # we will keep retrying to get the email every group sync.
+            email = None
+            # We may want to just return a string that indicates failure so we dont
+            # keep retrying
+            # email = f"FAILED TO GET CONFLUENCE EMAIL FOR {user_name}"
+        _USER_EMAIL_CACHE[user_name] = email
+    return _USER_EMAIL_CACHE[user_name]
+
+
+def _get_user(confluence_client: OnyxConfluence, user_id: str) -> str:
+    """Get Confluence Display Name based on the account-id or userkey value
+
+    Args:
+        user_id (str): The user id (i.e: the account-id or userkey)
+        confluence_client (Confluence): The Confluence Client
+
+    Returns:
+        str: The User Display Name. 'Unknown User' if the user is deactivated or not found
+    """
+    global _USER_ID_TO_DISPLAY_NAME_CACHE
+    if _USER_ID_TO_DISPLAY_NAME_CACHE.get(user_id) is None:
+        try:
+            result = confluence_client.get_user_details_by_userkey(user_id)
+            found_display_name = result.get("displayName")
+        except Exception:
+            found_display_name = None
+
+        if not found_display_name:
+            try:
+                result = confluence_client.get_user_details_by_accountid(user_id)
+                found_display_name = result.get("displayName")
+            except Exception:
+                found_display_name = None
+
+        _USER_ID_TO_DISPLAY_NAME_CACHE[user_id] = found_display_name
+
+    return _USER_ID_TO_DISPLAY_NAME_CACHE.get(user_id) or _USER_NOT_FOUND
+
+
+def attachment_to_content(
+    confluence_client: OnyxConfluence,
+    attachment: dict[str, Any],
+    parent_content_id: str | None = None,
+) -> str | None:
+    """If it returns None, assume that we should skip this attachment."""
+    if not validate_attachment_filetype(attachment):
+        return None
+
+    if "api.atlassian.com" in confluence_client.url:
+        # https://developer.atlassian.com/cloud/confluence/rest/v1/api-group-content---attachments/#api-wiki-rest-api-content-id-child-attachment-attachmentid-download-get
+        if not parent_content_id:
+            logger.warning(
+                "parent_content_id is required to download attachments from Confluence Cloud!"
+            )
+            return None
+
+        download_link = (
+            confluence_client.url
+            + f"/rest/api/content/{parent_content_id}/child/attachment/{attachment['id']}/download"
+        )
+    else:
+        download_link = confluence_client.url + attachment["_links"]["download"]
+
+    attachment_size = attachment["extensions"]["fileSize"]
+    if attachment_size > CONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD:
+        logger.warning(
+            f"Skipping {download_link} due to size. "
+            f"size={attachment_size} "
+            f"threshold={CONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD}"
+        )
+        return None
+
+    logger.info(f"_attachment_to_content - _session.get: link={download_link}")
+
+    # why are we using session.get here? we probably won't retry these ... is that ok?
+    response = confluence_client._session.get(download_link)
+    if response.status_code != 200:
+        logger.warning(
+            f"Failed to fetch {download_link} with invalid status code {response.status_code}"
+        )
+        return None
+
+    extracted_text = extract_file_text(
+        io.BytesIO(response.content),
+        file_name=attachment["title"],
+        break_on_unprocessable=False,
     )
-    spaces = confluence_client_with_minimal_retries.get_all_spaces(limit=1)
+    if len(extracted_text) > CONFLUENCE_CONNECTOR_ATTACHMENT_CHAR_COUNT_THRESHOLD:
+        logger.warning(
+            f"Skipping {download_link} due to char count. "
+            f"char count={len(extracted_text)} "
+            f"threshold={CONFLUENCE_CONNECTOR_ATTACHMENT_CHAR_COUNT_THRESHOLD}"
+        )
+        return None
 
-    # uncomment the following for testing
-    # the following is an attempt to retrieve the user's timezone
-    # Unfornately, all data is returned in UTC regardless of the user's time zone
-    # even tho CQL parses incoming times based on the user's time zone
-    # space_key = spaces["results"][0]["key"]
-    # space_details = confluence_client_with_minimal_retries.cql(f"space.key={space_key}+AND+type=space")
+    return extracted_text
 
-    if not spaces:
-        raise RuntimeError(
-            f"No spaces found at {wiki_base}! "
-            "Check your credentials and wiki_base and make sure "
-            "is_cloud is set correctly."
+
+def extract_text_from_confluence_html(
+    confluence_client: OnyxConfluence,
+    confluence_object: dict[str, Any],
+    fetched_titles: set[str],
+) -> str:
+    """Parse a Confluence html page and replace the 'user Id' by the real
+        User Display Name
+
+    Args:
+        confluence_object (dict): The confluence object as a dict
+        confluence_client (Confluence): Confluence client
+        fetched_titles (set[str]): The titles of the pages that have already been fetched
+    Returns:
+        str: loaded and formated Confluence page
+    """
+    body = confluence_object["body"]
+    object_html = body.get("storage", body.get("view", {})).get("value")
+
+    soup = bs4.BeautifulSoup(object_html, "html.parser")
+    for user in soup.findAll("ri:user"):
+        user_id = (
+            user.attrs["ri:account-id"]
+            if "ri:account-id" in user.attrs
+            else user.get("ri:userkey")
+        )
+        if not user_id:
+            logger.warning(
+                "ri:userkey not found in ri:user element. " f"Found attrs: {user.attrs}"
+            )
+            continue
+        # Include @ sign for tagging, more clear for LLM
+        user.replaceWith("@" + _get_user(confluence_client, user_id))
+
+    for html_page_reference in soup.findAll("ac:structured-macro"):
+        # Here, we only want to process page within page macros
+        if html_page_reference.attrs.get("ac:name") != "include":
+            continue
+
+        page_data = html_page_reference.find("ri:page")
+        if not page_data:
+            logger.warning(
+                f"Skipping retrieval of {html_page_reference} because because page data is missing"
+            )
+            continue
+
+        page_title = page_data.attrs.get("ri:content-title")
+        if not page_title:
+            # only fetch pages that have a title
+            logger.warning(
+                f"Skipping retrieval of {html_page_reference} because it has no title"
+            )
+            continue
+
+        if page_title in fetched_titles:
+            # prevent recursive fetching of pages
+            logger.debug(f"Skipping {page_title} because it has already been fetched")
+            continue
+
+        fetched_titles.add(page_title)
+
+        # Wrap this in a try-except because there are some pages that might not exist
+        try:
+            page_query = f"type=page and title='{quote(page_title)}'"
+
+            page_contents: dict[str, Any] | None = None
+            # Confluence enforces title uniqueness, so we should only get one result here
+            for page in confluence_client.paginated_cql_retrieval(
+                cql=page_query,
+                expand="body.storage.value",
+                limit=1,
+            ):
+                page_contents = page
+                break
+        except Exception as e:
+            logger.warning(
+                f"Error getting page contents for object {confluence_object}: {e}"
+            )
+            continue
+
+        if not page_contents:
+            continue
+
+        text_from_page = extract_text_from_confluence_html(
+            confluence_client=confluence_client,
+            confluence_object=page_contents,
+            fetched_titles=fetched_titles,
         )
 
+        html_page_reference.replaceWith(text_from_page)
 
-def build_confluence_client(
-    credentials: dict[str, Any],
-    is_cloud: bool,
-    wiki_base: str,
-) -> OnyxConfluence:
-    try:
-        _validate_connector_configuration(
-            credentials=credentials,
-            is_cloud=is_cloud,
-            wiki_base=wiki_base,
-        )
-    except Exception as e:
-        raise ConnectorValidationError(str(e))
+    for html_link_body in soup.findAll("ac:link-body"):
+        # This extracts the text from inline links in the page so they can be
+        # represented in the document text as plain text
+        try:
+            text_from_link = html_link_body.text
+            html_link_body.replaceWith(f"(LINK TEXT: {text_from_link})")
+        except Exception as e:
+            logger.warning(f"Error processing ac:link-body: {e}")
 
-    return OnyxConfluence(
-        api_version="cloud" if is_cloud else "latest",
-        # Remove trailing slash from wiki_base if present
-        url=wiki_base.rstrip("/"),
-        # passing in username causes issues for Confluence data center
-        username=credentials["confluence_username"] if is_cloud else None,
-        password=credentials["confluence_access_token"] if is_cloud else None,
-        token=credentials["confluence_access_token"] if not is_cloud else None,
-        backoff_and_retry=True,
-        max_backoff_retries=10,
-        max_backoff_seconds=60,
-        cloud=is_cloud,
-    )
+    return format_document_soup(soup)

backend/onyx/connectors/confluence/utils.py

@@ -1,185 +1,38 @@
-import io
+import math
+import time
+from collections.abc import Callable
 from datetime import datetime
+from datetime import timedelta
 from datetime import timezone
 from typing import Any
+from typing import cast
 from typing import TYPE_CHECKING
+from typing import TypeVar
 from urllib.parse import parse_qs
 from urllib.parse import quote
 from urllib.parse import urlparse
 
 import bs4
+import requests
+from pydantic import BaseModel
 
-from onyx.configs.app_configs import (
-    CONFLUENCE_CONNECTOR_ATTACHMENT_CHAR_COUNT_THRESHOLD,
-)
-from onyx.configs.app_configs import CONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD
-from onyx.file_processing.extract_file_text import extract_file_text
-from onyx.file_processing.html_utils import format_document_soup
 from onyx.utils.logger import setup_logger
 
 if TYPE_CHECKING:
-    from onyx.connectors.confluence.onyx_confluence import OnyxConfluence
+    pass
 
 logger = setup_logger()
 
-
-_USER_EMAIL_CACHE: dict[str, str | None] = {}
+CONFLUENCE_OAUTH_TOKEN_URL = "https://auth.atlassian.com/oauth/token"
+RATE_LIMIT_MESSAGE_LOWERCASE = "Rate limit exceeded".lower()
 
 
-def get_user_email_from_username__server(
-    confluence_client: "OnyxConfluence", user_name: str
-) -> str | None:
-    global _USER_EMAIL_CACHE
-    if _USER_EMAIL_CACHE.get(user_name) is None:
-        try:
-            response = confluence_client.get_mobile_parameters(user_name)
-            email = response.get("email")
-        except Exception:
-            logger.warning(f"failed to get confluence email for {user_name}")
-            # For now, we'll just return None and log a warning. This means
-            # we will keep retrying to get the email every group sync.
-            email = None
-            # We may want to just return a string that indicates failure so we dont
-            # keep retrying
-            # email = f"FAILED TO GET CONFLUENCE EMAIL FOR {user_name}"
-        _USER_EMAIL_CACHE[user_name] = email
-    return _USER_EMAIL_CACHE[user_name]
-
-
-_USER_NOT_FOUND = "Unknown Confluence User"
-_USER_ID_TO_DISPLAY_NAME_CACHE: dict[str, str | None] = {}
-
-
-def _get_user(confluence_client: "OnyxConfluence", user_id: str) -> str:
-    """Get Confluence Display Name based on the account-id or userkey value
-
-    Args:
-        user_id (str): The user id (i.e: the account-id or userkey)
-        confluence_client (Confluence): The Confluence Client
-
-    Returns:
-        str: The User Display Name. 'Unknown User' if the user is deactivated or not found
-    """
-    global _USER_ID_TO_DISPLAY_NAME_CACHE
-    if _USER_ID_TO_DISPLAY_NAME_CACHE.get(user_id) is None:
-        try:
-            result = confluence_client.get_user_details_by_userkey(user_id)
-            found_display_name = result.get("displayName")
-        except Exception:
-            found_display_name = None
-
-        if not found_display_name:
-            try:
-                result = confluence_client.get_user_details_by_accountid(user_id)
-                found_display_name = result.get("displayName")
-            except Exception:
-                found_display_name = None
-
-        _USER_ID_TO_DISPLAY_NAME_CACHE[user_id] = found_display_name
-
-    return _USER_ID_TO_DISPLAY_NAME_CACHE.get(user_id) or _USER_NOT_FOUND
-
-
-def extract_text_from_confluence_html(
-    confluence_client: "OnyxConfluence",
-    confluence_object: dict[str, Any],
-    fetched_titles: set[str],
-) -> str:
-    """Parse a Confluence html page and replace the 'user Id' by the real
-        User Display Name
-
-    Args:
-        confluence_object (dict): The confluence object as a dict
-        confluence_client (Confluence): Confluence client
-        fetched_titles (set[str]): The titles of the pages that have already been fetched
-    Returns:
-        str: loaded and formated Confluence page
-    """
-    body = confluence_object["body"]
-    object_html = body.get("storage", body.get("view", {})).get("value")
-
-    soup = bs4.BeautifulSoup(object_html, "html.parser")
-    for user in soup.findAll("ri:user"):
-        user_id = (
-            user.attrs["ri:account-id"]
-            if "ri:account-id" in user.attrs
-            else user.get("ri:userkey")
-        )
-        if not user_id:
-            logger.warning(
-                "ri:userkey not found in ri:user element. " f"Found attrs: {user.attrs}"
-            )
-            continue
-        # Include @ sign for tagging, more clear for LLM
-        user.replaceWith("@" + _get_user(confluence_client, user_id))
-
-    for html_page_reference in soup.findAll("ac:structured-macro"):
-        # Here, we only want to process page within page macros
-        if html_page_reference.attrs.get("ac:name") != "include":
-            continue
-
-        page_data = html_page_reference.find("ri:page")
-        if not page_data:
-            logger.warning(
-                f"Skipping retrieval of {html_page_reference} because because page data is missing"
-            )
-            continue
-
-        page_title = page_data.attrs.get("ri:content-title")
-        if not page_title:
-            # only fetch pages that have a title
-            logger.warning(
-                f"Skipping retrieval of {html_page_reference} because it has no title"
-            )
-            continue
-
-        if page_title in fetched_titles:
-            # prevent recursive fetching of pages
-            logger.debug(f"Skipping {page_title} because it has already been fetched")
-            continue
-
-        fetched_titles.add(page_title)
-
-        # Wrap this in a try-except because there are some pages that might not exist
-        try:
-            page_query = f"type=page and title='{quote(page_title)}'"
-
-            page_contents: dict[str, Any] | None = None
-            # Confluence enforces title uniqueness, so we should only get one result here
-            for page in confluence_client.paginated_cql_retrieval(
-                cql=page_query,
-                expand="body.storage.value",
-                limit=1,
-            ):
-                page_contents = page
-                break
-        except Exception as e:
-            logger.warning(
-                f"Error getting page contents for object {confluence_object}: {e}"
-            )
-            continue
-
-        if not page_contents:
-            continue
-
-        text_from_page = extract_text_from_confluence_html(
-            confluence_client=confluence_client,
-            confluence_object=page_contents,
-            fetched_titles=fetched_titles,
-        )
-
-        html_page_reference.replaceWith(text_from_page)
-
-    for html_link_body in soup.findAll("ac:link-body"):
-        # This extracts the text from inline links in the page so they can be
-        # represented in the document text as plain text
-        try:
-            text_from_link = html_link_body.text
-            html_link_body.replaceWith(f"(LINK TEXT: {text_from_link})")
-        except Exception as e:
-            logger.warning(f"Error processing ac:link-body: {e}")
-
-    return format_document_soup(soup)
+class TokenResponse(BaseModel):
+    access_token: str
+    expires_in: int
+    token_type: str
+    refresh_token: str
+    scope: str
 
 
 def validate_attachment_filetype(attachment: dict[str, Any]) -> bool:
@@ -193,49 +46,6 @@ def validate_attachment_filetype(attachment: dict[str, Any]) -> bool:
     ]
 
 
-def attachment_to_content(
-    confluence_client: "OnyxConfluence",
-    attachment: dict[str, Any],
-) -> str | None:
-    """If it returns None, assume that we should skip this attachment."""
-    if not validate_attachment_filetype(attachment):
-        return None
-
-    download_link = confluence_client.url + attachment["_links"]["download"]
-
-    attachment_size = attachment["extensions"]["fileSize"]
-    if attachment_size > CONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD:
-        logger.warning(
-            f"Skipping {download_link} due to size. "
-            f"size={attachment_size} "
-            f"threshold={CONFLUENCE_CONNECTOR_ATTACHMENT_SIZE_THRESHOLD}"
-        )
-        return None
-
-    logger.info(f"_attachment_to_content - _session.get: link={download_link}")
-    response = confluence_client._session.get(download_link)
-    if response.status_code != 200:
-        logger.warning(
-            f"Failed to fetch {download_link} with invalid status code {response.status_code}"
-        )
-        return None
-
-    extracted_text = extract_file_text(
-        io.BytesIO(response.content),
-        file_name=attachment["title"],
-        break_on_unprocessable=False,
-    )
-    if len(extracted_text) > CONFLUENCE_CONNECTOR_ATTACHMENT_CHAR_COUNT_THRESHOLD:
-        logger.warning(
-            f"Skipping {download_link} due to char count. "
-            f"char count={len(extracted_text)} "
-            f"threshold={CONFLUENCE_CONNECTOR_ATTACHMENT_CHAR_COUNT_THRESHOLD}"
-        )
-        return None
-
-    return extracted_text
-
-
 def build_confluence_document_id(
     base_url: str, content_url: str, is_cloud: bool
 ) -> str:
@@ -284,6 +94,137 @@ def datetime_from_string(datetime_string: str) -> datetime:
     return datetime_object
 
 
+def confluence_refresh_tokens(
+    client_id: str, client_secret: str, cloud_id: str, refresh_token: str
+) -> dict[str, Any]:
+    # rotate the refresh and access token
+    # Note that access tokens are only good for an hour in confluence cloud,
+    # so we're going to have problems if the connector runs for longer
+    # https://developer.atlassian.com/cloud/confluence/oauth-2-3lo-apps/#use-a-refresh-token-to-get-another-access-token-and-refresh-token-pair
+    response = requests.post(
+        CONFLUENCE_OAUTH_TOKEN_URL,
+        headers={"Content-Type": "application/x-www-form-urlencoded"},
+        data={
+            "grant_type": "refresh_token",
+            "client_id": client_id,
+            "client_secret": client_secret,
+            "refresh_token": refresh_token,
+        },
+    )
+
+    try:
+        token_response = TokenResponse.model_validate_json(response.text)
+    except Exception:
+        raise RuntimeError("Confluence Cloud token refresh failed.")
+
+    now = datetime.now(timezone.utc)
+    expires_at = now + timedelta(seconds=token_response.expires_in)
+
+    new_credentials: dict[str, Any] = {}
+    new_credentials["confluence_access_token"] = token_response.access_token
+    new_credentials["confluence_refresh_token"] = token_response.refresh_token
+    new_credentials["created_at"] = now.isoformat()
+    new_credentials["expires_at"] = expires_at.isoformat()
+    new_credentials["expires_in"] = token_response.expires_in
+    new_credentials["scope"] = token_response.scope
+    new_credentials["cloud_id"] = cloud_id
+    return new_credentials
+
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+# https://developer.atlassian.com/cloud/confluence/rate-limiting/
+# this uses the native rate limiting option provided by the
+# confluence client and otherwise applies a simpler set of error handling
+def handle_confluence_rate_limit(confluence_call: F) -> F:
+    def wrapped_call(*args: list[Any], **kwargs: Any) -> Any:
+        MAX_RETRIES = 5
+
+        TIMEOUT = 600
+        timeout_at = time.monotonic() + TIMEOUT
+
+        for attempt in range(MAX_RETRIES):
+            if time.monotonic() > timeout_at:
+                raise TimeoutError(
+                    f"Confluence call attempts took longer than {TIMEOUT} seconds."
+                )
+
+            try:
+                # we're relying more on the client to rate limit itself
+                # and applying our own retries in a more specific set of circumstances
+                return confluence_call(*args, **kwargs)
+            except requests.HTTPError as e:
+                delay_until = _handle_http_error(e, attempt)
+                logger.warning(
+                    f"HTTPError in confluence call. "
+                    f"Retrying in {delay_until} seconds..."
+                )
+                while time.monotonic() < delay_until:
+                    # in the future, check a signal here to exit
+                    time.sleep(1)
+            except AttributeError as e:
+                # Some error within the Confluence library, unclear why it fails.
+                # Users reported it to be intermittent, so just retry
+                if attempt == MAX_RETRIES - 1:
+                    raise e
+
+                logger.exception(
+                    "Confluence Client raised an AttributeError. Retrying..."
+                )
+                time.sleep(5)
+
+    return cast(F, wrapped_call)
+
+
+def _handle_http_error(e: requests.HTTPError, attempt: int) -> int:
+    MIN_DELAY = 2
+    MAX_DELAY = 60
+    STARTING_DELAY = 5
+    BACKOFF = 2
+
+    # Check if the response or headers are None to avoid potential AttributeError
+    if e.response is None or e.response.headers is None:
+        logger.warning("HTTPError with `None` as response or as headers")
+        raise e
+
+    if (
+        e.response.status_code != 429
+        and RATE_LIMIT_MESSAGE_LOWERCASE not in e.response.text.lower()
+    ):
+        raise e
+
+    retry_after = None
+
+    retry_after_header = e.response.headers.get("Retry-After")
+    if retry_after_header is not None:
+        try:
+            retry_after = int(retry_after_header)
+            if retry_after > MAX_DELAY:
+                logger.warning(
+                    f"Clamping retry_after from {retry_after} to {MAX_DELAY} seconds..."
+                )
+                retry_after = MAX_DELAY
+            if retry_after < MIN_DELAY:
+                retry_after = MIN_DELAY
+        except ValueError:
+            pass
+
+    if retry_after is not None:
+        logger.warning(
+            f"Rate limiting with retry header. Retrying after {retry_after} seconds..."
+        )
+        delay = retry_after
+    else:
+        logger.warning(
+            "Rate limiting without retry header. Retrying with exponential backoff..."
+        )
+        delay = min(STARTING_DELAY * (BACKOFF**attempt), MAX_DELAY)
+
+    delay_until = math.ceil(time.monotonic() + delay)
+    return delay_until
+
+
 def get_single_param_from_url(url: str, param: str) -> str | None:
     """Get a parameter from a url"""
     parsed_url = urlparse(url)

backend/onyx/connectors/credentials_provider.py

@@ -0,0 +1,135 @@
+import uuid
+from types import TracebackType
+from typing import Any
+
+from redis.lock import Lock as RedisLock
+from sqlalchemy import select
+
+from onyx.connectors.interfaces import CredentialsProviderInterface
+from onyx.db.models import Credential
+from onyx.db.session import get_session_with_tenant
+from onyx.redis.redis_pool import get_redis_client
+
+
+class OnyxDBCredentialsProvider(
+    CredentialsProviderInterface["OnyxDBCredentialsProvider"]
+):
+    """Implementation to allow the connector to callback and update credentials in the db.
+    Required in cases where credentials can rotate while the connector is running.
+    """
+
+    LOCK_TTL = 900  # TTL of the lock
+
+    def __init__(self, tenant_id: str, connector_name: str, credential_id: int):
+        self._tenant_id = tenant_id
+        self._connector_name = connector_name
+        self._credential_id = credential_id
+
+        self.redis_client = get_redis_client(tenant_id=tenant_id)
+
+        # lock used to prevent overlapping renewal of credentials
+        self.lock_key = f"da_lock:connector:{connector_name}:credential_{credential_id}"
+        self._lock: RedisLock = self.redis_client.lock(self.lock_key, self.LOCK_TTL)
+
+    def __enter__(self) -> "OnyxDBCredentialsProvider":
+        acquired = self._lock.acquire(blocking_timeout=self.LOCK_TTL)
+        if not acquired:
+            raise RuntimeError(f"Could not acquire lock for key: {self.lock_key}")
+
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        """Release the lock when exiting the context."""
+        if self._lock and self._lock.owned():
+            self._lock.release()
+
+    def get_tenant_id(self) -> str | None:
+        return self._tenant_id
+
+    def get_provider_key(self) -> str:
+        return str(self._credential_id)
+
+    def get_credentials(self) -> dict[str, Any]:
+        with get_session_with_tenant(tenant_id=self._tenant_id) as db_session:
+            credential = db_session.execute(
+                select(Credential).where(Credential.id == self._credential_id)
+            ).scalar_one()
+
+            if credential is None:
+                raise ValueError(
+                    f"No credential found: credential={self._credential_id}"
+                )
+
+            return credential.credential_json
+
+    def set_credentials(self, credential_json: dict[str, Any]) -> None:
+        with get_session_with_tenant(tenant_id=self._tenant_id) as db_session:
+            try:
+                credential = db_session.execute(
+                    select(Credential)
+                    .where(Credential.id == self._credential_id)
+                    .with_for_update()
+                ).scalar_one()
+
+                if credential is None:
+                    raise ValueError(
+                        f"No credential found: credential={self._credential_id}"
+                    )
+
+                credential.credential_json = credential_json
+                db_session.commit()
+            except Exception:
+                db_session.rollback()
+                raise
+
+    def is_dynamic(self) -> bool:
+        return True
+
+
+class OnyxStaticCredentialsProvider(
+    CredentialsProviderInterface["OnyxStaticCredentialsProvider"]
+):
+    """Implementation (a very simple one!) to handle static credentials."""
+
+    def __init__(
+        self,
+        tenant_id: str | None,
+        connector_name: str,
+        credential_json: dict[str, Any],
+    ):
+        self._tenant_id = tenant_id
+        self._connector_name = connector_name
+        self._credential_json = credential_json
+
+        self._provider_key = str(uuid.uuid4())
+
+    def __enter__(self) -> "OnyxStaticCredentialsProvider":
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        pass
+
+    def get_tenant_id(self) -> str | None:
+        return self._tenant_id
+
+    def get_provider_key(self) -> str:
+        return self._provider_key
+
+    def get_credentials(self) -> dict[str, Any]:
+        return self._credential_json
+
+    def set_credentials(self, credential_json: dict[str, Any]) -> None:
+        self._credential_json = credential_json
+
+    def is_dynamic(self) -> bool:
+        return False

backend/onyx/connectors/factory.py

@@ -12,6 +12,7 @@ from onyx.connectors.blob.connector import BlobStorageConnector
 from onyx.connectors.bookstack.connector import BookstackConnector
 from onyx.connectors.clickup.connector import ClickupConnector
 from onyx.connectors.confluence.connector import ConfluenceConnector
+from onyx.connectors.credentials_provider import OnyxDBCredentialsProvider
 from onyx.connectors.discord.connector import DiscordConnector
 from onyx.connectors.discourse.connector import DiscourseConnector
 from onyx.connectors.document360.connector import Document360Connector
@@ -32,6 +33,7 @@ from onyx.connectors.guru.connector import GuruConnector
 from onyx.connectors.hubspot.connector import HubSpotConnector
 from onyx.connectors.interfaces import BaseConnector
 from onyx.connectors.interfaces import CheckpointConnector
+from onyx.connectors.interfaces import CredentialsConnector
 from onyx.connectors.interfaces import EventConnector
 from onyx.connectors.interfaces import LoadConnector
 from onyx.connectors.interfaces import PollConnector
@@ -57,6 +59,7 @@ from onyx.db.connector import fetch_connector_by_id
 from onyx.db.credentials import backend_update_credential_json
 from onyx.db.credentials import fetch_credential_by_id
 from onyx.db.models import Credential
+from shared_configs.contextvars import get_current_tenant_id
 
 
 class ConnectorMissingException(Exception):
@@ -167,10 +170,17 @@ def instantiate_connector(
     connector_class = identify_connector_class(source, input_type)
 
     connector = connector_class(**connector_specific_config)
-    new_credentials = connector.load_credentials(credential.credential_json)
 
-    if new_credentials is not None:
-        backend_update_credential_json(credential, new_credentials, db_session)
+    if isinstance(connector, CredentialsConnector):
+        provider = OnyxDBCredentialsProvider(
+            get_current_tenant_id(), str(source), credential.id
+        )
+        connector.set_credentials_provider(provider)
+    else:
+        new_credentials = connector.load_credentials(credential.credential_json)
+
+        if new_credentials is not None:
+            backend_update_credential_json(credential, new_credentials, db_session)
 
     return connector
 

backend/onyx/connectors/file/connector.py

@@ -16,7 +16,7 @@ from onyx.connectors.interfaces import LoadConnector
 from onyx.connectors.models import BasicExpertInfo
 from onyx.connectors.models import Document
 from onyx.connectors.models import Section
-from onyx.db.engine import get_session_with_current_tenant
+from onyx.db.session import get_session_with_current_tenant
 from onyx.file_processing.extract_file_text import detect_encoding
 from onyx.file_processing.extract_file_text import extract_file_text
 from onyx.file_processing.extract_file_text import get_file_ext

backend/onyx/connectors/interfaces.py

@@ -1,7 +1,10 @@
 import abc
 from collections.abc import Generator
 from collections.abc import Iterator
+from types import TracebackType
 from typing import Any
+from typing import Generic
+from typing import TypeVar
 
 from pydantic import BaseModel
 
@@ -111,6 +114,69 @@ class OAuthConnector(BaseConnector):
         raise NotImplementedError
 
 
+T = TypeVar("T", bound="CredentialsProviderInterface")
+
+
+class CredentialsProviderInterface(abc.ABC, Generic[T]):
+    @abc.abstractmethod
+    def __enter__(self) -> T:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def get_tenant_id(self) -> str | None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def get_provider_key(self) -> str:
+        """a unique key that the connector can use to lock around a credential
+        that might be used simultaneously.
+
+        Will typically be the credential id, but can also just be something random
+        in cases when there is nothing to lock (aka static credentials)
+        """
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def get_credentials(self) -> dict[str, Any]:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def set_credentials(self, credential_json: dict[str, Any]) -> None:
+        raise NotImplementedError
+
+    @abc.abstractmethod
+    def is_dynamic(self) -> bool:
+        """If dynamic, the credentials may change during usage ... maening the client
+        needs to use the locking features of the credentials provider to operate
+        correctly.
+
+        If static, the client can simply reference the credentials once and use them
+        through the entire indexing run.
+        """
+        raise NotImplementedError
+
+
+class CredentialsConnector(BaseConnector):
+    """Implement this if the connector needs to be able to read and write credentials
+    on the fly. Typically used with shared credentials/tokens that might be renewed
+    at any time."""
+
+    @abc.abstractmethod
+    def set_credentials_provider(
+        self, credentials_provider: CredentialsProviderInterface
+    ) -> None:
+        raise NotImplementedError
+
+
 # Event driven
 class EventConnector(BaseConnector):
     @abc.abstractmethod

backend/onyx/connectors/slack/utils.py

@@ -72,6 +72,7 @@ def make_slack_api_rate_limited(
     @wraps(call)
     def rate_limited_call(**kwargs: Any) -> SlackResponse:
         last_exception = None
+
         for _ in range(max_retries):
             try:
                 # Make the API call

backend/onyx/connectors/web/connector.py

@@ -42,6 +42,10 @@ from shared_configs.configs import MULTI_TENANT
 logger = setup_logger()
 
 WEB_CONNECTOR_MAX_SCROLL_ATTEMPTS = 20
+# Threshold for determining when to replace vs append iframe content
+IFRAME_TEXT_LENGTH_THRESHOLD = 700
+# Message indicating JavaScript is disabled, which often appears when scraping fails
+JAVASCRIPT_DISABLED_MESSAGE = "You have JavaScript disabled in your browser"
 
 
 class WEB_CONNECTOR_VALID_SETTINGS(str, Enum):
@@ -138,7 +142,8 @@ def get_internal_links(
         # Account for malformed backslashes in URLs
         href = href.replace("\\", "/")
 
-        if should_ignore_pound and "#" in href:
+        # "#!" indicates the page is using a hashbang URL, which is a client-side routing technique
+        if should_ignore_pound and "#" in href and "#!" not in href:
             href = href.split("#")[0]
 
         if not is_valid_url(href):
@@ -288,6 +293,7 @@ class WebConnector(LoadConnector):
         and converts them into documents"""
         visited_links: set[str] = set()
         to_visit: list[str] = self.to_visit_list
+        content_hashes = set()
 
         if not to_visit:
             raise ValueError("No URLs to visit")
@@ -302,29 +308,30 @@ class WebConnector(LoadConnector):
         playwright, context = start_playwright()
         restart_playwright = False
         while to_visit:
-            current_url = to_visit.pop()
-            if current_url in visited_links:
+            initial_url = to_visit.pop()
+            if initial_url in visited_links:
                 continue
-            visited_links.add(current_url)
+            visited_links.add(initial_url)
 
             try:
-                protected_url_check(current_url)
+                protected_url_check(initial_url)
             except Exception as e:
-                last_error = f"Invalid URL {current_url} due to {e}"
+                last_error = f"Invalid URL {initial_url} due to {e}"
                 logger.warning(last_error)
                 continue
 
-            logger.info(f"Visiting {current_url}")
+            index = len(visited_links)
+            logger.info(f"{index}: Visiting {initial_url}")
 
             try:
-                check_internet_connection(current_url)
+                check_internet_connection(initial_url)
                 if restart_playwright:
                     playwright, context = start_playwright()
                     restart_playwright = False
 
-                if current_url.split(".")[-1] == "pdf":
+                if initial_url.split(".")[-1] == "pdf":
                     # PDF files are not checked for links
-                    response = requests.get(current_url)
+                    response = requests.get(initial_url)
                     page_text, metadata = read_pdf_file(
                         file=io.BytesIO(response.content)
                     )
@@ -332,10 +339,10 @@ class WebConnector(LoadConnector):
 
                     doc_batch.append(
                         Document(
-                            id=current_url,
-                            sections=[Section(link=current_url, text=page_text)],
+                            id=initial_url,
+                            sections=[Section(link=initial_url, text=page_text)],
                             source=DocumentSource.WEB,
-                            semantic_identifier=current_url.split("/")[-1],
+                            semantic_identifier=initial_url.split("/")[-1],
                             metadata=metadata,
                             doc_updated_at=_get_datetime_from_last_modified_header(
                                 last_modified
@@ -347,21 +354,29 @@ class WebConnector(LoadConnector):
                     continue
 
                 page = context.new_page()
-                page_response = page.goto(current_url)
+
+                # Can't use wait_until="networkidle" because it interferes with the scrolling behavior
+                page_response = page.goto(
+                    initial_url,
+                    timeout=30000,  # 30 seconds
+                )
+
                 last_modified = (
                     page_response.header_value("Last-Modified")
                     if page_response
                     else None
                 )
-                final_page = page.url
-                if final_page != current_url:
-                    logger.info(f"Redirected to {final_page}")
-                    protected_url_check(final_page)
-                    current_url = final_page
-                    if current_url in visited_links:
-                        logger.info("Redirected page already indexed")
+                final_url = page.url
+                if final_url != initial_url:
+                    protected_url_check(final_url)
+                    initial_url = final_url
+                    if initial_url in visited_links:
+                        logger.info(
+                            f"{index}: {initial_url} redirected to {final_url} - already indexed"
+                        )
                         continue
-                    visited_links.add(current_url)
+                    logger.info(f"{index}: {initial_url} redirected to {final_url}")
+                    visited_links.add(initial_url)
 
                 if self.scroll_before_scraping:
                     scroll_attempts = 0
@@ -379,26 +394,58 @@ class WebConnector(LoadConnector):
                 soup = BeautifulSoup(content, "html.parser")
 
                 if self.recursive:
-                    internal_links = get_internal_links(base_url, current_url, soup)
+                    internal_links = get_internal_links(base_url, initial_url, soup)
                     for link in internal_links:
                         if link not in visited_links:
                             to_visit.append(link)
 
                 if page_response and str(page_response.status)[0] in ("4", "5"):
-                    last_error = f"Skipped indexing {current_url} due to HTTP {page_response.status} response"
+                    last_error = f"Skipped indexing {initial_url} due to HTTP {page_response.status} response"
                     logger.info(last_error)
                     continue
 
                 parsed_html = web_html_cleanup(soup, self.mintlify_cleanup)
 
+                """For websites containing iframes that need to be scraped,
+                the code below can extract text from within these iframes.
+                """
+                logger.debug(
+                    f"{index}: Length of cleaned text {len(parsed_html.cleaned_text)}"
+                )
+                if JAVASCRIPT_DISABLED_MESSAGE in parsed_html.cleaned_text:
+                    iframe_count = page.frame_locator("iframe").locator("html").count()
+                    if iframe_count > 0:
+                        iframe_texts = (
+                            page.frame_locator("iframe")
+                            .locator("html")
+                            .all_inner_texts()
+                        )
+                        document_text = "\n".join(iframe_texts)
+                        """ 700 is the threshold value for the length of the text extracted
+                        from the iframe based on the issue faced """
+                        if len(parsed_html.cleaned_text) < IFRAME_TEXT_LENGTH_THRESHOLD:
+                            parsed_html.cleaned_text = document_text
+                        else:
+                            parsed_html.cleaned_text += "\n" + document_text
+
+                # Sometimes pages with #! will serve duplicate content
+                # There are also just other ways this can happen
+                hashed_text = hash((parsed_html.title, parsed_html.cleaned_text))
+                if hashed_text in content_hashes:
+                    logger.info(
+                        f"{index}: Skipping duplicate title + content for {initial_url}"
+                    )
+                    continue
+                content_hashes.add(hashed_text)
+
                 doc_batch.append(
                     Document(
-                        id=current_url,
+                        id=initial_url,
                         sections=[
-                            Section(link=current_url, text=parsed_html.cleaned_text)
+                            Section(link=initial_url, text=parsed_html.cleaned_text)
                         ],
                         source=DocumentSource.WEB,
-                        semantic_identifier=parsed_html.title or current_url,
+                        semantic_identifier=parsed_html.title or initial_url,
                         metadata={},
                         doc_updated_at=_get_datetime_from_last_modified_header(
                             last_modified
@@ -410,7 +457,7 @@ class WebConnector(LoadConnector):
 
                 page.close()
             except Exception as e:
-                last_error = f"Failed to fetch '{current_url}': {e}"
+                last_error = f"Failed to fetch '{initial_url}': {e}"
                 logger.exception(last_error)
                 playwright.stop()
                 restart_playwright = True

backend/onyx/context/search/models.py

@@ -76,6 +76,10 @@ class SavedSearchSettings(InferenceSettings, IndexingSetting):
             provider_type=search_settings.provider_type,
             index_name=search_settings.index_name,
             multipass_indexing=search_settings.multipass_indexing,
+            embedding_precision=search_settings.embedding_precision,
+            reduced_dimension=search_settings.reduced_dimension,
+            # Whether switching to this model requires re-indexing
+            background_reindex_enabled=search_settings.background_reindex_enabled,
             # Reranking Details
             rerank_model_name=search_settings.rerank_model_name,
             rerank_provider_type=search_settings.rerank_provider_type,

backend/onyx/db/auth.py

@@ -16,11 +16,11 @@ from sqlalchemy.orm import Session
 from onyx.auth.invited_users import get_invited_users
 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.models import AccessToken
 from onyx.db.models import OAuthAccount
 from onyx.db.models import User
+from onyx.db.session import get_async_session
+from onyx.db.session import get_async_session_with_tenant
 from onyx.utils.variable_functionality import (
     fetch_versioned_implementation_with_fallback,
 )

backend/onyx/db/chat.py

@@ -168,7 +168,7 @@ def get_chat_sessions_by_user(
     if not include_onyxbot_flows:
         stmt = stmt.where(ChatSession.onyxbot_flow.is_(False))
 
-    stmt = stmt.order_by(desc(ChatSession.time_created))
+    stmt = stmt.order_by(desc(ChatSession.time_updated))
 
     if deleted is not None:
         stmt = stmt.where(ChatSession.deleted == deleted)
@@ -962,6 +962,7 @@ def translate_db_message_to_chat_message_detail(
             chat_message.sub_questions
         ),
         refined_answer_improvement=chat_message.refined_answer_improvement,
+        is_agentic=chat_message.is_agentic,
         error=chat_message.error,
     )
 

backend/onyx/db/chat_search.py

@@ -3,14 +3,13 @@ from typing import Optional
 from typing import Tuple
 from uuid import UUID
 
+from sqlalchemy import column
 from sqlalchemy import desc
 from sqlalchemy import func
-from sqlalchemy import literal
-from sqlalchemy import Select
 from sqlalchemy import select
-from sqlalchemy import union_all
 from sqlalchemy.orm import joinedload
 from sqlalchemy.orm import Session
+from sqlalchemy.sql.expression import ColumnClause
 
 from onyx.db.models import ChatMessage
 from onyx.db.models import ChatSession
@@ -26,127 +25,87 @@ def search_chat_sessions(
     include_onyxbot_flows: bool = False,
 ) -> Tuple[List[ChatSession], bool]:
     """
-    Search for chat sessions based on the provided query.
-    If no query is provided, returns recent chat sessions.
+    Fast full-text search on ChatSession + ChatMessage using tsvectors.
 
-    Returns a tuple of (chat_sessions, has_more)
+    If no query is provided, returns the most recent chat sessions.
+    Otherwise, searches both chat messages and session descriptions.
+
+    Returns a tuple of (sessions, has_more) where has_more indicates if
+    there are additional results beyond the requested page.
     """
-    offset = (page - 1) * page_size
+    offset_val = (page - 1) * page_size
 
-    # If no search query, we use standard SQLAlchemy pagination
+    # If no query, just return the most recent sessions
     if not query or not query.strip():
-        stmt = select(ChatSession)
-        if user_id:
+        stmt = (
+            select(ChatSession)
+            .order_by(desc(ChatSession.time_created))
+            .offset(offset_val)
+            .limit(page_size + 1)
+        )
+        if user_id is not None:
             stmt = stmt.where(ChatSession.user_id == user_id)
         if not include_onyxbot_flows:
             stmt = stmt.where(ChatSession.onyxbot_flow.is_(False))
         if not include_deleted:
             stmt = stmt.where(ChatSession.deleted.is_(False))
 
-        stmt = stmt.order_by(desc(ChatSession.time_created))
-
-        # Apply pagination
-        stmt = stmt.offset(offset).limit(page_size + 1)
         result = db_session.execute(stmt.options(joinedload(ChatSession.persona)))
-        chat_sessions = result.scalars().all()
+        sessions = result.scalars().all()
 
-        has_more = len(chat_sessions) > page_size
+        has_more = len(sessions) > page_size
         if has_more:
-            chat_sessions = chat_sessions[:page_size]
+            sessions = sessions[:page_size]
 
-        return list(chat_sessions), has_more
+        return list(sessions), has_more
 
-    words = query.lower().strip().split()
+    # Otherwise, proceed with full-text search
+    query = query.strip()
 
-    # Message mach subquery
-    message_matches = []
-    for word in words:
-        word_like = f"%{word}%"
-        message_match: Select = (
-            select(ChatMessage.chat_session_id, literal(1.0).label("search_rank"))
-            .join(ChatSession, ChatSession.id == ChatMessage.chat_session_id)
-            .where(func.lower(ChatMessage.message).like(word_like))
-        )
-
-        if user_id:
-            message_match = message_match.where(ChatSession.user_id == user_id)
-
-        message_matches.append(message_match)
-
-    if message_matches:
-        message_matches_query = union_all(*message_matches).alias("message_matches")
-    else:
-        return [], False
-
-    # Description matches
-    description_match: Select = select(
-        ChatSession.id.label("chat_session_id"), literal(0.5).label("search_rank")
-    ).where(func.lower(ChatSession.description).like(f"%{query.lower()}%"))
-
-    if user_id:
-        description_match = description_match.where(ChatSession.user_id == user_id)
+    base_conditions = []
+    if user_id is not None:
+        base_conditions.append(ChatSession.user_id == user_id)
     if not include_onyxbot_flows:
-        description_match = description_match.where(ChatSession.onyxbot_flow.is_(False))
+        base_conditions.append(ChatSession.onyxbot_flow.is_(False))
     if not include_deleted:
-        description_match = description_match.where(ChatSession.deleted.is_(False))
+        base_conditions.append(ChatSession.deleted.is_(False))
 
-    # Combine all match sources
-    combined_matches = union_all(
-        message_matches_query.select(), description_match
-    ).alias("combined_matches")
+    message_tsv: ColumnClause = column("message_tsv")
+    description_tsv: ColumnClause = column("description_tsv")
 
-    # Use CTE to group and get max rank
-    session_ranks = (
-        select(
-            combined_matches.c.chat_session_id,
-            func.max(combined_matches.c.search_rank).label("rank"),
-        )
-        .group_by(combined_matches.c.chat_session_id)
-        .alias("session_ranks")
+    ts_query = func.plainto_tsquery("english", query)
+
+    description_session_ids = (
+        select(ChatSession.id)
+        .where(*base_conditions)
+        .where(description_tsv.op("@@")(ts_query))
     )
 
-    # Get ranked sessions with pagination
-    ranked_query = (
-        db_session.query(session_ranks.c.chat_session_id, session_ranks.c.rank)
-        .order_by(desc(session_ranks.c.rank), session_ranks.c.chat_session_id)
-        .offset(offset)
+    message_session_ids = (
+        select(ChatMessage.chat_session_id)
+        .join(ChatSession, ChatMessage.chat_session_id == ChatSession.id)
+        .where(*base_conditions)
+        .where(message_tsv.op("@@")(ts_query))
+    )
+
+    combined_ids = description_session_ids.union(message_session_ids).alias(
+        "combined_ids"
+    )
+
+    final_stmt = (
+        select(ChatSession)
+        .join(combined_ids, ChatSession.id == combined_ids.c.id)
+        .order_by(desc(ChatSession.time_created))
+        .distinct()
+        .offset(offset_val)
         .limit(page_size + 1)
+        .options(joinedload(ChatSession.persona))
     )
 
-    result = ranked_query.all()
+    session_objs = db_session.execute(final_stmt).scalars().all()
 
-    # Extract session IDs and ranks
-    session_ids_with_ranks = {row.chat_session_id: row.rank for row in result}
-    session_ids = list(session_ids_with_ranks.keys())
-
-    if not session_ids:
-        return [], False
-
-    # Now, let's query the actual ChatSession objects using the IDs
-    stmt = select(ChatSession).where(ChatSession.id.in_(session_ids))
-
-    if user_id:
-        stmt = stmt.where(ChatSession.user_id == user_id)
-    if not include_onyxbot_flows:
-        stmt = stmt.where(ChatSession.onyxbot_flow.is_(False))
-    if not include_deleted:
-        stmt = stmt.where(ChatSession.deleted.is_(False))
-
-    # Full objects with eager loading
-    result = db_session.execute(stmt.options(joinedload(ChatSession.persona)))
-    chat_sessions = result.scalars().all()
-
-    # Sort based on above ranking
-    chat_sessions = sorted(
-        chat_sessions,
-        key=lambda session: (
-            -session_ids_with_ranks.get(session.id, 0),  # Rank (higher first)
-            session.time_created.timestamp() * -1,  # Then by time (newest first)
-        ),
-    )
-
-    has_more = len(chat_sessions) > page_size
+    has_more = len(session_objs) > page_size
     if has_more:
-        chat_sessions = chat_sessions[:page_size]
+        session_objs = session_objs[:page_size]
 
-    return chat_sessions, has_more
+    return list(session_objs), has_more

backend/onyx/db/connector_credential_pair.py

@@ -16,7 +16,6 @@ from onyx.configs.app_configs import DISABLE_AUTH
 from onyx.db.connector import fetch_connector_by_id
 from onyx.db.credentials import fetch_credential_by_id
 from onyx.db.credentials import fetch_credential_by_id_for_user
-from onyx.db.engine import get_session_context_manager
 from onyx.db.enums import AccessType
 from onyx.db.enums import ConnectorCredentialPairStatus
 from onyx.db.models import ConnectorCredentialPair
@@ -29,6 +28,7 @@ from onyx.db.models import User
 from onyx.db.models import User__UserGroup
 from onyx.db.models import UserGroup__ConnectorCredentialPair
 from onyx.db.models import UserRole
+from onyx.db.session import get_session_context_manager
 from onyx.server.models import StatusResponse
 from onyx.utils.logger import setup_logger
 from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop

backend/onyx/db/credentials.py

@@ -360,18 +360,13 @@ def backend_update_credential_json(
     db_session.commit()
 
 
-def delete_credential(
+def _delete_credential_internal(
+    credential: Credential,
     credential_id: int,
-    user: User | None,
     db_session: Session,
     force: bool = False,
 ) -> None:
-    credential = fetch_credential_by_id_for_user(credential_id, user, db_session)
-    if credential is None:
-        raise ValueError(
-            f"Credential by provided id {credential_id} does not exist or does not belong to user"
-        )
-
+    """Internal utility function to handle the actual deletion of a credential"""
     associated_connectors = (
         db_session.query(ConnectorCredentialPair)
         .filter(ConnectorCredentialPair.credential_id == credential_id)
@@ -416,6 +411,35 @@ def delete_credential(
     db_session.commit()
 
 
+def delete_credential_for_user(
+    credential_id: int,
+    user: User,
+    db_session: Session,
+    force: bool = False,
+) -> None:
+    """Delete a credential that belongs to a specific user"""
+    credential = fetch_credential_by_id_for_user(credential_id, user, db_session)
+    if credential is None:
+        raise ValueError(
+            f"Credential by provided id {credential_id} does not exist or does not belong to user"
+        )
+
+    _delete_credential_internal(credential, credential_id, db_session, force)
+
+
+def delete_credential(
+    credential_id: int,
+    db_session: Session,
+    force: bool = False,
+) -> None:
+    """Delete a credential regardless of ownership (admin function)"""
+    credential = fetch_credential_by_id(credential_id, db_session)
+    if credential is None:
+        raise ValueError(f"Credential by provided id {credential_id} does not exist")
+
+    _delete_credential_internal(credential, credential_id, db_session, force)
+
+
 def create_initial_public_credential(db_session: Session) -> None:
     error_msg = (
         "DB is not in a valid initial state."

backend/onyx/db/document.py

@@ -24,7 +24,6 @@ from sqlalchemy.sql.expression import null
 from onyx.configs.constants import DEFAULT_BOOST
 from onyx.configs.constants import DocumentSource
 from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id
-from onyx.db.engine import get_session_context_manager
 from onyx.db.enums import AccessType
 from onyx.db.enums import ConnectorCredentialPairStatus
 from onyx.db.feedback import delete_document_feedback_for_documents__no_commit
@@ -34,6 +33,7 @@ from onyx.db.models import Credential
 from onyx.db.models import Document as DbDocument
 from onyx.db.models import DocumentByConnectorCredentialPair
 from onyx.db.models import User
+from onyx.db.session import get_session_context_manager
 from onyx.db.tag import delete_document_tags_for_documents__no_commit
 from onyx.db.utils import model_to_dict
 from onyx.document_index.interfaces import DocumentMetadata

backend/onyx/db/document_set.py

@@ -323,15 +323,27 @@ def update_document_set(
         _mark_document_set_cc_pairs_as_outdated__no_commit(
             db_session=db_session, document_set_id=document_set_row.id
         )
+
+        # commit this before performing more updates on the row id or else
+        # we'll have conflicting updates in the same commit
+        db_session.commit()
+
         # add in rows for the new CC pairs
-        ds_cc_pairs = [
-            DocumentSet__ConnectorCredentialPair(
+        existing_cc_pair_ids: set[int] = set()  # use to avoid duplicates
+
+        ds_cc_pairs: list[DocumentSet__ConnectorCredentialPair] = []
+        for cc_pair_id in document_set_update_request.cc_pair_ids:
+            if cc_pair_id in existing_cc_pair_ids:
+                continue
+
+            item = DocumentSet__ConnectorCredentialPair(
                 document_set_id=document_set_update_request.id,
                 connector_credential_pair_id=cc_pair_id,
                 is_current=True,
             )
-            for cc_pair_id in document_set_update_request.cc_pair_ids
-        ]
+            ds_cc_pairs.append(item)
+            existing_cc_pair_ids.add(cc_pair_id)
+
         db_session.add_all(ds_cc_pairs)
         db_session.commit()
     except:

backend/onyx/db/engine.py

@@ -1,30 +1,20 @@
-import contextlib
 import os
-import re
 import ssl
 import threading
 import time
-from collections.abc import AsyncGenerator
-from collections.abc import Generator
-from contextlib import asynccontextmanager
-from contextlib import contextmanager
 from datetime import datetime
 from typing import Any
-from typing import ContextManager
 
 import asyncpg  # type: ignore
 import boto3
-from fastapi import HTTPException
 from sqlalchemy import event
 from sqlalchemy import pool
 from sqlalchemy import text
 from sqlalchemy.engine import create_engine
 from sqlalchemy.engine import Engine
 from sqlalchemy.ext.asyncio import AsyncEngine
-from sqlalchemy.ext.asyncio import AsyncSession
 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_NAME
 from onyx.configs.app_configs import LOG_POSTGRES_CONN_COUNTS
@@ -33,7 +23,6 @@ from onyx.configs.app_configs import POSTGRES_API_SERVER_POOL_OVERFLOW
 from onyx.configs.app_configs import POSTGRES_API_SERVER_POOL_SIZE
 from onyx.configs.app_configs import POSTGRES_DB
 from onyx.configs.app_configs import POSTGRES_HOST
-from onyx.configs.app_configs import POSTGRES_IDLE_SESSIONS_TIMEOUT
 from onyx.configs.app_configs import POSTGRES_PASSWORD
 from onyx.configs.app_configs import POSTGRES_POOL_PRE_PING
 from onyx.configs.app_configs import POSTGRES_POOL_RECYCLE
@@ -42,13 +31,7 @@ from onyx.configs.app_configs import POSTGRES_USE_NULL_POOL
 from onyx.configs.app_configs import POSTGRES_USER
 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
-from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
-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
 
 logger = setup_logger()
 
@@ -59,7 +42,6 @@ USE_IAM_AUTH = os.getenv("USE_IAM_AUTH", "False").lower() == "true"
 
 # 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:
@@ -176,13 +158,6 @@ def get_db_current_time(db_session: Session) -> datetime:
     return result
 
 
-SCHEMA_NAME_REGEX = re.compile(r"^[a-zA-Z0-9_-]+$")
-
-
-def is_valid_schema_name(name: str) -> bool:
-    return SCHEMA_NAME_REGEX.match(name) is not None
-
-
 class SqlEngine:
     _engine: Engine | None = None
     _lock: threading.Lock = threading.Lock()
@@ -258,31 +233,6 @@ class SqlEngine:
                 cls._engine = None
 
 
-def get_all_tenant_ids() -> list[str]:
-    """Returning [None] means the only tenant is the 'public' or self hosted tenant."""
-
-    if not MULTI_TENANT:
-        return [POSTGRES_DEFAULT_SCHEMA]
-
-    with get_session_with_shared_schema() as session:
-        result = session.execute(
-            text(
-                f"""
-                SELECT schema_name
-                FROM information_schema.schemata
-                WHERE schema_name NOT IN ('pg_catalog', 'information_schema', '{POSTGRES_DEFAULT_SCHEMA}')"""
-            )
-        )
-        tenant_ids = [row[0] for row in result]
-
-    valid_tenants = [
-        tenant
-        for tenant in tenant_ids
-        if tenant is None or tenant.startswith(TENANT_ID_PREFIX)
-    ]
-    return valid_tenants
-
-
 def get_sqlalchemy_engine() -> Engine:
     return SqlEngine.get_engine()
 
@@ -352,164 +302,6 @@ def get_sqlalchemy_async_engine() -> AsyncEngine:
     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,
-    class_=AsyncSession,
-    expire_on_commit=False,
-)
-
-
-@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()
-
-    with get_session_with_tenant(tenant_id=tenant_id) as session:
-        yield session
-
-
-# Used in multi tenant mode when need to refer to the shared `public` schema
-@contextmanager
-def get_session_with_shared_schema() -> Generator[Session, None, None]:
-    token = CURRENT_TENANT_ID_CONTEXTVAR.set(POSTGRES_DEFAULT_SCHEMA)
-    with get_session_with_tenant(tenant_id=POSTGRES_DEFAULT_SCHEMA) as session:
-        yield session
-    CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
-
-
-@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)
-
-    if not is_valid_schema_name(tenant_id):
-        raise HTTPException(status_code=400, detail="Invalid tenant ID")
-
-    with engine.connect() as connection:
-        dbapi_connection = connection.connection
-        cursor = dbapi_connection.cursor()
-        try:
-            cursor.execute(f'SET search_path = "{tenant_id}"')
-            if POSTGRES_IDLE_SESSIONS_TIMEOUT:
-                cursor.execute(
-                    text(
-                        f"SET SESSION idle_in_transaction_session_timeout = {POSTGRES_IDLE_SESSIONS_TIMEOUT}"
-                    )
-                )
-        finally:
-            cursor.close()
-
-        with Session(bind=connection, expire_on_commit=False) as session:
-            try:
-                yield session
-            finally:
-                if MULTI_TENANT:
-                    cursor = dbapi_connection.cursor()
-                    try:
-                        cursor.execute('SET search_path TO "$user", public')
-                    finally:
-                        cursor.close()
-
-
-def set_search_path_on_checkout(
-    dbapi_conn: Any, connection_record: Any, connection_proxy: Any
-) -> None:
-    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}"')
-
-
-def get_session_generator_with_tenant() -> Generator[Session, None, None]:
-    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")
-
-    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() -> 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
-
-
-def get_session_context_manager() -> ContextManager[Session]:
-    """Context manager for database sessions."""
-    return contextlib.contextmanager(get_session_generator_with_tenant)()
-
-
-def get_session_factory() -> sessionmaker[Session]:
-    global SessionFactory
-    if SessionFactory is None:
-        SessionFactory = sessionmaker(bind=get_sqlalchemy_engine())
-    return SessionFactory
-
-
 async def warm_up_connections(
     sync_connections_to_warm_up: int = 20, async_connections_to_warm_up: int = 20
 ) -> None:

backend/onyx/db/enums.py

@@ -63,6 +63,9 @@ class IndexModelStatus(str, PyEnum):
     PRESENT = "PRESENT"
     FUTURE = "FUTURE"
 
+    def is_current(self) -> bool:
+        return self == IndexModelStatus.PRESENT
+
 
 class ChatSessionSharedStatus(str, PyEnum):
     PUBLIC = "public"
@@ -83,3 +86,11 @@ class AccessType(str, PyEnum):
     PUBLIC = "public"
     PRIVATE = "private"
     SYNC = "sync"
+
+
+class EmbeddingPrecision(str, PyEnum):
+    # matches vespa tensor type
+    # only support float / bfloat16 for now, since there's not a
+    # good reason to specify anything else
+    BFLOAT16 = "bfloat16"
+    FLOAT = "float"

backend/onyx/db/index_attempt.py

@@ -16,12 +16,12 @@ from sqlalchemy.orm import Session
 from sqlalchemy.sql import Select
 
 from onyx.connectors.models import ConnectorFailure
-from onyx.db.engine import get_session_context_manager
 from onyx.db.models import IndexAttempt
 from onyx.db.models import IndexAttemptError
 from onyx.db.models import IndexingStatus
 from onyx.db.models import IndexModelStatus
 from onyx.db.models import SearchSettings
+from onyx.db.session import get_session_context_manager
 from onyx.server.documents.models import ConnectorCredentialPair
 from onyx.server.documents.models import ConnectorCredentialPairIdentifier
 from onyx.utils.logger import setup_logger

backend/onyx/db/models.py

@@ -7,6 +7,7 @@ from typing import Optional
 from uuid import uuid4
 
 from pydantic import BaseModel
+from sqlalchemy.orm import validates
 from typing_extensions import TypedDict  # noreorder
 from uuid import UUID
 
@@ -25,6 +26,7 @@ from sqlalchemy import ForeignKey
 from sqlalchemy import func
 from sqlalchemy import Index
 from sqlalchemy import Integer
+
 from sqlalchemy import Sequence
 from sqlalchemy import String
 from sqlalchemy import Text
@@ -44,7 +46,13 @@ from onyx.configs.constants import DEFAULT_BOOST, MilestoneRecordType
 from onyx.configs.constants import DocumentSource
 from onyx.configs.constants import FileOrigin
 from onyx.configs.constants import MessageType
-from onyx.db.enums import AccessType, IndexingMode, SyncType, SyncStatus
+from onyx.db.enums import (
+    AccessType,
+    EmbeddingPrecision,
+    IndexingMode,
+    SyncType,
+    SyncStatus,
+)
 from onyx.configs.constants import NotificationType
 from onyx.configs.constants import SearchFeedbackType
 from onyx.configs.constants import TokenRateLimitScope
@@ -205,6 +213,10 @@ class User(SQLAlchemyBaseUserTableUUID, Base):
         primaryjoin="User.id == foreign(ConnectorCredentialPair.creator_id)",
     )
 
+    @validates("email")
+    def validate_email(self, key: str, value: str) -> str:
+        return value.lower() if value else value
+
     @property
     def password_configured(self) -> bool:
         """
@@ -710,6 +722,23 @@ class SearchSettings(Base):
         ForeignKey("embedding_provider.provider_type"), nullable=True
     )
 
+    # Whether switching to this model should re-index all connectors in the background
+    # if no re-index is needed, will be ignored. Only used during the switch-over process.
+    background_reindex_enabled: Mapped[bool] = mapped_column(Boolean, default=True)
+
+    # allows for quantization -> less memory usage for a small performance hit
+    embedding_precision: Mapped[EmbeddingPrecision] = mapped_column(
+        Enum(EmbeddingPrecision, native_enum=False)
+    )
+
+    # can be used to reduce dimensionality of vectors and save memory with
+    # a small performance hit. More details in the `Reducing embedding dimensions`
+    # section here:
+    # https://platform.openai.com/docs/guides/embeddings#embedding-models
+    # If not specified, will just use the model_dim without any reduction.
+    # NOTE: this is only currently available for OpenAI models
+    reduced_dimension: Mapped[int | None] = mapped_column(Integer, nullable=True)
+
     # Mini and Large Chunks (large chunk also checks for model max context)
     multipass_indexing: Mapped[bool] = mapped_column(Boolean, default=True)
 
@@ -791,6 +820,12 @@ class SearchSettings(Base):
             self.multipass_indexing, self.model_name, self.provider_type
         )
 
+    @property
+    def final_embedding_dim(self) -> int:
+        if self.reduced_dimension:
+            return self.reduced_dimension
+        return self.model_dim
+
     @staticmethod
     def can_use_large_chunks(
         multipass: bool, model_name: str, provider_type: EmbeddingProvider | None
@@ -1755,6 +1790,7 @@ class ChannelConfig(TypedDict):
     channel_name: str | None  # None for default channel config
     respond_tag_only: NotRequired[bool]  # defaults to False
     respond_to_bots: NotRequired[bool]  # defaults to False
+    is_ephemeral: NotRequired[bool]  # defaults to False
     respond_member_group_list: NotRequired[list[str]]
     answer_filters: NotRequired[list[AllowedAnswerFilters]]
     # If None then no follow up
@@ -2269,6 +2305,10 @@ class UserTenantMapping(Base):
     email: Mapped[str] = mapped_column(String, nullable=False, primary_key=True)
     tenant_id: Mapped[str] = mapped_column(String, nullable=False)
 
+    @validates("email")
+    def validate_email(self, key: str, value: str) -> str:
+        return value.lower() if value else value
+
 
 # This is a mapping from tenant IDs to anonymous user paths
 class TenantAnonymousUserPath(Base):

backend/onyx/db/persona.py

@@ -209,13 +209,21 @@ def create_update_persona(
         if not all_prompt_ids:
             raise ValueError("No prompt IDs provided")
 
+        is_default_persona: bool | None = create_persona_request.is_default_persona
         # Default persona validation
         if create_persona_request.is_default_persona:
             if not create_persona_request.is_public:
                 raise ValueError("Cannot make a default persona non public")
 
-            if user and user.role != UserRole.ADMIN:
-                raise ValueError("Only admins can make a default persona")
+            if user:
+                # Curators can edit default personas, but not make them
+                if (
+                    user.role == UserRole.CURATOR
+                    or user.role == UserRole.GLOBAL_CURATOR
+                ):
+                    is_default_persona = None
+                elif user.role != UserRole.ADMIN:
+                    raise ValueError("Only admins can make a default persona")
 
         persona = upsert_persona(
             persona_id=persona_id,
@@ -241,7 +249,7 @@ def create_update_persona(
             num_chunks=create_persona_request.num_chunks,
             llm_relevance_filter=create_persona_request.llm_relevance_filter,
             llm_filter_extraction=create_persona_request.llm_filter_extraction,
-            is_default_persona=create_persona_request.is_default_persona,
+            is_default_persona=is_default_persona,
         )
 
         versioned_make_persona_private = fetch_versioned_implementation(
@@ -428,7 +436,7 @@ def upsert_persona(
     remove_image: bool | None = None,
     search_start_date: datetime | None = None,
     builtin_persona: bool = False,
-    is_default_persona: bool = False,
+    is_default_persona: bool | None = None,
     label_ids: list[int] | None = None,
     chunks_above: int = CONTEXT_CHUNKS_ABOVE,
     chunks_below: int = CONTEXT_CHUNKS_BELOW,
@@ -523,7 +531,11 @@ def upsert_persona(
         existing_persona.is_visible = is_visible
         existing_persona.search_start_date = search_start_date
         existing_persona.labels = labels or []
-        existing_persona.is_default_persona = is_default_persona
+        existing_persona.is_default_persona = (
+            is_default_persona
+            if is_default_persona is not None
+            else existing_persona.is_default_persona
+        )
         # Do not delete any associations manually added unless
         # a new updated list is provided
         if document_sets is not None:
@@ -575,7 +587,9 @@ def upsert_persona(
             display_priority=display_priority,
             is_visible=is_visible,
             search_start_date=search_start_date,
-            is_default_persona=is_default_persona,
+            is_default_persona=is_default_persona
+            if is_default_persona is not None
+            else False,
             labels=labels or [],
         )
         db_session.add(new_persona)

backend/onyx/db/search_settings.py

@@ -13,12 +13,13 @@ from onyx.configs.model_configs import OLD_DEFAULT_DOCUMENT_ENCODER_MODEL
 from onyx.configs.model_configs import OLD_DEFAULT_MODEL_DOC_EMBEDDING_DIM
 from onyx.configs.model_configs import OLD_DEFAULT_MODEL_NORMALIZE_EMBEDDINGS
 from onyx.context.search.models import SavedSearchSettings
-from onyx.db.engine import get_session_with_current_tenant
+from onyx.db.enums import EmbeddingPrecision
 from onyx.db.llm import fetch_embedding_provider
 from onyx.db.models import CloudEmbeddingProvider
 from onyx.db.models import IndexAttempt
 from onyx.db.models import IndexModelStatus
 from onyx.db.models import SearchSettings
+from onyx.db.session import get_session_with_current_tenant
 from onyx.indexing.models import IndexingSetting
 from onyx.natural_language_processing.search_nlp_models import clean_model_name
 from onyx.natural_language_processing.search_nlp_models import warm_up_cross_encoder
@@ -59,12 +60,15 @@ def create_search_settings(
         index_name=search_settings.index_name,
         provider_type=search_settings.provider_type,
         multipass_indexing=search_settings.multipass_indexing,
+        embedding_precision=search_settings.embedding_precision,
+        reduced_dimension=search_settings.reduced_dimension,
         multilingual_expansion=search_settings.multilingual_expansion,
         disable_rerank_for_streaming=search_settings.disable_rerank_for_streaming,
         rerank_model_name=search_settings.rerank_model_name,
         rerank_provider_type=search_settings.rerank_provider_type,
         rerank_api_key=search_settings.rerank_api_key,
         num_rerank=search_settings.num_rerank,
+        background_reindex_enabled=search_settings.background_reindex_enabled,
     )
 
     db_session.add(embedding_model)
@@ -305,6 +309,7 @@ def get_old_default_embedding_model() -> IndexingSetting:
         model_dim=(
             DOC_EMBEDDING_DIM if is_overridden else OLD_DEFAULT_MODEL_DOC_EMBEDDING_DIM
         ),
+        embedding_precision=(EmbeddingPrecision.FLOAT),
         normalize=(
             NORMALIZE_EMBEDDINGS
             if is_overridden
@@ -322,6 +327,7 @@ def get_new_default_embedding_model() -> IndexingSetting:
     return IndexingSetting(
         model_name=DOCUMENT_ENCODER_MODEL,
         model_dim=DOC_EMBEDDING_DIM,
+        embedding_precision=(EmbeddingPrecision.FLOAT),
         normalize=NORMALIZE_EMBEDDINGS,
         query_prefix=ASYM_QUERY_PREFIX,
         passage_prefix=ASYM_PASSAGE_PREFIX,

backend/onyx/db/session.py

@@ -0,0 +1,110 @@
+import contextlib
+from collections.abc import AsyncGenerator
+from collections.abc import Generator
+from contextlib import asynccontextmanager
+from contextlib import contextmanager
+from typing import ContextManager
+
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import Session
+from sqlalchemy.orm import sessionmaker
+
+from onyx.configs.app_configs import POSTGRES_IDLE_SESSIONS_TIMEOUT
+from onyx.db.engine import get_sqlalchemy_engine
+from onyx.db.session_schema_translate_map import (
+    OnyxSchemaTranslateMapSession as OnyxSession,
+)
+from onyx.db.utils import is_valid_schema_name
+from onyx.utils.logger import setup_logger
+from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
+from shared_configs.contextvars import CURRENT_TENANT_ID_CONTEXTVAR
+from shared_configs.contextvars import get_current_tenant_id
+
+
+logger = setup_logger()
+
+SessionFactory: sessionmaker[Session] | None = None
+
+
+@contextmanager
+def get_session_with_current_tenant() -> Generator[Session, None, None]:
+    tenant_id = get_current_tenant_id()
+
+    with OnyxSession.get_session_with_tenant(tenant_id=tenant_id) as session:
+        yield session
+
+
+# Used in multi tenant mode when need to refer to the shared `public` schema
+@contextmanager
+def get_session_with_shared_schema() -> Generator[Session, None, None]:
+    token = CURRENT_TENANT_ID_CONTEXTVAR.set(POSTGRES_DEFAULT_SCHEMA)
+    with OnyxSession.get_session_with_tenant(
+        tenant_id=POSTGRES_DEFAULT_SCHEMA
+    ) as session:
+        yield session
+    CURRENT_TENANT_ID_CONTEXTVAR.reset(token)
+
+
+def get_session_generator_with_tenant() -> Generator[Session, None, None]:
+    tenant_id = get_current_tenant_id()
+    with OnyxSession.get_session_with_tenant(tenant_id=tenant_id) as session:
+        yield session
+
+
+def get_session_context_manager() -> ContextManager[Session]:
+    """Context manager for database sessions."""
+    return contextlib.contextmanager(get_session_generator_with_tenant)()
+
+
+def get_session_factory() -> sessionmaker[Session]:
+    global SessionFactory
+    if SessionFactory is None:
+        SessionFactory = sessionmaker(bind=get_sqlalchemy_engine())
+    return SessionFactory
+
+
+@contextmanager
+def get_session_with_tenant(*, tenant_id: str | None) -> Generator[Session, None, None]:
+    with OnyxSession.get_session_with_tenant(tenant_id=tenant_id) as session:
+        yield session
+
+
+def get_session() -> Generator[Session, None, None]:
+    yield from OnyxSession.get_session()
+
+
+def get_multi_tenant_session(tenant_id: str) -> Generator[Session, None, None]:
+    yield from OnyxSession.get_multi_tenant_session(tenant_id)
+
+
+def get_single_tenant_session() -> Generator[Session, None, None]:
+    yield from OnyxSession.get_single_tenant_session()
+
+
+async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
+    """Proxy method that simply delegates to `get_async_session`."""
+    async for session in OnyxSession.get_async_session():
+        yield session
+
+
+@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 for session in OnyxSession.get_multi_tenant_async_session(tenant_id):
+        if POSTGRES_IDLE_SESSIONS_TIMEOUT:
+            await session.execute(
+                text(
+                    f"SET idle_in_transaction_session_timeout = {POSTGRES_IDLE_SESSIONS_TIMEOUT}"
+                )
+            )
+
+        yield session

backend/onyx/db/session_schema_translate_map.py

@@ -0,0 +1,140 @@
+"""
+Implements multi-tenant / schema handling for a session via
+SQLAlchemy's schema_translate_map feature.
+
+This is better for us than
+SET search_path because that approach pins the connection in RDS proxy since it
+alters the connection state.
+"""
+from collections.abc import AsyncGenerator
+from collections.abc import Generator
+from contextlib import contextmanager
+
+from fastapi import HTTPException
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import Session
+
+from onyx.configs.app_configs import POSTGRES_IDLE_SESSIONS_TIMEOUT
+from onyx.db.engine import get_sqlalchemy_async_engine
+from onyx.db.engine import get_sqlalchemy_engine
+from onyx.db.utils import is_valid_schema_name
+from onyx.server.utils import BasicAuthenticationError
+from shared_configs.configs import MULTI_TENANT
+from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
+from shared_configs.contextvars import get_current_tenant_id
+
+
+class OnyxSchemaTranslateMapSession:
+    @contextmanager
+    @staticmethod
+    def get_session_with_tenant(
+        *, tenant_id: str | None
+    ) -> Generator[Session, None, None]:
+        """
+        Generate a database session for a specific tenant.
+        """
+        if tenant_id is None:
+            tenant_id = POSTGRES_DEFAULT_SCHEMA
+
+        schema_translate_map = {None: tenant_id}
+
+        engine = get_sqlalchemy_engine()
+
+        if not is_valid_schema_name(tenant_id):
+            raise HTTPException(status_code=400, detail="Invalid tenant ID")
+
+        with engine.connect().execution_options(
+            schema_translate_map=schema_translate_map
+        ) as connection:
+            dbapi_connection = connection.connection
+            if POSTGRES_IDLE_SESSIONS_TIMEOUT:
+                try:
+                    cursor = dbapi_connection.cursor()
+                    cursor.execute(
+                        text(
+                            f"SET SESSION idle_in_transaction_session_timeout = {POSTGRES_IDLE_SESSIONS_TIMEOUT}"
+                        )
+                    )
+                finally:
+                    cursor.close()
+
+            with Session(bind=connection, expire_on_commit=False) as session:
+                yield session
+
+    @staticmethod
+    def get_session() -> Generator[Session, None, None]:
+        if MULTI_TENANT:
+            tenant_id = get_current_tenant_id()
+            yield from OnyxSchemaTranslateMapSession.get_multi_tenant_session(tenant_id)
+            return
+
+        yield from OnyxSchemaTranslateMapSession.get_single_tenant_session()
+        return
+
+    @staticmethod
+    def get_multi_tenant_session(tenant_id: str) -> Generator[Session, None, None]:
+        schema_translate_map = {None: 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")
+
+        engine = get_sqlalchemy_engine()
+        with engine.connect().execution_options(
+            schema_translate_map=schema_translate_map
+        ) as connection:
+            with Session(bind=connection, expire_on_commit=False) as session:
+                yield session
+
+    @staticmethod
+    def get_single_tenant_session() -> Generator[Session, None, None]:
+        engine = get_sqlalchemy_engine()
+
+        # single tenant
+        with engine.connect() as connection:
+            with Session(bind=connection, expire_on_commit=False) as session:
+                yield session
+
+    @staticmethod
+    async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
+        if MULTI_TENANT:
+            tenant_id = get_current_tenant_id()
+            async for session in OnyxSchemaTranslateMapSession.get_multi_tenant_async_session(
+                tenant_id
+            ):
+                yield session
+            return
+
+        async for session in OnyxSchemaTranslateMapSession.get_single_tenant_async_session():
+            yield session
+
+    @staticmethod
+    async def get_multi_tenant_async_session(
+        tenant_id: str,
+    ) -> AsyncGenerator[AsyncSession, None]:
+        engine = get_sqlalchemy_async_engine()
+
+        if not is_valid_schema_name(tenant_id):
+            raise HTTPException(status_code=400, detail="Invalid tenant ID")
+
+        # Create connection with schema translation
+        schema_translate_map = {None: tenant_id}
+        async with engine.connect() as connection:
+            connection = await connection.execution_options(
+                schema_translate_map=schema_translate_map
+            )
+            async with AsyncSession(
+                bind=connection, expire_on_commit=False
+            ) as async_session:
+                yield async_session
+
+    @staticmethod
+    async def get_single_tenant_async_session() -> AsyncGenerator[AsyncSession, None]:
+        engine = get_sqlalchemy_async_engine()
+
+        # single tenant
+        async with AsyncSession(engine, expire_on_commit=False) as async_session:
+            yield async_session

backend/onyx/db/session_search_path.py

@@ -0,0 +1,185 @@
+"""
+Implements multi-tenant / schema handling for a session via "SET search_path".
+
+This is worse for us than schema_translate_map because this approach pins the connection
+ in RDS proxy since it alters the connection state.
+
+Keeping this approach here while we test/iterate.
+"""
+from collections.abc import AsyncGenerator
+from collections.abc import Generator
+from contextlib import asynccontextmanager
+from contextlib import contextmanager
+from typing import Any
+
+from fastapi import HTTPException
+from sqlalchemy import event
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import AsyncSession
+from sqlalchemy.orm import Session
+from sqlalchemy.orm import sessionmaker
+
+from onyx.configs.app_configs import POSTGRES_IDLE_SESSIONS_TIMEOUT
+from onyx.db.engine import get_sqlalchemy_async_engine
+from onyx.db.engine import get_sqlalchemy_engine
+from onyx.db.utils import is_valid_schema_name
+from onyx.server.utils import BasicAuthenticationError
+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.contextvars import get_current_tenant_id
+
+
+logger = setup_logger()
+
+
+AsyncSessionLocal = sessionmaker(  # type: ignore
+    bind=get_sqlalchemy_async_engine(),
+    class_=AsyncSession,
+    expire_on_commit=False,
+)
+
+
+class OnyxSearchPathSession:
+    @staticmethod
+    def _set_search_path_on_checkout(
+        dbapi_conn: Any, connection_record: Any, connection_proxy: Any
+    ) -> None:
+        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
+    @staticmethod
+    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", OnyxSearchPathSession._set_search_path_on_checkout
+        )
+
+        if not is_valid_schema_name(tenant_id):
+            raise HTTPException(status_code=400, detail="Invalid tenant ID")
+
+        with engine.connect() as connection:
+            dbapi_connection = connection.connection
+            cursor = dbapi_connection.cursor()
+            try:
+                cursor.execute(f'SET search_path = "{tenant_id}"')
+                if POSTGRES_IDLE_SESSIONS_TIMEOUT:
+                    cursor.execute(
+                        text(
+                            f"SET SESSION idle_in_transaction_session_timeout = {POSTGRES_IDLE_SESSIONS_TIMEOUT}"
+                        )
+                    )
+            finally:
+                cursor.close()
+
+            with Session(bind=connection, expire_on_commit=False) as session:
+                try:
+                    yield session
+                finally:
+                    if MULTI_TENANT:
+                        cursor = dbapi_connection.cursor()
+                        try:
+                            cursor.execute('SET search_path TO "$user", public')
+                        finally:
+                            cursor.close()
+
+    @staticmethod
+    def get_session() -> Generator[Session, None, None]:
+        if MULTI_TENANT:
+            tenant_id = get_current_tenant_id()
+            yield from OnyxSearchPathSession.get_multi_tenant_session(tenant_id)
+            return
+
+        yield from OnyxSearchPathSession.get_single_tenant_session()
+        return
+
+    @staticmethod
+    def get_multi_tenant_session(tenant_id: str) -> Generator[Session, None, None]:
+        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")
+
+        engine = get_sqlalchemy_engine()
+        with Session(engine, expire_on_commit=False) as session:
+            session.execute(text(f'SET search_path = "{tenant_id}"'))
+            yield session
+
+    @staticmethod
+    def get_single_tenant_session() -> Generator[Session, None, None]:
+        engine = get_sqlalchemy_engine()
+
+        with Session(engine, expire_on_commit=False) as session:
+            yield session
+
+    @staticmethod
+    async def get_async_session() -> AsyncGenerator[AsyncSession, None]:
+        if MULTI_TENANT:
+            tenant_id = get_current_tenant_id()
+            async for session in OnyxSearchPathSession.get_multi_tenant_async_session(
+                tenant_id
+            ):
+                yield session
+            return
+
+        async for session in OnyxSearchPathSession.get_single_tenant_async_session():
+            yield session
+
+    @staticmethod
+    async def get_multi_tenant_async_session(
+        tenant_id: str,
+    ) -> AsyncGenerator[AsyncSession, None]:
+        engine = get_sqlalchemy_async_engine()
+
+        if not is_valid_schema_name(tenant_id):
+            raise HTTPException(status_code=400, detail="Invalid tenant ID")
+
+        async with AsyncSession(engine, expire_on_commit=False) as async_session:
+            await async_session.execute(text(f'SET search_path = "{tenant_id}"'))
+            yield async_session
+
+    @staticmethod
+    async def get_single_tenant_async_session() -> AsyncGenerator[AsyncSession, None]:
+        engine = get_sqlalchemy_async_engine()
+
+        # single tenant
+        async with AsyncSession(engine, expire_on_commit=False) as async_session:
+            yield async_session
+
+    @asynccontextmanager
+    @staticmethod
+    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

backend/onyx/db/swap_index.py

@@ -8,10 +8,12 @@ from onyx.db.index_attempt import cancel_indexing_attempts_past_model
 from onyx.db.index_attempt import (
     count_unique_cc_pairs_with_successful_index_attempts,
 )
+from onyx.db.models import ConnectorCredentialPair
 from onyx.db.models import SearchSettings
 from onyx.db.search_settings import get_current_search_settings
 from onyx.db.search_settings import get_secondary_search_settings
 from onyx.db.search_settings import update_search_settings_status
+from onyx.document_index.factory import get_default_document_index
 from onyx.key_value_store.factory import get_kv_store
 from onyx.utils.logger import setup_logger
 
@@ -19,7 +21,49 @@ from onyx.utils.logger import setup_logger
 logger = setup_logger()
 
 
-def check_index_swap(db_session: Session) -> SearchSettings | None:
+def _perform_index_swap(
+    db_session: Session,
+    current_search_settings: SearchSettings,
+    secondary_search_settings: SearchSettings,
+    all_cc_pairs: list[ConnectorCredentialPair],
+) -> None:
+    """Swap the indices and expire the old one."""
+    current_search_settings = get_current_search_settings(db_session)
+    update_search_settings_status(
+        search_settings=current_search_settings,
+        new_status=IndexModelStatus.PAST,
+        db_session=db_session,
+    )
+
+    update_search_settings_status(
+        search_settings=secondary_search_settings,
+        new_status=IndexModelStatus.PRESENT,
+        db_session=db_session,
+    )
+
+    if len(all_cc_pairs) > 0:
+        kv_store = get_kv_store()
+        kv_store.store(KV_REINDEX_KEY, False)
+
+        # Expire jobs for the now past index/embedding model
+        cancel_indexing_attempts_past_model(db_session)
+
+        # Recount aggregates
+        for cc_pair in all_cc_pairs:
+            resync_cc_pair(cc_pair, db_session=db_session)
+
+    # remove the old index from the vector db
+    document_index = get_default_document_index(secondary_search_settings, None)
+    document_index.ensure_indices_exist(
+        primary_embedding_dim=secondary_search_settings.final_embedding_dim,
+        primary_embedding_precision=secondary_search_settings.embedding_precision,
+        # just finished swap, no more secondary index
+        secondary_index_embedding_dim=None,
+        secondary_index_embedding_precision=None,
+    )
+
+
+def check_and_perform_index_swap(db_session: Session) -> SearchSettings | None:
     """Get count of cc-pairs and count of successful index_attempts for the
     new model grouped by connector + credential, if it's the same, then assume
     new index is done building. If so, swap the indices and expire the old one.
@@ -27,52 +71,45 @@ def check_index_swap(db_session: Session) -> SearchSettings | None:
     Returns None if search settings did not change, or the old search settings if they
     did change.
     """
-
-    old_search_settings = None
-
     # Default CC-pair created for Ingestion API unused here
     all_cc_pairs = get_connector_credential_pairs(db_session)
     cc_pair_count = max(len(all_cc_pairs) - 1, 0)
-    search_settings = get_secondary_search_settings(db_session)
+    secondary_search_settings = get_secondary_search_settings(db_session)
 
-    if not search_settings:
+    if not secondary_search_settings:
         return None
 
+    # If the secondary search settings are not configured to reindex in the background,
+    # we can just swap over instantly
+    if not secondary_search_settings.background_reindex_enabled:
+        current_search_settings = get_current_search_settings(db_session)
+        _perform_index_swap(
+            db_session=db_session,
+            current_search_settings=current_search_settings,
+            secondary_search_settings=secondary_search_settings,
+            all_cc_pairs=all_cc_pairs,
+        )
+        return current_search_settings
+
     unique_cc_indexings = count_unique_cc_pairs_with_successful_index_attempts(
-        search_settings_id=search_settings.id, db_session=db_session
+        search_settings_id=secondary_search_settings.id, db_session=db_session
     )
 
     # Index Attempts are cleaned up as well when the cc-pair is deleted so the logic in this
     # function is correct. The unique_cc_indexings are specifically for the existing cc-pairs
+    old_search_settings = None
     if unique_cc_indexings > cc_pair_count:
         logger.error("More unique indexings than cc pairs, should not occur")
 
     if cc_pair_count == 0 or cc_pair_count == unique_cc_indexings:
         # Swap indices
         current_search_settings = get_current_search_settings(db_session)
-        update_search_settings_status(
-            search_settings=current_search_settings,
-            new_status=IndexModelStatus.PAST,
+        _perform_index_swap(
             db_session=db_session,
+            current_search_settings=current_search_settings,
+            secondary_search_settings=secondary_search_settings,
+            all_cc_pairs=all_cc_pairs,
         )
-
-        update_search_settings_status(
-            search_settings=search_settings,
-            new_status=IndexModelStatus.PRESENT,
-            db_session=db_session,
-        )
-
-        if cc_pair_count > 0:
-            kv_store = get_kv_store()
-            kv_store.store(KV_REINDEX_KEY, False)
-
-            # Expire jobs for the now past index/embedding model
-            cancel_indexing_attempts_past_model(db_session)
-
-            # Recount aggregates
-            for cc_pair in all_cc_pairs:
-                resync_cc_pair(cc_pair, db_session=db_session)
-
-            old_search_settings = current_search_settings
+        old_search_settings = current_search_settings
 
     return old_search_settings

backend/onyx/db/tenant.py

@@ -0,0 +1,31 @@
+from sqlalchemy import text
+
+from onyx.db.session import get_session_with_shared_schema
+from shared_configs.configs import MULTI_TENANT
+from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
+from shared_configs.configs import TENANT_ID_PREFIX
+
+
+def get_all_tenant_ids() -> list[str]:
+    """Returning [None] means the only tenant is the 'public' or self hosted tenant."""
+
+    if not MULTI_TENANT:
+        return [POSTGRES_DEFAULT_SCHEMA]
+
+    with get_session_with_shared_schema() as session:
+        result = session.execute(
+            text(
+                f"""
+                SELECT schema_name
+                FROM information_schema.schemata
+                WHERE schema_name NOT IN ('pg_catalog', 'information_schema', '{POSTGRES_DEFAULT_SCHEMA}')"""
+            )
+        )
+        tenant_ids = [row[0] for row in result]
+
+    valid_tenants = [
+        tenant
+        for tenant in tenant_ids
+        if tenant is None or tenant.startswith(TENANT_ID_PREFIX)
+    ]
+    return valid_tenants