diff --git a/api/db/services/document_service.py b/api/db/services/document_service.py index aa532af62..a2fee62d0 100644 --- a/api/db/services/document_service.py +++ b/api/db/services/document_service.py @@ -361,9 +361,10 @@ class DocumentService(CommonService): @DB.connection_context() def remove_document(cls, doc, tenant_id): from api.db.services.task_service import TaskService, cancel_all_task_of - cls.clear_chunk_num(doc.id) + if not cls.delete_document_and_update_kb_counts(doc.id): + return True - # Cancel all running tasks first Using preset function in task_service.py --- set cancel flag in Redis + # Cancel all running tasks first Using preset function in task_service.py --- set cancel flag in Redis try: cancel_all_task_of(doc.id) logging.info(f"Cancelled all tasks for document {doc.id}") @@ -419,7 +420,7 @@ class DocumentService(CommonService): except Exception as e: logging.warning(f"Failed to cleanup knowledge graph for document {doc.id}: {e}") - return cls.delete_by_id(doc.id) + return True @classmethod @DB.connection_context() @@ -521,9 +522,32 @@ class DocumentService(CommonService): Knowledgebase.id == kb_id).execute() return num + @classmethod + @DB.connection_context() + def delete_document_and_update_kb_counts(cls, doc_id) -> bool: + """Atomically delete the document row and update KB counters. + + Returns True if the document was deleted by this call, False if it was + already deleted by a concurrent request (idempotent). + """ + with DB.atomic(): + doc = cls.model.get_or_none(cls.model.id == doc_id) + if doc is None: + return False + deleted = cls.model.delete().where(cls.model.id == doc_id).execute() + if not deleted: + return False + Knowledgebase.update( + token_num=Knowledgebase.token_num - doc.token_num, + chunk_num=Knowledgebase.chunk_num - doc.chunk_num, + doc_num=Knowledgebase.doc_num - 1, + ).where(Knowledgebase.id == doc.kb_id).execute() + return True + @classmethod @DB.connection_context() def clear_chunk_num(cls, doc_id): + """Deprecated: use delete_document_and_update_kb_counts instead.""" doc = cls.model.get_by_id(doc_id) assert doc, "Can't fine document in database."