From fe44e36609dc2a97a22ecfe0f08c6facfd244760 Mon Sep 17 00:00:00 2001
From: huweiqiang <huweiqiang@idea.edu.cn>
Date: Wed, 13 Dec 2023 11:13:09 +0800
Subject: [PATCH] refactor(delete dataset): add model API to cascade delete
 dataset

---
 deepdataspace/model/dataset.py        | 23 +++++++++++++++++++++++
 deepdataspace/scripts/dataset_cmds.py | 19 +------------------
 2 files changed, 24 insertions(+), 18 deletions(-)

diff --git a/deepdataspace/model/dataset.py b/deepdataspace/model/dataset.py
index dba8e38..b99f690 100644
--- a/deepdataspace/model/dataset.py
+++ b/deepdataspace/model/dataset.py
@@ -439,3 +439,26 @@ def eval_description(self):
             except:
                 logger.warning(f"Failed to eval description_func[{self.description_func}] for dataset[{self.id}]")
                 return self.description or self.path
+
+    @staticmethod
+    def cascade_delete(dataset: "DataSet"):
+        """
+        Cascade delete the dataset, along with all its images, labels, categories and objects.
+        """
+        if dataset is None:
+            return
+
+        dataset_id = dataset.id
+        print(f"dataset [{dataset_id}] is found, deleting...")
+
+        print(f"dataset [{dataset_id}] is found, deleting categories...")
+        Category.delete_many({"dataset_id": dataset_id})
+
+        print(f"dataset [{dataset_id}] is found, deleting labels...")
+        Label.delete_many({"dataset_id": dataset_id})
+
+        print(f"dataset [{dataset_id}] is found, deleting images...")
+        Image(dataset_id).get_collection().drop()
+
+        DataSet.delete_many({"id": dataset_id})
+        print(f"dataset [{dataset_id}] is deleted.")
diff --git a/deepdataspace/scripts/dataset_cmds.py b/deepdataspace/scripts/dataset_cmds.py
index 3168d6d..427f747 100644
--- a/deepdataspace/scripts/dataset_cmds.py
+++ b/deepdataspace/scripts/dataset_cmds.py
@@ -5,7 +5,6 @@
 """
 
 import os
-import shutil
 
 import pkg_resources
 
@@ -17,10 +16,7 @@
 @ddsop.command("delete_one", help="Delete a dataset.")
 @click.argument("dataset_dir")
 def delete_one(dataset_dir):
-    from deepdataspace.model import Image
-    from deepdataspace.model import Label
     from deepdataspace.model import DataSet
-    from deepdataspace.model import Category
     from deepdataspace.utils.string import get_str_md5
 
     dataset_dir = os.path.abspath(dataset_dir)
@@ -30,20 +26,7 @@ def delete_one(dataset_dir):
         print(f"dataset [{dataset_dir}] is not imported before, skip...")
         return
 
-    dataset_id = dataset.id
-    print(f"dataset [{dataset_id}] is found, deleting...")
-
-    print(f"dataset [{dataset_id}] is found, deleting categories...")
-    Category.delete_many({"dataset_id": dataset_id})
-
-    print(f"dataset [{dataset_id}] is found, deleting labels...")
-    Label.delete_many({"dataset_id": dataset_id})
-
-    print(f"dataset [{dataset_id}] is found, deleting images...")
-    Image(dataset_id).get_collection().drop()
-
-    DataSet.delete_many({"id": dataset_id})
-    print(f"dataset [{dataset_id}] is deleted.")
+    DataSet.cascade_delete(dataset)
 
 
 @ddsop.command("import_all", help="Trigger a background task of importing all datasets in a data dir.")