Bläddra i källkod

refactor: Delete the dataset to verify whether it is in use (#5112)

Summer-Gu 10 månader sedan
förälder
incheckning
7f98c2ea3f

+ 8 - 5
api/controllers/console/datasets/datasets.py

@@ -8,7 +8,7 @@ import services
 from controllers.console import api
 from controllers.console.apikey import api_key_fields, api_key_list
 from controllers.console.app.error import ProviderNotInitializeError
-from controllers.console.datasets.error import DatasetNameDuplicateError
+from controllers.console.datasets.error import DatasetInUseError, DatasetNameDuplicateError
 from controllers.console.setup import setup_required
 from controllers.console.wraps import account_initialization_required
 from core.errors.error import LLMBadRequestError, ProviderTokenNotInitError
@@ -217,10 +217,13 @@ class DatasetApi(Resource):
         if not current_user.is_admin_or_owner:
             raise Forbidden()
 
-        if DatasetService.delete_dataset(dataset_id_str, current_user):
-            return {'result': 'success'}, 204
-        else:
-            raise NotFound("Dataset not found.")
+        try:
+            if DatasetService.delete_dataset(dataset_id_str, current_user):
+                return {'result': 'success'}, 204
+            else:
+                raise NotFound("Dataset not found.")
+        except services.errors.dataset.DatasetInUseError:
+            raise DatasetInUseError()
 
 
 class DatasetQueryApi(Resource):

+ 6 - 0
api/controllers/console/datasets/error.py

@@ -71,3 +71,9 @@ class InvalidMetadataError(BaseHTTPException):
     error_code = 'invalid_metadata'
     description = "The metadata content is incorrect. Please check and verify."
     code = 400
+
+
+class DatasetInUseError(BaseHTTPException):
+    error_code = 'dataset_in_use'
+    description = "The dataset is being used by some apps. Please remove the dataset from the apps before deleting it."
+    code = 409

+ 8 - 5
api/controllers/service_api/dataset/dataset.py

@@ -4,7 +4,7 @@ from werkzeug.exceptions import NotFound
 
 import services.dataset_service
 from controllers.service_api import api
-from controllers.service_api.dataset.error import DatasetNameDuplicateError
+from controllers.service_api.dataset.error import DatasetInUseError, DatasetNameDuplicateError
 from controllers.service_api.wraps import DatasetApiResource
 from core.model_runtime.entities.model_entities import ModelType
 from core.provider_manager import ProviderManager
@@ -113,10 +113,13 @@ class DatasetApi(DatasetApiResource):
 
         dataset_id_str = str(dataset_id)
 
-        if DatasetService.delete_dataset(dataset_id_str, current_user):
-            return {'result': 'success'}, 204
-        else:
-            raise NotFound("Dataset not found.")
+        try:
+            if DatasetService.delete_dataset(dataset_id_str, current_user):
+                return {'result': 'success'}, 204
+            else:
+                raise NotFound("Dataset not found.")
+        except services.errors.dataset.DatasetInUseError:
+            raise DatasetInUseError()
 
 api.add_resource(DatasetListApi, '/datasets')
 api.add_resource(DatasetApi, '/datasets/<uuid:dataset_id>')

+ 6 - 0
api/controllers/service_api/dataset/error.py

@@ -71,3 +71,9 @@ class InvalidMetadataError(BaseHTTPException):
     error_code = 'invalid_metadata'
     description = "The metadata content is incorrect. Please check and verify."
     code = 400
+
+
+class DatasetInUseError(BaseHTTPException):
+    error_code = 'dataset_in_use'
+    description = "The dataset is being used by some apps. Please remove the dataset from the apps before deleting it."
+    code = 409

+ 4 - 2
api/services/dataset_service.py

@@ -33,7 +33,7 @@ from models.dataset import (
 from models.model import UploadFile
 from models.source import DataSourceBinding
 from services.errors.account import NoPermissionError
-from services.errors.dataset import DatasetNameDuplicateError
+from services.errors.dataset import DatasetInUseError, DatasetNameDuplicateError
 from services.errors.document import DocumentIndexingError
 from services.errors.file import FileNotExistsError
 from services.feature_service import FeatureModel, FeatureService
@@ -232,7 +232,9 @@ class DatasetService:
 
     @staticmethod
     def delete_dataset(dataset_id, user):
-        # todo: cannot delete dataset if it is being processed
+        count = AppDatasetJoin.query.filter_by(dataset_id=dataset_id).count()
+        if count > 0:
+            raise DatasetInUseError()
 
         dataset = DatasetService.get_dataset(dataset_id)
 

+ 4 - 0
api/services/errors/dataset.py

@@ -3,3 +3,7 @@ from services.errors.base import BaseServiceError
 
 class DatasetNameDuplicateError(BaseServiceError):
     pass
+
+
+class DatasetInUseError(BaseServiceError):
+    pass