highperfocused/danswer
1
0
Fork 0
mirror of https://github.com/danswer-ai/danswer.git synced 2025-06-29 17:20:44 +02:00
Code Issues Packages Projects Releases Wiki Activity

Automatically delete document set relationships when deleting a ConnectorCredentialPair

Browse Source
This commit is contained in:
Weves
2023-10-10 15:09:59 -07:00
committed by Chris Weaver
parent 3142e2eed2
commit 091c2c8a80
3 changed files with 35 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
backend/danswer/background/connector_deletion.py
View File

@ -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 DocumentIndex
from danswer.datastores.interfaces import UpdateRequest from danswer.datastores.interfaces import UpdateRequest
from danswer.db.connector import fetch_connector_by_id 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.connector_credential_pair import get_connector_credential_pair
from danswer.db.deletion_attempt import check_deletion_attempt_is_allowed from danswer.db.deletion_attempt import check_deletion_attempt_is_allowed
from danswer.db.document import delete_document_by_connector_credential_pair 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.document import prepare_to_modify_documents
from danswer.db.engine import get_sqlalchemy_engine from danswer.db.engine import get_sqlalchemy_engine
from danswer.db.index_attempt import delete_index_attempts from danswer.db.index_attempt import delete_index_attempts
from danswer.db.models import ConnectorCredentialPair
from danswer.server.models import ConnectorCredentialPairIdentifier from danswer.server.models import ConnectorCredentialPairIdentifier
from danswer.utils.logger import setup_logger from danswer.utils.logger import setup_logger
@ -96,9 +102,11 @@ def _delete_connector_credential_pair_batch(
def _delete_connector_credential_pair( def _delete_connector_credential_pair(
db_session: Session, db_session: Session,
document_index: DocumentIndex, document_index: DocumentIndex,
connector_id: int, cc_pair: ConnectorCredentialPair,
credential_id: int,
) -> int: ) -> int:
connector_id = cc_pair.connector_id
credential_id = cc_pair.credential_id
num_docs_deleted = 0 num_docs_deleted = 0
while True: while True:
documents = get_documents_for_connector_credential_pair( documents = get_documents_for_connector_credential_pair(
@ -124,7 +132,11 @@ def _delete_connector_credential_pair(
connector_id=connector_id, connector_id=connector_id,
credential_id=credential_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, db_session=db_session,
connector_id=connector_id, connector_id=connector_id,
credential_id=credential_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( return _delete_connector_credential_pair(
db_session=db_session, db_session=db_session,
document_index=get_default_document_index(), document_index=get_default_document_index(),
connector_id=connector_id, cc_pair=cc_pair,
credential_id=credential_id,
) )
except Exception as e: except Exception as e:
logger.exception(f"Failed to run connector_deletion due to {e}") logger.exception(f"Failed to run connector_deletion due to {e}")

13
backend/danswer/db/connector_credential_pair.py
View File

@ -9,6 +9,7 @@ from sqlalchemy.orm import Session
from danswer.db.connector import fetch_connector_by_id from danswer.db.connector import fetch_connector_by_id
from danswer.db.credentials import fetch_credential_by_id from danswer.db.credentials import fetch_credential_by_id
from danswer.db.models import ConnectorCredentialPair from danswer.db.models import ConnectorCredentialPair
from danswer.db.models import DocumentSet__ConnectorCredentialPair
from danswer.db.models import IndexingStatus from danswer.db.models import IndexingStatus
from danswer.db.models import User from danswer.db.models import User
from danswer.server.models import StatusResponse from danswer.server.models import StatusResponse
@ -81,7 +82,17 @@ def update_connector_credential_pair(
db_session.commit() 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, db_session: Session,
connector_id: int, connector_id: int,
credential_id: int, credential_id: int,

12
backend/danswer/db/document_set.py
View File

@ -18,7 +18,7 @@ from danswer.server.models import DocumentSetCreationRequest
from danswer.server.models import DocumentSetUpdateRequest 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 db_session: Session, document_set_id: int, is_current: bool | None = None
) -> None: ) -> None:
"""NOTE: does not commit transaction, this must be done by the caller""" """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) 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 db_session: Session, document_set_id: int
) -> None: ) -> None:
"""NOTE: does not commit transaction, this must be done by the caller""" """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 # update the attached CC pairs
# first, mark all existing CC pairs as not current # 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 db_session=db_session, document_set_id=document_set_row.id
) )
# add in rows for the new CC pairs # 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 # mark as up to date
document_set.is_up_to_date = True document_set.is_up_to_date = True
# delete outdated relationship table rows # 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=db_session, document_set_id=document_set_id, is_current=False
) )
db_session.commit() db_session.commit()
@ -155,7 +155,7 @@ def delete_document_set(
document_set_row: DocumentSetDBModel, db_session: Session document_set_row: DocumentSetDBModel, db_session: Session
) -> None: ) -> None:
# delete all relationships to CC pairs # 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=db_session, document_set_id=document_set_row.id
) )
db_session.delete(document_set_row) 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 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 db_session=db_session, document_set_id=document_set_id
) )
# mark the row as needing a sync, it will be deleted there since there # mark the row as needing a sync, it will be deleted there since there