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

Update force delete script to handle user groups

Browse Source
This commit is contained in:
Weves 2024-07-21 22:21:40 -07:00 committed by Chris Weaver
parent 1b49d17239
commit 6222f533be
1 changed files with 18 additions and 4 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

22
backend/scripts/force_delete_connector_by_id.py
View File

@ -14,7 +14,10 @@ sys.path.append(parent_dir)
# flake8: noqa: E402 # flake8: noqa: E402
# Now import Danswer modules # Now import Danswer modules
from danswer.db.models import DocumentSet__ConnectorCredentialPair from danswer.db.models import (
DocumentSet__ConnectorCredentialPair,
UserGroup__ConnectorCredentialPair,
)
from danswer.db.connector import fetch_connector_by_id from danswer.db.connector import fetch_connector_by_id
from danswer.db.document import get_documents_for_connector_credential_pair from danswer.db.document import get_documents_for_connector_credential_pair
from danswer.db.index_attempt import ( from danswer.db.index_attempt import (
@ -44,7 +47,7 @@ logger = setup_logger()
_DELETION_BATCH_SIZE = 1000 _DELETION_BATCH_SIZE = 1000
def unsafe_deletion( def _unsafe_deletion(
db_session: Session, db_session: Session,
document_index: DocumentIndex, document_index: DocumentIndex,
cc_pair: ConnectorCredentialPair, cc_pair: ConnectorCredentialPair,
@ -82,11 +85,22 @@ def unsafe_deletion(
credential_id=credential_id, credential_id=credential_id,
) )
# Delete document sets + connector / credential Pairs # Delete document sets
stmt = delete(DocumentSet__ConnectorCredentialPair).where( stmt = delete(DocumentSet__ConnectorCredentialPair).where(
DocumentSet__ConnectorCredentialPair.connector_credential_pair_id == pair_id DocumentSet__ConnectorCredentialPair.connector_credential_pair_id == pair_id
) )
db_session.execute(stmt) db_session.execute(stmt)
# delete user group associations
stmt = delete(UserGroup__ConnectorCredentialPair).where(
UserGroup__ConnectorCredentialPair.cc_pair_id == pair_id
)
db_session.execute(stmt)
# need to flush to avoid foreign key violations
db_session.flush()
# delete the actual connector credential pair
stmt = delete(ConnectorCredentialPair).where( stmt = delete(ConnectorCredentialPair).where(
ConnectorCredentialPair.connector_id == connector_id, ConnectorCredentialPair.connector_id == connector_id,
ConnectorCredentialPair.credential_id == credential_id, ConnectorCredentialPair.credential_id == credential_id,
@ -168,7 +182,7 @@ def _delete_connector(cc_pair_id: int, db_session: Session) -> None:
primary_index_name=curr_ind_name, secondary_index_name=sec_ind_name primary_index_name=curr_ind_name, secondary_index_name=sec_ind_name
) )
files_deleted_count = unsafe_deletion( files_deleted_count = _unsafe_deletion(
db_session=db_session, db_session=db_session,
document_index=document_index, document_index=document_index,
cc_pair=cc_pair, cc_pair=cc_pair,