From f7cc7190fe81f415a652fb3be84ba3e571f2aa73 Mon Sep 17 00:00:00 2001 From: Weves Date: Mon, 4 Sep 2023 15:44:31 -0700 Subject: [PATCH] Allow connectors that have documents with feedback to be deleted --- backend/danswer/db/document.py | 4 ++++ backend/danswer/db/feedback.py | 12 ++++++++++++ 2 files changed, 16 insertions(+) diff --git a/backend/danswer/db/document.py b/backend/danswer/db/document.py index 951931b9c..f27d5cb59 100644 --- a/backend/danswer/db/document.py +++ b/backend/danswer/db/document.py @@ -9,6 +9,7 @@ from sqlalchemy.orm import Session from danswer.configs.constants import DEFAULT_BOOST from danswer.datastores.interfaces import DocumentMetadata +from danswer.db.feedback import delete_document_feedback_for_documents from danswer.db.models import Document as DbDocument from danswer.db.models import DocumentByConnectorCredentialPair from danswer.db.utils import model_to_dict @@ -168,5 +169,8 @@ def delete_documents(db_session: Session, document_ids: list[str]) -> None: def delete_documents_complete(db_session: Session, document_ids: list[str]) -> None: logger.info(f"Deleting {len(document_ids)} documents from the DB") delete_document_by_connector_credential_pair(db_session, document_ids) + delete_document_feedback_for_documents( + document_ids=document_ids, db_session=db_session + ) delete_documents(db_session, document_ids) db_session.commit() diff --git a/backend/danswer/db/feedback.py b/backend/danswer/db/feedback.py index 3cfcf3a22..0adef3ade 100644 --- a/backend/danswer/db/feedback.py +++ b/backend/danswer/db/feedback.py @@ -1,6 +1,7 @@ from uuid import UUID from sqlalchemy import asc +from sqlalchemy import delete from sqlalchemy import desc from sqlalchemy import select from sqlalchemy.orm import Session @@ -151,3 +152,14 @@ def create_doc_retrieval_feedback( db_session.add(retrieval_feedback) db_session.commit() + + +def delete_document_feedback_for_documents( + document_ids: list[str], db_session: Session +) -> None: + """NOTE: does not commit transaction so that this can be used as part of a + larger transaction block.""" + stmt = delete(DocumentRetrievalFeedback).where( + DocumentRetrievalFeedback.document_id.in_(document_ids) + ) + db_session.execute(stmt)