From 9bffeb65af7ae1ab1a09eb5aa7051cf51ae0def9 Mon Sep 17 00:00:00 2001 From: pablonyx Date: Sun, 29 Dec 2024 10:58:38 -0500 Subject: [PATCH] Eagerly load CCpair connectors (#3531) * remove left over vim command * eager loading * Revert "remove left over vim command" This reverts commit 184a134ae08e8e6daa48de9a4c19b3e4ac860aff. --- backend/onyx/db/connector_credential_pair.py | 10 +++++++++- backend/onyx/server/documents/connector.py | 2 +- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/backend/onyx/db/connector_credential_pair.py b/backend/onyx/db/connector_credential_pair.py index fd9b3bc6a07e..3c796492b11e 100644 --- a/backend/onyx/db/connector_credential_pair.py +++ b/backend/onyx/db/connector_credential_pair.py @@ -7,6 +7,7 @@ from sqlalchemy import exists from sqlalchemy import Select from sqlalchemy import select from sqlalchemy.orm import aliased +from sqlalchemy.orm import joinedload from sqlalchemy.orm import Session from onyx.configs.constants import DocumentSource @@ -90,15 +91,22 @@ def get_connector_credential_pairs( user: User | None = None, get_editable: bool = True, ids: list[int] | None = None, + eager_load_connector: bool = False, ) -> list[ConnectorCredentialPair]: stmt = select(ConnectorCredentialPair).distinct() + + if eager_load_connector: + stmt = stmt.options(joinedload(ConnectorCredentialPair.connector)) + stmt = _add_user_filters(stmt, user, get_editable) + if not include_disabled: stmt = stmt.where( ConnectorCredentialPair.status == ConnectorCredentialPairStatus.ACTIVE - ) # noqa + ) if ids: stmt = stmt.where(ConnectorCredentialPair.id.in_(ids)) + return list(db_session.scalars(stmt).all()) diff --git a/backend/onyx/server/documents/connector.py b/backend/onyx/server/documents/connector.py index 065576f73d7e..f7f4633e50e7 100644 --- a/backend/onyx/server/documents/connector.py +++ b/backend/onyx/server/documents/connector.py @@ -1058,7 +1058,7 @@ def get_basic_connector_indexing_status( _: User = Depends(current_user), db_session: Session = Depends(get_session), ) -> list[BasicCCPairInfo]: - cc_pairs = get_connector_credential_pairs(db_session) + cc_pairs = get_connector_credential_pairs(db_session, eager_load_connector=True) return [ BasicCCPairInfo( has_successful_run=cc_pair.last_successful_index_time is not None,