From 8cfe80c53aed24352824239fe3eb40fbd5566c65 Mon Sep 17 00:00:00 2001 From: hagen-danswer Date: Tue, 24 Sep 2024 09:09:52 -0700 Subject: [PATCH] Added doc_set__user_group cleanup for user_group deletion (#2551) --- backend/ee/danswer/db/user_group.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/backend/ee/danswer/db/user_group.py b/backend/ee/danswer/db/user_group.py index f62fe19a7d7..529112004b0 100644 --- a/backend/ee/danswer/db/user_group.py +++ b/backend/ee/danswer/db/user_group.py @@ -16,6 +16,7 @@ from danswer.db.models import Credential__UserGroup from danswer.db.models import Document from danswer.db.models import DocumentByConnectorCredentialPair +from danswer.db.models import DocumentSet__UserGroup from danswer.db.models import LLMProvider__UserGroup from danswer.db.models import TokenRateLimit__UserGroup from danswer.db.models import User @@ -532,6 +533,17 @@ def _cleanup_user_group__cc_pair_relationships__no_commit( db_session.delete(user_group__cc_pair_relationship) +def _cleanup_document_set__user_group_relationships__no_commit( + db_session: Session, user_group_id: int +) -> None: + """NOTE: does not commit the transaction.""" + db_session.execute( + delete(DocumentSet__UserGroup).where( + DocumentSet__UserGroup.user_group_id == user_group_id + ) + ) + + def mark_user_group_as_synced(db_session: Session, user_group: UserGroup) -> None: # cleanup outdated relationships _cleanup_user_group__cc_pair_relationships__no_commit( @@ -553,6 +565,9 @@ def delete_user_group(db_session: Session, user_group: UserGroup) -> None: user_group_id=user_group.id, outdated_only=False, ) + _cleanup_document_set__user_group_relationships__no_commit( + db_session=db_session, user_group_id=user_group.id + ) # need to flush so that we don't get a foreign key error when deleting the user group row db_session.flush()