fix: handle missing documents in retry indexing task and streamline e…

…rror handling Signed-off-by: -LAN- <[email protected]>
langgenius · Dec 24, 2024 · af2a26e · af2a26e
1 parent 98d8c44
commit af2a26e
Showing 1 changed file with 25 additions and 24 deletions.
diff --git a/api/tasks/retry_document_indexing_task.py b/api/tasks/retry_document_indexing_task.py
@@ -58,36 +58,37 @@ def retry_document_indexing_task(dataset_id: str, document_ids: list[str]):
         document = (
             db.session.query(Document).filter(Document.id == document_id, Document.dataset_id == dataset_id).first()
         )
+        if not document:
+            logging.info(click.style("Document not found: {}".format(document_id), fg="yellow"))
+            return
         try:
-            if document:
-                # clean old data
-                index_processor = IndexProcessorFactory(document.doc_form).init_index_processor()
-
-                segments = db.session.query(DocumentSegment).filter(DocumentSegment.document_id == document_id).all()
-                if segments:
-                    index_node_ids = [segment.index_node_id for segment in segments]
-                    # delete from vector index
-                    index_processor.clean(dataset, index_node_ids)
+            # clean old data
+            index_processor = IndexProcessorFactory(document.doc_form).init_index_processor()
 
-                    for segment in segments:
-                        db.session.delete(segment)
-                    db.session.commit()
+            segments = db.session.query(DocumentSegment).filter(DocumentSegment.document_id == document_id).all()
+            if segments:
+                index_node_ids = [segment.index_node_id for segment in segments]
+                # delete from vector index
+                index_processor.clean(dataset, index_node_ids)
 
-                document.indexing_status = "parsing"
-                document.processing_started_at = datetime.datetime.utcnow()
-                db.session.add(document)
+                for segment in segments:
+                    db.session.delete(segment)
                 db.session.commit()
 
-                indexing_runner = IndexingRunner()
-                indexing_runner.run([document])
-                redis_client.delete(retry_indexing_cache_key)
+            document.indexing_status = "parsing"
+            document.processing_started_at = datetime.datetime.utcnow()
+            db.session.add(document)
+            db.session.commit()
+
+            indexing_runner = IndexingRunner()
+            indexing_runner.run([document])
+            redis_client.delete(retry_indexing_cache_key)
         except Exception as ex:
-            if document:
-                document.indexing_status = "error"
-                document.error = str(ex)
-                document.stopped_at = datetime.datetime.utcnow()
-                db.session.add(document)
-                db.session.commit()
+            document.indexing_status = "error"
+            document.error = str(ex)
+            document.stopped_at = datetime.datetime.utcnow()
+            db.session.add(document)
+            db.session.commit()
             logging.info(click.style(str(ex), fg="yellow"))
             redis_client.delete(retry_indexing_cache_key)
             pass