From e0c8189f1ab6985598c93abc5374f21c5c89236d Mon Sep 17 00:00:00 2001
From: crazywoola <100913391+crazywoola@users.noreply.github.com>
Date: Fri, 11 Oct 2024 11:28:09 +0800
Subject: [PATCH] feat: add NotFound error for dataset service (#9215)

---
 api/services/dataset_service.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/api/services/dataset_service.py b/api/services/dataset_service.py
index b8f80a9f77..4f77fdc18c 100644
--- a/api/services/dataset_service.py
+++ b/api/services/dataset_service.py
@@ -8,6 +8,7 @@ from typing import Optional
 
 from flask_login import current_user
 from sqlalchemy import func
+from werkzeug.exceptions import NotFound
 
 from configs import dify_config
 from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError
@@ -975,6 +976,8 @@ class DocumentService:
     ):
         DatasetService.check_dataset_model_setting(dataset)
         document = DocumentService.get_document(dataset.id, document_data["original_document_id"])
+        if document is None:
+            raise NotFound("Document not found")
         if document.display_status != "available":
             raise ValueError("Document is not available")
         # update document name