diff --git a/backend/danswer/background/connector_deletion.py b/backend/danswer/background/connector_deletion.py index 09312fc701e..dcb4c5aaabc 100644 --- a/backend/danswer/background/connector_deletion.py +++ b/backend/danswer/background/connector_deletion.py @@ -17,7 +17,12 @@ from danswer.datastores.document_index import get_default_document_index from danswer.datastores.interfaces import DocumentIndex from danswer.datastores.interfaces import UpdateRequest from danswer.db.connector import fetch_connector_by_id -from danswer.db.connector_credential_pair import delete_connector_credential_pair +from danswer.db.connector_credential_pair import ( + delete_connector_credential_pair__no_commit, +) +from danswer.db.connector_credential_pair import ( + delete_document_set_relationships_for_cc_pair__no_commit, +) from danswer.db.connector_credential_pair import get_connector_credential_pair from danswer.db.deletion_attempt import check_deletion_attempt_is_allowed from danswer.db.document import delete_document_by_connector_credential_pair @@ -27,6 +32,7 @@ from danswer.db.document import get_documents_for_connector_credential_pair from danswer.db.document import prepare_to_modify_documents from danswer.db.engine import get_sqlalchemy_engine from danswer.db.index_attempt import delete_index_attempts +from danswer.db.models import ConnectorCredentialPair from danswer.server.models import ConnectorCredentialPairIdentifier from danswer.utils.logger import setup_logger @@ -96,9 +102,11 @@ def _delete_connector_credential_pair_batch( def _delete_connector_credential_pair( db_session: Session, document_index: DocumentIndex, - connector_id: int, - credential_id: int, + cc_pair: ConnectorCredentialPair, ) -> int: + connector_id = cc_pair.connector_id + credential_id = cc_pair.credential_id + num_docs_deleted = 0 while True: documents = get_documents_for_connector_credential_pair( @@ -124,7 +132,11 @@ def _delete_connector_credential_pair( connector_id=connector_id, credential_id=credential_id, ) - delete_connector_credential_pair( + delete_document_set_relationships_for_cc_pair__no_commit( + cc_pair_id=cc_pair.id, + db_session=db_session, + ) + delete_connector_credential_pair__no_commit( db_session=db_session, connector_id=connector_id, credential_id=credential_id, @@ -168,8 +180,7 @@ def cleanup_connector_credential_pair(connector_id: int, credential_id: int) -> return _delete_connector_credential_pair( db_session=db_session, document_index=get_default_document_index(), - connector_id=connector_id, - credential_id=credential_id, + cc_pair=cc_pair, ) except Exception as e: logger.exception(f"Failed to run connector_deletion due to {e}") diff --git a/backend/danswer/db/connector_credential_pair.py b/backend/danswer/db/connector_credential_pair.py index 2b1767f1e23..35a9c70289e 100644 --- a/backend/danswer/db/connector_credential_pair.py +++ b/backend/danswer/db/connector_credential_pair.py @@ -9,6 +9,7 @@ from sqlalchemy.orm import Session from danswer.db.connector import fetch_connector_by_id from danswer.db.credentials import fetch_credential_by_id from danswer.db.models import ConnectorCredentialPair +from danswer.db.models import DocumentSet__ConnectorCredentialPair from danswer.db.models import IndexingStatus from danswer.db.models import User from danswer.server.models import StatusResponse @@ -81,7 +82,17 @@ def update_connector_credential_pair( db_session.commit() -def delete_connector_credential_pair( +def delete_document_set_relationships_for_cc_pair__no_commit( + cc_pair_id: int, db_session: Session +) -> None: + """NOTE: does not commit transaction, this must be done by the caller""" + stmt = delete(DocumentSet__ConnectorCredentialPair).where( + DocumentSet__ConnectorCredentialPair.connector_credential_pair_id == cc_pair_id + ) + db_session.execute(stmt) + + +def delete_connector_credential_pair__no_commit( db_session: Session, connector_id: int, credential_id: int, diff --git a/backend/danswer/db/document_set.py b/backend/danswer/db/document_set.py index 1e2f581b056..10ee5aa8845 100644 --- a/backend/danswer/db/document_set.py +++ b/backend/danswer/db/document_set.py @@ -18,7 +18,7 @@ from danswer.server.models import DocumentSetCreationRequest from danswer.server.models import DocumentSetUpdateRequest -def _delete_document_set_cc_pairs( +def _delete_document_set_cc_pairs__no_commit( db_session: Session, document_set_id: int, is_current: bool | None = None ) -> None: """NOTE: does not commit transaction, this must be done by the caller""" @@ -30,7 +30,7 @@ def _delete_document_set_cc_pairs( db_session.execute(stmt) -def _mark_document_set_cc_pairs_as_outdated( +def _mark_document_set_cc_pairs_as_outdated__no_commit( db_session: Session, document_set_id: int ) -> None: """NOTE: does not commit transaction, this must be done by the caller""" @@ -115,7 +115,7 @@ def update_document_set( # update the attached CC pairs # first, mark all existing CC pairs as not current - _mark_document_set_cc_pairs_as_outdated( + _mark_document_set_cc_pairs_as_outdated__no_commit( db_session=db_session, document_set_id=document_set_row.id ) # add in rows for the new CC pairs @@ -145,7 +145,7 @@ def mark_document_set_as_synced(document_set_id: int, db_session: Session) -> No # mark as up to date document_set.is_up_to_date = True # delete outdated relationship table rows - _delete_document_set_cc_pairs( + _delete_document_set_cc_pairs__no_commit( db_session=db_session, document_set_id=document_set_id, is_current=False ) db_session.commit() @@ -155,7 +155,7 @@ def delete_document_set( document_set_row: DocumentSetDBModel, db_session: Session ) -> None: # delete all relationships to CC pairs - _delete_document_set_cc_pairs( + _delete_document_set_cc_pairs__no_commit( db_session=db_session, document_set_id=document_set_row.id ) db_session.delete(document_set_row) @@ -184,7 +184,7 @@ def mark_document_set_as_to_be_deleted( ) # delete all relationships to CC pairs - _delete_document_set_cc_pairs( + _delete_document_set_cc_pairs__no_commit( db_session=db_session, document_set_id=document_set_id ) # mark the row as needing a sync, it will be deleted there since there