Added doc_set__user_group cleanup for user_group deletion (#2551)

This commit is contained in:
hagen-danswer 2024-09-24 09:09:52 -07:00 committed by GitHub
parent 487250320b
commit 8cfe80c53a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194

View File

@ -16,6 +16,7 @@ from danswer.db.models import ConnectorCredentialPair
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()