remove oauth.py again

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

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/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/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/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/api.py

@@ -12,10 +12,10 @@ 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.engine import get_current_tenant_id
 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()
 

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

@@ -25,12 +25,12 @@ 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.engine import get_current_tenant_id
-from onyx.db.engine import get_session
 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()
 

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

@@ -33,11 +33,11 @@ from onyx.connectors.google_utils.shared_constants import (
     GoogleOAuthAuthenticationMethod,
 )
 from onyx.db.credentials import create_credential
-from onyx.db.engine import get_current_tenant_id
-from onyx.db.engine import get_session
 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:

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

@@ -17,11 +17,11 @@ 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.engine import get_current_tenant_id
-from onyx.db.engine import get_session
 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:

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

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/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

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/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

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

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

@@ -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,6 +60,7 @@ 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.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

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

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(

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

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/credentials_provider.py

@@ -6,8 +6,8 @@ from redis.lock import Lock as RedisLock
 from sqlalchemy import select
 
 from onyx.connectors.interfaces import CredentialsProviderInterface
-from onyx.db.engine import get_session_with_tenant
 from onyx.db.models import Credential
+from onyx.db.session import get_session_with_tenant
 from onyx.redis.redis_pool import get_redis_client
 
 

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/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/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/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/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/search_settings.py

@@ -13,13 +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

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/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

backend/onyx/db/utils.py

@@ -1,3 +1,4 @@
+import re
 from typing import Any
 
 from sqlalchemy import inspect
@@ -5,5 +6,12 @@ from sqlalchemy import inspect
 from onyx.db.models import Base
 
 
+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
+
+
 def model_to_dict(model: Base) -> dict[str, Any]:
     return {c.key: getattr(model, c.key) for c in inspect(model).mapper.column_attrs}  # type: ignore

backend/onyx/file_store/utils.py

@@ -8,8 +8,8 @@ import requests
 from sqlalchemy.orm import Session
 
 from onyx.configs.constants import FileOrigin
-from onyx.db.engine import get_session_with_current_tenant
 from onyx.db.models import ChatMessage
+from onyx.db.session import get_session_with_current_tenant
 from onyx.file_store.file_store import get_default_file_store
 from onyx.file_store.models import FileDescriptor
 from onyx.file_store.models import InMemoryChatFile

backend/onyx/key_value_store/store.py

@@ -5,12 +5,11 @@ from typing import cast
 
 from fastapi import HTTPException
 from redis.client import Redis
-from sqlalchemy import text
 from sqlalchemy.orm import Session
 
-from onyx.db.engine import get_sqlalchemy_engine
-from onyx.db.engine import is_valid_schema_name
 from onyx.db.models import KVStore
+from onyx.db.session import get_multi_tenant_session
+from onyx.db.session import get_single_tenant_session
 from onyx.key_value_store.interface import KeyValueStore
 from onyx.key_value_store.interface import KvKeyNotFoundError
 from onyx.redis.redis_pool import get_redis_client
@@ -18,7 +17,6 @@ from onyx.server.utils import BasicAuthenticationError
 from onyx.utils.logger import setup_logger
 from onyx.utils.special_types import JSON_ro
 from shared_configs.configs import MULTI_TENANT
-from shared_configs.configs import POSTGRES_DEFAULT_SCHEMA
 from shared_configs.contextvars import get_current_tenant_id
 
 
@@ -41,16 +39,13 @@ class PgRedisKVStore(KeyValueStore):
 
     @contextmanager
     def _get_session(self) -> Iterator[Session]:
-        engine = get_sqlalchemy_engine()
-        with Session(engine, expire_on_commit=False) as session:
+        try:
             if MULTI_TENANT:
-                if self.tenant_id == POSTGRES_DEFAULT_SCHEMA:
-                    raise BasicAuthenticationError(detail="User must authenticate")
-                if not is_valid_schema_name(self.tenant_id):
-                    raise HTTPException(status_code=400, detail="Invalid tenant ID")
-                # Set the search_path to the tenant's schema
-                session.execute(text(f'SET search_path = "{self.tenant_id}"'))
-            yield session
+                yield from get_multi_tenant_session(self.tenant_id)
+            else:
+                yield from get_single_tenant_session()
+        except BasicAuthenticationError as e:
+            raise HTTPException(status_code=401, detail=str(e))
 
     def store(self, key: str, val: JSON_ro, encrypt: bool = False) -> None:
         # Not encrypted in Redis, but encrypted in Postgres

backend/onyx/llm/factory.py

@@ -4,10 +4,10 @@ from onyx.chat.models import PersonaOverrideConfig
 from onyx.configs.app_configs import DISABLE_GENERATIVE_AI
 from onyx.configs.model_configs import GEN_AI_MODEL_FALLBACK_MAX_TOKENS
 from onyx.configs.model_configs import GEN_AI_TEMPERATURE
-from onyx.db.engine import get_session_context_manager
 from onyx.db.llm import fetch_default_provider
 from onyx.db.llm import fetch_provider
 from onyx.db.models import Persona
+from onyx.db.session import get_session_context_manager
 from onyx.llm.chat_llm import DefaultMultiLLM
 from onyx.llm.exceptions import GenAIDisabledException
 from onyx.llm.interfaces import LLM

backend/onyx/onyxbot/slack/blocks.py

@@ -23,8 +23,8 @@ from onyx.configs.constants import SearchFeedbackType
 from onyx.configs.onyxbot_configs import DANSWER_BOT_NUM_DOCS_TO_DISPLAY
 from onyx.context.search.models import SavedSearchDoc
 from onyx.db.chat import get_chat_session_by_message_id
-from onyx.db.engine import get_session_with_current_tenant
 from onyx.db.models import ChannelConfig
+from onyx.db.session import get_session_with_current_tenant
 from onyx.onyxbot.slack.constants import CONTINUE_IN_WEB_UI_ACTION_ID
 from onyx.onyxbot.slack.constants import DISLIKE_BLOCK_ACTION_ID
 from onyx.onyxbot.slack.constants import FEEDBACK_DOC_BUTTON_BLOCK_ACTION_ID

backend/onyx/onyxbot/slack/handlers/handle_buttons.py

@@ -19,9 +19,9 @@ from onyx.connectors.slack.utils import make_slack_api_rate_limited
 from onyx.context.search.models import SavedSearchDoc
 from onyx.db.chat import get_chat_message
 from onyx.db.chat import translate_db_message_to_chat_message_detail
-from onyx.db.engine import get_session_with_current_tenant
 from onyx.db.feedback import create_chat_message_feedback
 from onyx.db.feedback import create_doc_retrieval_feedback
+from onyx.db.session import get_session_with_current_tenant
 from onyx.db.users import get_user_by_email
 from onyx.onyxbot.slack.blocks import build_follow_up_resolved_blocks
 from onyx.onyxbot.slack.blocks import build_slack_response_blocks

backend/onyx/onyxbot/slack/handlers/handle_message.py

@@ -5,8 +5,8 @@ from slack_sdk.errors import SlackApiError
 
 from onyx.configs.onyxbot_configs import DANSWER_BOT_FEEDBACK_REMINDER
 from onyx.configs.onyxbot_configs import DANSWER_REACT_EMOJI
-from onyx.db.engine import get_session_with_current_tenant
 from onyx.db.models import SlackChannelConfig
+from onyx.db.session import get_session_with_current_tenant
 from onyx.db.users import add_slack_user_if_not_exists
 from onyx.onyxbot.slack.blocks import get_feedback_reminder_blocks
 from onyx.onyxbot.slack.handlers.handle_regular_answer import (

backend/onyx/onyxbot/slack/handlers/handle_regular_answer.py

@@ -23,11 +23,11 @@ from onyx.configs.onyxbot_configs import MAX_THREAD_CONTEXT_PERCENTAGE
 from onyx.context.search.enums import OptionalSearchSetting
 from onyx.context.search.models import BaseFilters
 from onyx.context.search.models import RetrievalDetails
-from onyx.db.engine import get_session_with_current_tenant
 from onyx.db.models import SlackChannelConfig
 from onyx.db.models import User
 from onyx.db.persona import get_persona_by_id
 from onyx.db.persona import persona_has_search_tool
+from onyx.db.session import get_session_with_current_tenant
 from onyx.db.users import get_user_by_email
 from onyx.onyxbot.slack.blocks import build_slack_response_blocks
 from onyx.onyxbot.slack.handlers.utils import send_team_member_message

backend/onyx/onyxbot/slack/listener.py

@@ -36,12 +36,12 @@ from onyx.connectors.slack.utils import expert_info_from_slack_id
 from onyx.context.search.retrieval.search_runner import (
     download_nltk_data,
 )
-from onyx.db.engine import get_all_tenant_ids
-from onyx.db.engine import get_session_with_current_tenant
-from onyx.db.engine import get_session_with_tenant
 from onyx.db.models import SlackBot
 from onyx.db.search_settings import get_current_search_settings
+from onyx.db.session import get_session_with_current_tenant
+from onyx.db.session import get_session_with_tenant
 from onyx.db.slack_bot import fetch_slack_bots
+from onyx.db.tenant import get_all_tenant_ids
 from onyx.key_value_store.interface import KvKeyNotFoundError
 from onyx.natural_language_processing.search_nlp_models import EmbeddingModel
 from onyx.natural_language_processing.search_nlp_models import warm_up_bi_encoder

backend/onyx/onyxbot/slack/utils.py

@@ -32,7 +32,7 @@ from onyx.configs.onyxbot_configs import (
 )
 from onyx.connectors.slack.utils import make_slack_api_rate_limited
 from onyx.connectors.slack.utils import SlackTextCleaner
-from onyx.db.engine import get_session_with_current_tenant
+from onyx.db.session import get_session_with_current_tenant
 from onyx.db.users import get_user_by_email
 from onyx.llm.exceptions import GenAIDisabledException
 from onyx.llm.factory import get_default_llms

backend/onyx/server/api_key/api.py

@@ -9,8 +9,8 @@ from onyx.db.api_key import insert_api_key
 from onyx.db.api_key import regenerate_api_key
 from onyx.db.api_key import remove_api_key
 from onyx.db.api_key import update_api_key
-from onyx.db.engine import get_session
 from onyx.db.models import User
+from onyx.db.session import get_session
 from onyx.server.api_key.models import APIKeyArgs
 
 

backend/onyx/server/documents/cc_pair.py

@@ -38,7 +38,6 @@ from onyx.db.connector_credential_pair import (
 )
 from onyx.db.document import get_document_counts_for_cc_pairs
 from onyx.db.document import get_documents_for_cc_pair
-from onyx.db.engine import get_session
 from onyx.db.enums import AccessType
 from onyx.db.enums import ConnectorCredentialPairStatus
 from onyx.db.index_attempt import count_index_attempt_errors_for_cc_pair
@@ -50,6 +49,7 @@ from onyx.db.models import SearchSettings
 from onyx.db.models import User
 from onyx.db.search_settings import get_active_search_settings_list
 from onyx.db.search_settings import get_current_search_settings
+from onyx.db.session import get_session
 from onyx.redis.redis_connector import RedisConnector
 from onyx.redis.redis_pool import get_redis_client
 from onyx.server.documents.models import CCPairFullInfo
@@ -651,6 +651,8 @@ def associate_credential_to_connector(
         # Ensures we don't leave invalid connectors in the database
         # NOTE: consensus is that it makes sense to unify connector and ccpair creation flows
         # which would rid us of needing to handle cases like these
+        logger.exception("ValidationError")
+
         delete_connector(db_session, connector_id)
         db_session.commit()
 

backend/onyx/server/documents/connector.py

@@ -85,8 +85,6 @@ from onyx.db.credentials import delete_service_account_credentials
 from onyx.db.credentials import fetch_credential_by_id_for_user
 from onyx.db.deletion_attempt import check_deletion_attempt_is_allowed
 from onyx.db.document import get_document_counts_for_cc_pairs_parallel
-from onyx.db.engine import get_current_tenant_id
-from onyx.db.engine import get_session
 from onyx.db.enums import AccessType
 from onyx.db.enums import IndexingMode
 from onyx.db.index_attempt import get_index_attempts_for_cc_pair
@@ -99,6 +97,7 @@ from onyx.db.models import User
 from onyx.db.models import UserGroup__ConnectorCredentialPair
 from onyx.db.search_settings import get_current_search_settings
 from onyx.db.search_settings import get_secondary_search_settings
+from onyx.db.session import get_session
 from onyx.file_processing.extract_file_text import convert_docx_to_txt
 from onyx.file_store.file_store import get_default_file_store
 from onyx.key_value_store.interface import KvKeyNotFoundError
@@ -127,6 +126,7 @@ from onyx.utils.logger import setup_logger
 from onyx.utils.telemetry import create_milestone_and_report
 from onyx.utils.threadpool_concurrency import run_functions_tuples_in_parallel
 from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop
+from shared_configs.contextvars import get_current_tenant_id
 
 logger = setup_logger()
 

backend/onyx/server/documents/credential.py

@@ -19,9 +19,9 @@ from onyx.db.credentials import fetch_credentials_by_source_for_user
 from onyx.db.credentials import fetch_credentials_for_user
 from onyx.db.credentials import swap_credentials_connector
 from onyx.db.credentials import update_credential
-from onyx.db.engine import get_session
 from onyx.db.models import DocumentSource
 from onyx.db.models import User
+from onyx.db.session import get_session
 from onyx.server.documents.models import CredentialBase
 from onyx.server.documents.models import CredentialDataUpdateRequest
 from onyx.server.documents.models import CredentialSnapshot

backend/onyx/server/documents/document.py

@@ -9,9 +9,9 @@ from onyx.context.search.models import IndexFilters
 from onyx.context.search.preprocessing.access_filters import (
     build_access_filters_for_user,
 )
-from onyx.db.engine import get_session
 from onyx.db.models import User
 from onyx.db.search_settings import get_current_search_settings
+from onyx.db.session import get_session
 from onyx.document_index.factory import get_default_document_index
 from onyx.document_index.interfaces import VespaChunkRequest
 from onyx.natural_language_processing.utils import get_tokenizer

backend/onyx/server/documents/standard_oauth.py

@@ -17,8 +17,8 @@ from onyx.configs.app_configs import WEB_DOMAIN
 from onyx.configs.constants import DocumentSource
 from onyx.connectors.interfaces import OAuthConnector
 from onyx.db.credentials import create_credential
-from onyx.db.engine import get_session
 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

backend/onyx/server/features/document_set/api.py

@@ -14,16 +14,18 @@ from onyx.db.document_set import fetch_all_document_sets_for_user
 from onyx.db.document_set import insert_document_set
 from onyx.db.document_set import mark_document_set_as_to_be_deleted
 from onyx.db.document_set import update_document_set
-from onyx.db.engine import get_current_tenant_id
-from onyx.db.engine import get_session
 from onyx.db.models import User
+from onyx.db.session import get_session
 from onyx.server.features.document_set.models import CheckDocSetPublicRequest
 from onyx.server.features.document_set.models import CheckDocSetPublicResponse
 from onyx.server.features.document_set.models import DocumentSet
 from onyx.server.features.document_set.models import DocumentSetCreationRequest
 from onyx.server.features.document_set.models import DocumentSetUpdateRequest
+from onyx.utils.logger import setup_logger
 from onyx.utils.variable_functionality import fetch_ee_implementation_or_noop
+from shared_configs.contextvars import get_current_tenant_id
 
+logger = setup_logger()
 
 router = APIRouter(prefix="/manage")
 
@@ -83,6 +85,7 @@ def patch_document_set(
             user=user,
         )
     except Exception as e:
+        logger.exception("patch_document_set exceptioned")
         raise HTTPException(status_code=400, detail=str(e))
 
     primary_app.send_task(

backend/onyx/server/features/folder/api.py

@@ -6,7 +6,6 @@ from sqlalchemy.orm import Session
 
 from onyx.auth.users import current_user
 from onyx.db.chat import get_chat_session_by_id
-from onyx.db.engine import get_session
 from onyx.db.folder import add_chat_to_folder
 from onyx.db.folder import create_folder
 from onyx.db.folder import delete_folder
@@ -15,6 +14,7 @@ from onyx.db.folder import remove_chat_from_folder
 from onyx.db.folder import rename_folder
 from onyx.db.folder import update_folder_display_priority
 from onyx.db.models import User
+from onyx.db.session import get_session
 from onyx.server.features.folder.models import DeleteFolderOptions
 from onyx.server.features.folder.models import FolderChatSessionRequest
 from onyx.server.features.folder.models import FolderCreationRequest

backend/onyx/server/features/input_prompt/api.py

@@ -5,7 +5,6 @@ from sqlalchemy.orm import Session
 
 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.input_prompt import disable_input_prompt_for_user
 from onyx.db.input_prompt import fetch_input_prompt_by_id
 from onyx.db.input_prompt import fetch_input_prompts_by_user
@@ -15,6 +14,7 @@ from onyx.db.input_prompt import remove_public_input_prompt
 from onyx.db.input_prompt import update_input_prompt
 from onyx.db.models import InputPrompt__User
 from onyx.db.models import User
+from onyx.db.session import get_session
 from onyx.server.features.input_prompt.models import CreateInputPromptRequest
 from onyx.server.features.input_prompt.models import InputPromptSnapshot
 from onyx.server.features.input_prompt.models import UpdateInputPromptRequest

backend/onyx/server/features/notifications/api.py

@@ -4,11 +4,11 @@ from fastapi import HTTPException
 from sqlalchemy.orm import Session
 
 from onyx.auth.users import current_user
-from onyx.db.engine import get_session
 from onyx.db.models import User
 from onyx.db.notification import dismiss_notification
 from onyx.db.notification import get_notification_by_id
 from onyx.db.notification import get_notifications
+from onyx.db.session import get_session
 from onyx.server.settings.models import Notification as NotificationModel
 from onyx.utils.logger import setup_logger
 

backend/onyx/server/features/password/api.py

@@ -9,7 +9,7 @@ from onyx.auth.users import current_user
 from onyx.auth.users import get_user_manager
 from onyx.auth.users import User
 from onyx.auth.users import UserManager
-from onyx.db.engine import get_session
+from onyx.db.session import get_session
 from onyx.db.users import get_user_by_email
 from onyx.server.features.password.models import ChangePasswordRequest
 from onyx.server.features.password.models import UserResetRequest

backend/onyx/server/features/persona/api.py

@@ -18,7 +18,6 @@ from onyx.auth.users import current_user
 from onyx.configs.constants import FileOrigin
 from onyx.configs.constants import MilestoneRecordType
 from onyx.configs.constants import NotificationType
-from onyx.db.engine import get_session
 from onyx.db.models import StarterMessageModel as StarterMessage
 from onyx.db.models import User
 from onyx.db.notification import create_notification
@@ -38,6 +37,7 @@ from onyx.db.persona import update_persona_shared_users
 from onyx.db.persona import update_persona_visibility
 from onyx.db.prompts import build_prompt_name_from_persona_name
 from onyx.db.prompts import upsert_prompt
+from onyx.db.session import get_session
 from onyx.file_store.file_store import get_default_file_store
 from onyx.file_store.models import ChatFileType
 from onyx.secondary_llm_flows.starter_message_creation import (

backend/onyx/server/features/tool/api.py

@@ -8,8 +8,8 @@ from sqlalchemy.orm import Session
 
 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
 from onyx.db.tools import create_tool
 from onyx.db.tools import delete_tool
 from onyx.db.tools import get_tool_by_id

backend/onyx/server/gpts/api.py

@@ -8,8 +8,8 @@ from sqlalchemy.orm import Session
 
 from onyx.context.search.models import SearchRequest
 from onyx.context.search.pipeline import SearchPipeline
-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_default_llms
 from onyx.server.onyx_api.ingestion import api_key_dep
 from onyx.utils.logger import setup_logger

backend/onyx/server/manage/administrative.py

@@ -20,13 +20,13 @@ from onyx.db.connector_credential_pair import get_connector_credential_pair_for_
 from onyx.db.connector_credential_pair import (
     update_connector_credential_pair_from_id,
 )
-from onyx.db.engine import get_session
 from onyx.db.enums import ConnectorCredentialPairStatus
 from onyx.db.feedback import fetch_docs_ranked_by_boost_for_user
 from onyx.db.feedback import update_document_boost_for_user
 from onyx.db.feedback import update_document_hidden_for_user
 from onyx.db.index_attempt import cancel_indexing_attempts_for_ccpair
 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.key_value_store.factory import get_kv_store
 from onyx.key_value_store.interface import KvKeyNotFoundError

backend/onyx/server/manage/embedding/api.py

@@ -4,13 +4,13 @@ from fastapi import HTTPException
 from sqlalchemy.orm import Session
 
 from onyx.auth.users import current_admin_user
-from onyx.db.engine import get_session
 from onyx.db.llm import fetch_existing_embedding_providers
 from onyx.db.llm import remove_embedding_provider
 from onyx.db.llm import upsert_cloud_embedding_provider
 from onyx.db.models import User
 from onyx.db.search_settings import get_all_search_settings
 from onyx.db.search_settings import get_current_db_embedding_provider
+from onyx.db.session import get_session
 from onyx.indexing.models import EmbeddingModelDetail
 from onyx.natural_language_processing.search_nlp_models import EmbeddingModel
 from onyx.server.manage.embedding.models import CloudEmbeddingProvider

backend/onyx/server/manage/llm/api.py

@@ -8,7 +8,6 @@ from sqlalchemy.orm import Session
 
 from onyx.auth.users import current_admin_user
 from onyx.auth.users import current_chat_accesssible_user
-from onyx.db.engine import get_session
 from onyx.db.llm import fetch_existing_llm_providers
 from onyx.db.llm import fetch_existing_llm_providers_for_user
 from onyx.db.llm import fetch_provider
@@ -16,6 +15,7 @@ from onyx.db.llm import remove_llm_provider
 from onyx.db.llm import update_default_provider
 from onyx.db.llm import upsert_llm_provider
 from onyx.db.models import User
+from onyx.db.session import get_session
 from onyx.llm.factory import get_default_llms
 from onyx.llm.factory import get_llm
 from onyx.llm.llm_provider_options import fetch_available_well_known_llms

backend/onyx/server/manage/search_settings.py

@@ -11,7 +11,6 @@ from onyx.context.search.models import SavedSearchSettings
 from onyx.context.search.models import SearchSettingsCreationRequest
 from onyx.db.connector_credential_pair import get_connector_credential_pairs
 from onyx.db.connector_credential_pair import resync_cc_pair
-from onyx.db.engine import get_session
 from onyx.db.index_attempt import expire_index_attempts
 from onyx.db.models import IndexModelStatus
 from onyx.db.models import User
@@ -22,6 +21,7 @@ from onyx.db.search_settings import get_embedding_provider_from_provider_type
 from onyx.db.search_settings import get_secondary_search_settings
 from onyx.db.search_settings import update_current_search_settings
 from onyx.db.search_settings import update_search_settings_status
+from onyx.db.session import get_session
 from onyx.document_index.document_index_utils import get_multipass_config
 from onyx.document_index.factory import get_default_document_index
 from onyx.file_processing.unstructured import delete_unstructured_api_key

backend/onyx/server/manage/slack_bot.py

@@ -8,10 +8,10 @@ from sqlalchemy.orm import Session
 from onyx.auth.users import current_admin_user
 from onyx.configs.constants import MilestoneRecordType
 from onyx.db.constants import SLACK_BOT_PERSONA_PREFIX
-from onyx.db.engine import get_session
 from onyx.db.models import ChannelConfig
 from onyx.db.models import User
 from onyx.db.persona import get_persona_by_id
+from onyx.db.session import get_session
 from onyx.db.slack_bot import fetch_slack_bot
 from onyx.db.slack_bot import fetch_slack_bot_tokens
 from onyx.db.slack_bot import fetch_slack_bots

backend/onyx/server/manage/users.py

@@ -42,9 +42,9 @@ from onyx.configs.constants import AuthType
 from onyx.configs.constants import FASTAPI_USERS_AUTH_COOKIE_NAME
 from onyx.db.api_key import is_api_key_email_address
 from onyx.db.auth import get_total_users_count
-from onyx.db.engine import get_session
 from onyx.db.models import AccessToken
 from onyx.db.models import User
+from onyx.db.session import get_session
 from onyx.db.users import delete_user_from_db
 from onyx.db.users import get_all_users
 from onyx.db.users import get_page_of_filtered_users

backend/onyx/server/onyx_api/ingestion.py

@@ -11,11 +11,11 @@ from onyx.connectors.models import IndexAttemptMetadata
 from onyx.db.connector_credential_pair import get_connector_credential_pair_from_id
 from onyx.db.document import get_documents_by_cc_pair
 from onyx.db.document import get_ingestion_documents
-from onyx.db.engine import get_session
 from onyx.db.models import User
 from onyx.db.search_settings import get_active_search_settings
 from onyx.db.search_settings import get_current_search_settings
 from onyx.db.search_settings import get_secondary_search_settings
+from onyx.db.session import get_session
 from onyx.document_index.factory import get_default_document_index
 from onyx.indexing.embedder import DefaultIndexingEmbedder
 from onyx.indexing.indexing_pipeline import build_indexing_pipeline

backend/onyx/server/openai_assistants_api/asssistants_api.py

@@ -11,7 +11,6 @@ from sqlalchemy.orm import Session
 
 from onyx.auth.users import current_user
 from onyx.context.search.enums import RecencyBiasSetting
-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
@@ -19,6 +18,7 @@ from onyx.db.persona import get_personas_for_user
 from onyx.db.persona import mark_persona_as_deleted
 from onyx.db.persona import upsert_persona
 from onyx.db.prompts import upsert_prompt
+from onyx.db.session import get_session
 from onyx.db.tools import get_tool_by_name
 from onyx.utils.logger import setup_logger
 

backend/onyx/server/openai_assistants_api/messages_api.py

@@ -18,8 +18,8 @@ from onyx.db.chat import get_chat_message
 from onyx.db.chat import get_chat_messages_by_session
 from onyx.db.chat import get_chat_session_by_id
 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.utils import check_number_of_tokens
 
 router = APIRouter(prefix="")

backend/onyx/server/openai_assistants_api/runs_api.py

@@ -18,9 +18,9 @@ from onyx.db.chat import get_chat_message
 from onyx.db.chat import get_chat_messages_by_session
 from onyx.db.chat import get_chat_session_by_id
 from onyx.db.chat import get_or_create_root_message
-from onyx.db.engine import get_session
 from onyx.db.models import ChatMessage
 from onyx.db.models import User
+from onyx.db.session import get_session
 from onyx.server.query_and_chat.models import ChatMessageDetail
 from onyx.server.query_and_chat.models import CreateChatMessageRequest
 from onyx.tools.tool_implementations.search.search_tool import SearchTool

backend/onyx/server/openai_assistants_api/threads_api.py

@@ -13,8 +13,8 @@ from onyx.db.chat import delete_chat_session
 from onyx.db.chat import get_chat_session_by_id
 from onyx.db.chat import get_chat_sessions_by_user
 from onyx.db.chat import update_chat_session
-from onyx.db.engine import get_session
 from onyx.db.models import User
+from onyx.db.session import get_session
 from onyx.server.query_and_chat.models import ChatSessionDetails
 from onyx.server.query_and_chat.models import ChatSessionsResponse
 

backend/onyx/server/query_and_chat/chat_backend.py

@@ -48,12 +48,12 @@ from onyx.db.chat import set_as_latest_chat_message
 from onyx.db.chat import translate_db_message_to_chat_message_detail
 from onyx.db.chat import update_chat_session
 from onyx.db.chat_search import search_chat_sessions
-from onyx.db.engine import get_session
-from onyx.db.engine import get_session_with_tenant
 from onyx.db.feedback import create_chat_message_feedback
 from onyx.db.feedback import create_doc_retrieval_feedback
 from onyx.db.models import User
 from onyx.db.persona import get_persona_by_id
+from onyx.db.session import get_session
+from onyx.db.session import get_session_with_tenant
 from onyx.file_processing.extract_file_text import docx_to_txt_filename
 from onyx.file_processing.extract_file_text import extract_file_text
 from onyx.file_store.file_store import get_default_file_store

backend/onyx/server/query_and_chat/query_backend.py

@@ -22,9 +22,9 @@ from onyx.db.chat import get_search_docs_for_chat_message
 from onyx.db.chat import get_valid_messages_from_query_sessions
 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
 from onyx.db.models import User
 from onyx.db.search_settings import get_current_search_settings
+from onyx.db.session import get_session
 from onyx.db.tag import find_tags
 from onyx.document_index.factory import get_default_document_index
 from onyx.document_index.vespa.index import VespaIndex