Fix cancel ingest task, it will be stilling running and show internal server error (#836) (#16945)

api/db/services/document_service.py

@@ -917,7 +917,7 @@ class DocumentService(CommonService):
             info["run"] = TaskStatus.RUNNING.value
             # keep the doc in DONE state when keep_progress=True for GraphRAG, RAPTOR and Mindmap tasks
 
-        cls.update_by_id(doc_id, info)
+        cls.model.update(info).where((cls.model.id == doc_id) & ((cls.model.run.is_null(True)) | (cls.model.run != TaskStatus.CANCEL.value))).execute()
 
     @classmethod
     @DB.connection_context()

api/db/services/task_service.py

@@ -421,7 +421,9 @@ class TaskService(CommonService):
             doc_info = {"progress": -1, "run": TaskStatus.FAIL.value, "update_time": current_timestamp(), "update_date": get_format_time()}
             if info.get("progress_msg"):
                 doc_info["progress_msg"] = trim_header_by_lines((task.progress_msg or "") + "\n" + info["progress_msg"], TASK_MAX_LOG_LENGTH)
-            DocumentService.update_by_id(task.doc_id, doc_info)
+            DocumentService.model.update(doc_info).where(
+                (DocumentService.model.id == task.doc_id) & ((DocumentService.model.run.is_null(True)) | (DocumentService.model.run != TaskStatus.CANCEL.value))
+            ).execute()
 
     @classmethod
     @DB.connection_context()

rag/svr/task_executor_refactor/chunk_builder.py

@@ -27,6 +27,7 @@ from timeit import default_timer as timer
 from typing import Dict, List
 
 from common.constants import ParserType
+from common.exceptions import TaskCanceledException
 from common.misc_utils import thread_pool_exec
 from rag.svr.task_executor_refactor.task_context import TaskContext
 
@@ -103,6 +104,8 @@ async def run_chunking(
         logging.info("Chunking({}) {}/{} done".format(timer() - st, ctx.location, ctx.name))
         ctx.recording_context.record("parser_config_after_merge", parser_config)
         return cks
+    except TaskCanceledException:
+        raise
     except Exception as e:
         ctx.progress_cb(-1, msg="Internal server error while chunking: %s" % str(e).replace("'", ""))
         logging.exception("Chunking {}/{} got exception".format(ctx.location, ctx.name))