From 7f98c2ea3fec1670bd23ad0b956d082553ccba03 Mon Sep 17 00:00:00 2001 From: Summer-Gu <37869445+gubinjie@users.noreply.github.com> Date: Fri, 14 Jun 2024 03:25:38 +0800 Subject: [PATCH] refactor: Delete the dataset to verify whether it is in use (#5112) --- api/controllers/console/datasets/datasets.py | 13 ++++++++----- api/controllers/console/datasets/error.py | 6 ++++++ api/controllers/service_api/dataset/dataset.py | 13 ++++++++----- api/controllers/service_api/dataset/error.py | 6 ++++++ api/services/dataset_service.py | 6 ++++-- api/services/errors/dataset.py | 4 ++++ 6 files changed, 36 insertions(+), 12 deletions(-) diff --git a/api/controllers/console/datasets/datasets.py b/api/controllers/console/datasets/datasets.py index 6cbe454d37..c014fd6746 100644 --- a/api/controllers/console/datasets/datasets.py +++ b/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): diff --git a/api/controllers/console/datasets/error.py b/api/controllers/console/datasets/error.py index 29142b80e6..e77693b6c9 100644 --- a/api/controllers/console/datasets/error.py +++ b/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 diff --git a/api/controllers/service_api/dataset/dataset.py b/api/controllers/service_api/dataset/dataset.py index dcbfa88d99..8dd16c0787 100644 --- a/api/controllers/service_api/dataset/dataset.py +++ b/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/') diff --git a/api/controllers/service_api/dataset/error.py b/api/controllers/service_api/dataset/error.py index 29142b80e6..e77693b6c9 100644 --- a/api/controllers/service_api/dataset/error.py +++ b/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 diff --git a/api/services/dataset_service.py b/api/services/dataset_service.py index 06d3e9ec40..97a6f5d103 100644 --- a/api/services/dataset_service.py +++ b/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) diff --git a/api/services/errors/dataset.py b/api/services/errors/dataset.py index 254a70ffe3..d36cd1111c 100644 --- a/api/services/errors/dataset.py +++ b/api/services/errors/dataset.py @@ -3,3 +3,7 @@ from services.errors.base import BaseServiceError class DatasetNameDuplicateError(BaseServiceError): pass + + +class DatasetInUseError(BaseServiceError): + pass