feat: refactor document index reconciliation and deletion logic

- Introduced a new IndexAction class to standardize index operation types (CREATE, UPDATE, DELETE).
- Updated DocumentIndexReconciler to utilize the new IndexAction constants for improved clarity and maintainability.
- Refactored document deletion logic in DocumentService to streamline the deletion process and ensure proper index reconciliation.
- Enhanced error handling and logging during document deletion and index management operations.
- Updated documentation to reflect changes in the indexing architecture and operational flow.

Author: iziang

aperag/index/reconciler.py

@@ -29,6 +29,7 @@
 )
 from aperag.tasks.scheduler import TaskScheduler, create_task_scheduler
 from aperag.utils.utils import utc_now
+from aperag.utils.constant import IndexAction
 
 logger = logging.getLogger(__name__)
 
@@ -71,30 +72,30 @@ def _get_indexes_needing_reconciliation(self, session: Session) -> List[Document
         """
         from collections import defaultdict
 
-        operations = defaultdict(lambda: {"create": [], "update": [], "delete": []})
+        operations = defaultdict(lambda: {IndexAction.CREATE: [], IndexAction.UPDATE: [], IndexAction.DELETE: []})
 
         conditions = {
-            "create": and_(
+            IndexAction.CREATE: and_(
                 DocumentIndex.status == DocumentIndexStatus.PENDING,
                 DocumentIndex.observed_version < DocumentIndex.version,
                 DocumentIndex.version == 1,
             ),
-            "update": and_(
+            IndexAction.UPDATE: and_(
                 DocumentIndex.status == DocumentIndexStatus.PENDING,
                 DocumentIndex.observed_version < DocumentIndex.version,
                 DocumentIndex.version > 1,
             ),
-            "delete": and_(
+            IndexAction.DELETE: and_(
                 DocumentIndex.status == DocumentIndexStatus.DELETING,
             ),
         }
 
-        for operation_type, condition in conditions.items():
+        for action, condition in conditions.items():
             stmt = select(DocumentIndex).where(condition)
             result = session.execute(stmt)
             indexes = result.scalars().all()
             for index in indexes:
-                operations[index.document_id][operation_type].append(index)
+                operations[index.document_id][action].append(index)
 
         return operations
 
@@ -106,9 +107,9 @@ def _reconcile_single_document(self, document_id: str, operations: dict):
             # Collect indexes for this document that need claiming
             indexes_to_claim = []
 
-            for operation_type, doc_indexes in operations.items():
+            for action, doc_indexes in operations.items():
                 for doc_index in doc_indexes:
-                    indexes_to_claim.append((doc_index.id, doc_index.index_type, operation_type))
+                    indexes_to_claim.append((doc_index.id, doc_index.index_type, action))
 
             # Atomically claim the indexes for this document
             claimed_indexes = self._claim_document_indexes(session, document_id, indexes_to_claim)
@@ -129,10 +130,10 @@ def _claim_document_indexes(self, session: Session, document_id: str, indexes_to
         claimed_indexes = []
 
         try:
-            for index_id, index_type, operation_type in indexes_to_claim:
-                if operation_type in ["create", "update"]:
+            for index_id, index_type, action in indexes_to_claim:
+                if action in [IndexAction.CREATE, IndexAction.UPDATE]:
                     target_state = DocumentIndexStatus.CREATING
-                elif operation_type == "delete":
+                elif action == IndexAction.DELETE:
                     target_state = DocumentIndexStatus.DELETION_IN_PROGRESS
                 else:
                     continue
@@ -146,21 +147,21 @@ def _claim_document_indexes(self, session: Session, document_id: str, indexes_to
                     continue
 
                 # Build appropriate claiming conditions based on operation type
-                if operation_type == "create":
+                if action == IndexAction.CREATE:
                     claiming_conditions = [
                         DocumentIndex.id == index_id,
                         DocumentIndex.status == DocumentIndexStatus.PENDING,
                         DocumentIndex.observed_version < DocumentIndex.version,
                         DocumentIndex.version == 1,
                     ]
-                elif operation_type == "update":
+                elif action == IndexAction.UPDATE:
                     claiming_conditions = [
                         DocumentIndex.id == index_id,
                         DocumentIndex.status == DocumentIndexStatus.PENDING,
                         DocumentIndex.observed_version < DocumentIndex.version,
                         DocumentIndex.version > 1,
                     ]
-                elif operation_type == "delete":
+                elif action == IndexAction.DELETE:
                     claiming_conditions = [
                         DocumentIndex.id == index_id,
                         DocumentIndex.status == DocumentIndexStatus.DELETING,
@@ -180,10 +181,10 @@ def _claim_document_indexes(self, session: Session, document_id: str, indexes_to
                         'index_id': index_id,
                         'document_id': document_id,
                         'index_type': index_type,
-                        'operation_type': operation_type,
-                        'target_version': current_index.version if operation_type in ["create", "update"] else None,
+                        'action': action,
+                        'target_version': current_index.version if action in [IndexAction.CREATE, IndexAction.UPDATE] else None,
                     })
-                    logger.debug(f"Claimed index {index_id} for document {document_id} ({operation_type})")
+                    logger.debug(f"Claimed index {index_id} for document {document_id} ({action})")
                 else:
                     logger.debug(f"Could not claim index {index_id} for document {document_id}")
 
@@ -202,12 +203,12 @@ def _reconcile_document_operations(self, document_id: str, claimed_indexes: List
         # Group by operation type to batch operations
         operations_by_type = defaultdict(list)
         for claimed_index in claimed_indexes:
-            operation_type = claimed_index['operation_type']
-            operations_by_type[operation_type].append(claimed_index)
+            action = claimed_index['action']
+            operations_by_type[action].append(claimed_index)
 
         # Process create operations as a batch
-        if "create" in operations_by_type:
-            create_indexes = operations_by_type["create"]
+        if IndexAction.CREATE in operations_by_type:
+            create_indexes = operations_by_type[IndexAction.CREATE]
             create_types = [claimed_index['index_type'] for claimed_index in create_indexes]
             context = {}
 
@@ -227,8 +228,8 @@ def _reconcile_document_operations(self, document_id: str, claimed_indexes: List
             logger.info(f"Scheduled create task for document {document_id}, types: {create_types}")
 
         # Process update operations as a batch
-        if "update" in operations_by_type:
-            update_indexes = operations_by_type["update"]
+        if IndexAction.UPDATE in operations_by_type:
+            update_indexes = operations_by_type[IndexAction.UPDATE]
             update_types = [claimed_index['index_type'] for claimed_index in update_indexes]
             context = {}
 
@@ -248,8 +249,8 @@ def _reconcile_document_operations(self, document_id: str, claimed_indexes: List
             logger.info(f"Scheduled update task for document {document_id}, types: {update_types}")
 
         # Process delete operations as a batch
-        if "delete" in operations_by_type:
-            delete_indexes = operations_by_type["delete"]
+        if IndexAction.DELETE in operations_by_type:
+            delete_indexes = operations_by_type[IndexAction.DELETE]
             delete_types = [claimed_index['index_type'] for claimed_index in delete_indexes]
 
             task_id = self.task_scheduler.schedule_delete_index(

aperag/service/document_service.py

@@ -250,111 +250,67 @@ async def get_document(self, user: str, collection_id: str, document_id: str) ->
         async for session in get_async_session():
             return await self.build_document_response(document, session)
 
-    async def delete_document(self, user: str, collection_id: str, document_id: str) -> Optional[view_models.Document]:
-        """Delete document by ID (idempotent operation)
-
-        Returns the deleted document or None if already deleted/not found
+    async def _delete_document(self, session: AsyncSession, user: str, collection_id: str, document_id: str):
+        """
+        Core logic to delete a single document and its associated resources.
+        This method is designed to be called within a transaction.
         """
+        # Validate document existence and ownership
         document = await self.db_ops.query_document(user, collection_id, document_id)
         if document is None:
-            # Document already deleted or never existed - idempotent operation
-            return None
-
-        # Delete document and indexes atomically in a single transaction
-        async def _delete_document_atomically(session):
-            from sqlalchemy import select
-
-            from aperag.db.models import Document, DocumentStatus, utc_now
-
-            # Get and delete document
-            stmt = select(Document).where(
-                Document.id == document_id,
-                Document.collection_id == collection_id,
-                Document.user == user,
-                Document.status != DocumentStatus.DELETED,
-            )
-            result = await session.execute(stmt)
-            doc_to_delete = result.scalars().first()
+            # Silently ignore if document not found, as it might have been deleted by another process
+            logger.warning(f"Document {document_id} not found for deletion, skipping.")
+            return
 
-            if not doc_to_delete:
-                return None
-
-            # Soft delete document
-            doc_to_delete.status = DocumentStatus.DELETED
-            doc_to_delete.gmt_deleted = utc_now()
-            session.add(doc_to_delete)
-            await session.flush()
-            await session.refresh(doc_to_delete)
-
-            # Mark index specs for deletion
-            await document_index_manager.delete_document_indexes(session, document_id)
-
-            # Build response object
-            return await self.build_document_response(doc_to_delete, session)
-
-        result = await self.db_ops.execute_with_transaction(_delete_document_atomically)
+        # Use index manager to mark all related indexes for deletion
+        await document_index_manager.delete_document_indexes(
+            document_id=document.id, index_types=None, session=session
+        )
 
-        if result:
-            # Delete object storage files after successful database transaction
-            obj_store = get_object_store()
+        # Delete from object store
+        obj_store = get_object_store()
+        metadata = json.loads(document.doc_metadata) if document.doc_metadata else {}
+        if object_path := metadata.get("object_path"):
             try:
-                await sync_to_async(obj_store.delete_objects_by_prefix)(f"{document.object_store_base_path()}/")
+                # Use delete_objects_by_prefix to remove all related files (original, chunks, etc.)
+                await sync_to_async(obj_store.delete_objects_by_prefix)(document.object_store_base_path())
+                logger.info(f"Deleted objects from object store with prefix: {document.object_store_base_path()}")
             except Exception as e:
-                logger.warning(f"Failed to delete object storage files for document {document_id}: {e}")
-
-            # Trigger index reconciliation after successful document deletion
-            _trigger_index_reconciliation()
-
-            return result
+                logger.warning(f"Failed to delete objects for document {document.id} from object store: {e}")
 
-        return None
-
-    async def delete_documents(self, user: str, collection_id: str, document_ids: List[str]) -> dict:
-        # Delete documents and indexes atomically in a single transaction
-        async def _delete_documents_atomically(session):
-            from sqlalchemy import select
-
-            from aperag.db.models import Document, DocumentStatus, utc_now
-
-            # Get documents to delete
-            stmt = select(Document).where(
-                Document.id.in_(document_ids),
-                Document.collection_id == collection_id,
-                Document.user == user,
-                Document.status != DocumentStatus.DELETED,
-            )
-            result = await session.execute(stmt)
-            documents_to_delete = result.scalars().all()
+        # Delete the document record from the database
+        await session.delete(document)
+        await session.flush()
+        logger.info(f"Successfully marked document {document.id} and its indexes for deletion.")
 
-            if not documents_to_delete:
-                return [], list(document_ids)
+        return document
 
-            # Soft delete documents
-            success_ids = []
-            for doc in documents_to_delete:
-                doc.status = DocumentStatus.DELETED
-                doc.gmt_deleted = utc_now()
-                session.add(doc)
-                success_ids.append(doc.id)
+    async def delete_document(self, user: str, collection_id: str, document_id: str) -> dict:
+        """Delete a single document and trigger index reconciliation."""
 
-            await session.flush()
+        async def _delete_document_atomically(session: AsyncSession):
+            return await self._delete_document(session, user, collection_id, document_id)
 
-            # Delete indexes for all successful deletions
-            for doc_id in success_ids:
-                await document_index_manager.delete_document_indexes(session, doc_id)
+        result = await self.db_ops.execute_with_transaction(_delete_document_atomically)
 
-            # Calculate failed IDs
-            failed_ids = list(set(document_ids) - set(success_ids))
-            return success_ids, failed_ids
+        # Trigger reconciliation to process the deletion
+        _trigger_index_reconciliation()
+        return result
 
-        success_ids, failed_ids = await self.db_ops.execute_with_transaction(_delete_documents_atomically)
+    async def delete_documents(self, user: str, collection_id: str, document_ids: List[str]) -> dict:
+        """Delete multiple documents and trigger index reconciliation."""
 
-        result = {"success": success_ids, "failed": failed_ids}
+        async def _delete_documents_atomically(session: AsyncSession):
+            deleted_ids = []
+            for doc_id in document_ids:
+                await self._delete_document(session, user, collection_id, doc_id)
+                deleted_ids.append(doc_id)
+            return {"deleted_ids": deleted_ids, "status": "success"}
 
-        # Trigger index reconciliation after successful batch document deletion
-        if result.get("success"):  # Only trigger if at least one document was deleted successfully
-            _trigger_index_reconciliation()
+        result = await self.db_ops.execute_with_transaction(_delete_documents_atomically)
 
+        # Trigger reconciliation to process deletions
+        _trigger_index_reconciliation()
         return result
 
     async def rebuild_document_indexes(

aperag/utils/constant.py

@@ -39,3 +39,9 @@ class QuotaType:
     MAX_COLLECTION_COUNT = "max_collection_count"
     MAX_DOCUMENT_COUNT = "max_document_count"
     MAX_CONVERSATION_COUNT = "max_conversation_count"
+
+
+class IndexAction:
+    CREATE = "create"
+    UPDATE = "update"
+    DELETE = "delete"