From 71bcf75d9a089c9bfb9763de87907eba640e35fd Mon Sep 17 00:00:00 2001
From: Jyong <76649700+JohnJyong@users.noreply.github.com>
Date: Mon, 1 Jul 2024 17:06:51 +0800
Subject: [PATCH] Feat/add delete knowledge confirm (#5810)

---
 api/controllers/console/datasets/datasets.py | 10 ++++++++++
 api/services/dataset_service.py              | 12 ++++++++----
 2 files changed, 18 insertions(+), 4 deletions(-)

diff --git a/api/controllers/console/datasets/datasets.py b/api/controllers/console/datasets/datasets.py
index d5196aae61..fdd61b0a0c 100644
--- a/api/controllers/console/datasets/datasets.py
+++ b/api/controllers/console/datasets/datasets.py
@@ -226,6 +226,15 @@ class DatasetApi(Resource):
         except services.errors.dataset.DatasetInUseError:
             raise DatasetInUseError()
 
+class DatasetUseCheckApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    def get(self, dataset_id):
+        dataset_id_str = str(dataset_id)
+
+        dataset_is_using = DatasetService.dataset_use_check(dataset_id_str)
+        return {'is_using': dataset_is_using}, 200
 
 class DatasetQueryApi(Resource):
 
@@ -562,6 +571,7 @@ class DatasetErrorDocs(Resource):
 
 api.add_resource(DatasetListApi, '/datasets')
 api.add_resource(DatasetApi, '/datasets/<uuid:dataset_id>')
+api.add_resource(DatasetUseCheckApi, '/datasets/<uuid:dataset_id>/use-check')
 api.add_resource(DatasetQueryApi, '/datasets/<uuid:dataset_id>/queries')
 api.add_resource(DatasetErrorDocs, '/datasets/<uuid:dataset_id>/error-docs')
 api.add_resource(DatasetIndexingEstimateApi, '/datasets/indexing-estimate')
diff --git a/api/services/dataset_service.py b/api/services/dataset_service.py
index 38ef874af3..6207a1a45c 100644
--- a/api/services/dataset_service.py
+++ b/api/services/dataset_service.py
@@ -34,7 +34,7 @@ from models.dataset import (
 from models.model import UploadFile
 from models.source import DataSourceOauthBinding
 from services.errors.account import NoPermissionError
-from services.errors.dataset import DatasetInUseError, DatasetNameDuplicateError
+from services.errors.dataset import DatasetNameDuplicateError
 from services.errors.document import DocumentIndexingError
 from services.errors.file import FileNotExistsError
 from services.feature_service import FeatureModel, FeatureService
@@ -234,9 +234,6 @@ class DatasetService:
 
     @staticmethod
     def delete_dataset(dataset_id, user):
-        count = AppDatasetJoin.query.filter_by(dataset_id=dataset_id).count()
-        if count > 0:
-            raise DatasetInUseError()
 
         dataset = DatasetService.get_dataset(dataset_id)
 
@@ -251,6 +248,13 @@ class DatasetService:
         db.session.commit()
         return True
 
+    @staticmethod
+    def dataset_use_check(dataset_id) -> bool:
+        count = AppDatasetJoin.query.filter_by(dataset_id=dataset_id).count()
+        if count > 0:
+            return True
+        return False
+
     @staticmethod
     def check_dataset_permission(dataset, user):
         if dataset.tenant_id != user.current_tenant_id: