Update cleanup_lab.py

CederGroupHub · Mar 26, 2024 · a36009a · a36009a
1 parent 19a9951
commit a36009a
Showing 1 changed file with 9 additions and 5 deletions.
diff --git a/alab_management/scripts/cleanup_lab.py b/alab_management/scripts/cleanup_lab.py
@@ -26,6 +26,8 @@ def cleanup_lab(
         .get_collection("tasks")
         .count_documents({})
     )
+
+    whole_database_dropped = False
     if all_collections:
         if not _force_i_know_its_dangerous:
             if user_confirmation is None:
@@ -45,10 +47,12 @@ def cleanup_lab(
             print(f"Removing database {database_name}")
             _GetMongoCollection.init()
             _GetMongoCollection.client.drop_database(database_name)  # type: ignore
+            whole_database_dropped = True
         else:
             print(f"Removing database {database_name}")
             _GetMongoCollection.init()
             _GetMongoCollection.client.drop_database(database_name)
+            whole_database_dropped = True
 
         # if sim_mode != AlabOSConfig().is_sim_mode() or database_name == "Alab":
         #     print("Wrong name of database. Hence, not removed.")
@@ -60,9 +64,9 @@ def cleanup_lab(
         # else:
         #     print("Wrong name of database. Hence, not removed.")
         #     return False
-
-    DeviceView()._clean_up_device_collection()
-    SampleView().clean_up_sample_position_collection()
-    _GetMongoCollection.get_collection("_lock").drop()
-    _GetMongoCollection.get_collection("requests").drop()
+    if not whole_database_dropped:
+        DeviceView()._clean_up_device_collection()
+        SampleView().clean_up_sample_position_collection()
+        _GetMongoCollection.get_collection("_lock").drop()
+        _GetMongoCollection.get_collection("requests").drop()
     return True