From 22bdea79a86481960f68402ece9ef840d97cd8e7 Mon Sep 17 00:00:00 2001
From: -LAN- <laipz8200@outlook.com>
Date: Thu, 26 Dec 2024 19:31:51 +0800
Subject: [PATCH] fix: enhance file upload error handling and update base error
 class

Signed-off-by: -LAN- <laipz8200@outlook.com>
---
 .../service_api/dataset/document.py           | 29 ++++++++++++-------
 api/services/errors/base.py                   |  2 +-
 2 files changed, 20 insertions(+), 11 deletions(-)

diff --git a/api/controllers/service_api/dataset/document.py b/api/controllers/service_api/dataset/document.py
index 84c58c62df5b3c..ea664b8f1be4d4 100644
--- a/api/controllers/service_api/dataset/document.py
+++ b/api/controllers/service_api/dataset/document.py
@@ -8,12 +8,16 @@
 import services.dataset_service
 from controllers.common.errors import FilenameNotExistsError
 from controllers.service_api import api
-from controllers.service_api.app.error import ProviderNotInitializeError
+from controllers.service_api.app.error import (
+    FileTooLargeError,
+    NoFileUploadedError,
+    ProviderNotInitializeError,
+    TooManyFilesError,
+    UnsupportedFileTypeError,
+)
 from controllers.service_api.dataset.error import (
     ArchivedDocumentImmutableError,
     DocumentIndexingError,
-    NoFileUploadedError,
-    TooManyFilesError,
 )
 from controllers.service_api.wraps import DatasetApiResource, cloud_edition_billing_resource_check
 from core.errors.error import ProviderTokenNotInitError
@@ -238,13 +242,18 @@ def post(self, tenant_id, dataset_id, document_id):
             if not file.filename:
                 raise FilenameNotExistsError
 
-            upload_file = FileService.upload_file(
-                filename=file.filename,
-                content=file.read(),
-                mimetype=file.mimetype,
-                user=current_user,
-                source="datasets",
-            )
+            try:
+                upload_file = FileService.upload_file(
+                    filename=file.filename,
+                    content=file.read(),
+                    mimetype=file.mimetype,
+                    user=current_user,
+                    source="datasets",
+                )
+            except services.errors.file.FileTooLargeError as file_too_large_error:
+                raise FileTooLargeError(file_too_large_error.description)
+            except services.errors.file.UnsupportedFileTypeError:
+                raise UnsupportedFileTypeError()
             data_source = {"type": "upload_file", "info_list": {"file_info_list": {"file_ids": [upload_file.id]}}}
             args["data_source"] = data_source
         # validate args
diff --git a/api/services/errors/base.py b/api/services/errors/base.py
index 4d39f956b8c932..35ea28468e0d86 100644
--- a/api/services/errors/base.py
+++ b/api/services/errors/base.py
@@ -1,6 +1,6 @@
 from typing import Optional
 
 
-class BaseServiceError(Exception):
+class BaseServiceError(ValueError):
     def __init__(self, description: Optional[str] = None):
         self.description = description