cleanup

fiftyone/core/dataset.py

@@ -57,6 +57,7 @@
 
 fot = fou.lazy_import("fiftyone.core.stages")
 foud = fou.lazy_import("fiftyone.utils.data")
+foos = fou.lazy_import("fiftyone.operators.store")
 
 
 _SUMMARY_FIELD_KEY = "_summary_field"
@@ -5180,13 +5181,13 @@ def _delete(self):
             self._frame_collection.drop()
             fofr.Frame._reset_docs(self._frame_collection_name)
 
+        foos.cleanup_store_for_dataset(self._doc.id)
+
         # Update singleton
         self._instances.pop(self._doc.name, None)
         _delete_dataset_doc(self._doc)
         self._deleted = True
 
-        _cleanup_execution_store_for_dataset(self._doc.id)
-
     def add_dir(
         self,
         dataset_dir=None,
@@ -10177,11 +10178,3 @@ def _extract_archive_if_necessary(archive_path, cleanup):
         )
 
     return dataset_dir
-
-
-def _cleanup_execution_store_for_dataset(dataset_id):
-    """Cleans up the execution store for the given dataset."""
-    from fiftyone.operators.store import ExecutionStoreService
-
-    svc = ExecutionStoreService(dataset_id=dataset_id)
-    svc.cleanup_for_dataset()

fiftyone/core/odm/database.py

@@ -37,6 +37,7 @@
 fod = fou.lazy_import("fiftyone.core.dataset")
 foe = fou.lazy_import("fiftyone.core.evaluation")
 fors = fou.lazy_import("fiftyone.core.runs")
+foos = fou.lazy_import("fiftyone.operators.store")
 
 
 logger = logging.getLogger(__name__)
@@ -1126,6 +1127,16 @@ def delete_dataset(name, dry_run=False):
         if not dry_run:
             _delete_run_results(conn, result_ids)
 
+    dataset_id = dataset_dict["_id"]
+    svc = foos.ExecutionStoreService(dataset_id=dataset_id)
+    num_docs = sum(
+        len(svc.list_keys(store_name)) for store_name in svc.list_stores()
+    )
+    if num_docs > 0:
+        _logger.info("Deleting %d store doc(s)", num_docs)
+        if not dry_run:
+            foos.cleanup_store_for_dataset(dataset_id)
+
 
 def delete_saved_view(dataset_name, view_name, dry_run=False):
     """Deletes the saved view with the given name from the dataset with the

fiftyone/operators/store/__init__.py

@@ -5,14 +5,16 @@
 | `voxel51.com <https://voxel51.com/>`_
 |
 """
+import types
 
-from .service import ExecutionStoreService
+from .service import cleanup_store_for_dataset, ExecutionStoreService
 from .store import ExecutionStore
 from .models import StoreDocument, KeyDocument
 
+# This tells Sphinx to allow refs to imported objects in this module
+# https://stackoverflow.com/a/31594545/16823653
 __all__ = [
-    "ExecutionStoreService",
-    "StoreDocument",
-    "KeyDocument",
-    "ExecutionStore",
+    k
+    for k, v in globals().items()
+    if not k.startswith("_") and not isinstance(v, types.ModuleType)
 ]

fiftyone/operators/store/service.py

@@ -9,12 +9,25 @@
 import bson
 import logging
 from typing import Optional, List
+
 from fiftyone.operators.store.models import StoreDocument, KeyDocument
 
+
 logger = logging.getLogger(__name__)
 
 
-class ExecutionStoreService:
+def cleanup_store_for_dataset(dataset_id):
+    """Deletes all documents in the execution store associated with the given
+    dataset.
+
+    Args:
+        dataset_id: the dataset ID
+    """
+    svc = ExecutionStoreService(dataset_id=dataset_id)
+    svc.cleanup_for_dataset()
+
+
+class ExecutionStoreService(object):
     def __init__(
         self,
         repo: Optional["ExecutionStoreRepo"] = None,
@@ -23,8 +36,8 @@ def __init__(
         """Service for managing execution store operations.
 
         Args:
-            repo: Optional execution store repository instance
-            dataset_id: Optional bson.ObjectId of the dataset to scope operations to
+            repo (None): execution store repository instance
+            dataset_id (None): dataset ID to scope operations to
         """
         from fiftyone.factory.repo_factory import (
             RepositoryFactory,
@@ -139,5 +152,7 @@ def list_keys(self, store_name: str) -> List[str]:
         return self._repo.list_keys(store_name)
 
     def cleanup_for_dataset(self) -> None:
-        """Cleans up the execution store for the dataset specified during initialization."""
+        """Cleans up the execution store for the dataset specified during
+        initialization.
+        """
         self._repo.cleanup_for_dataset()

fiftyone/operators/store/store.py

@@ -8,21 +8,15 @@
 
 import bson
 import logging
-from fiftyone.operators.store.service import ExecutionStoreService
 from typing import Any, Optional
 
-logger = logging.getLogger(__name__)
+from fiftyone.operators.store.service import ExecutionStoreService
 
 
-class ExecutionStore:
-    @staticmethod
-    def create(
-        store_name: str, dataset_id: Optional[bson.ObjectId] = None
-    ) -> "ExecutionStore":
-        return ExecutionStore(
-            store_name, ExecutionStoreService(dataset_id=dataset_id)
-        )
+logger = logging.getLogger(__name__)
+
 
+class ExecutionStore(object):
     def __init__(self, store_name: str, store_service: ExecutionStoreService):
         """
         Args:
@@ -32,22 +26,30 @@ def __init__(self, store_name: str, store_service: ExecutionStoreService):
         self.store_name: str = store_name
         self._store_service: ExecutionStoreService = store_service
 
+    @staticmethod
+    def create(
+        store_name: str, dataset_id: Optional[bson.ObjectId] = None
+    ) -> "ExecutionStore":
+        return ExecutionStore(
+            store_name, ExecutionStoreService(dataset_id=dataset_id)
+        )
+
     def list_all_stores(self) -> list[str]:
         """Lists all stores in the execution store.
 
         Returns:
-            list: A list of store names.
+            list: a list of store names
         """
         return self._store_service.list_stores()
 
     def get(self, key: str) -> Optional[Any]:
         """Retrieves a value from the store by its key.
 
         Args:
-            key (str): The key to retrieve the value for.
+            key: the key to retrieve the value for
 
         Returns:
-            Optional[Any]: The value stored under the given key, or None if not found.
+            the value stored under the given key, or None if not found
         """
         key_doc = self._store_service.get_key(self.store_name, key)
         if key_doc is None:
@@ -58,31 +60,31 @@ def set(self, key: str, value: Any, ttl: Optional[int] = None) -> None:
         """Sets a value in the store with an optional TTL.
 
         Args:
-            key (str): The key to store the value under.
-            value (Any): The value to store.
-            ttl (Optional[int], optional): The time-to-live in seconds. Defaults to None.
+            key: the key to store the value under
+            value: the value to store
+            ttl (None): the time-to-live in seconds
         """
         self._store_service.set_key(self.store_name, key, value, ttl)
 
     def delete(self, key: str) -> bool:
         """Deletes a key from the store.
 
         Args:
-            key (str): The key to delete.
+            key: the key to delete.
 
         Returns:
-            bool: True if the key was deleted, False otherwise.
+            True/False whether the key was deleted
         """
         return self._store_service.delete_key(self.store_name, key)
 
     def has(self, key: str) -> bool:
         """Checks if the store has a specific key.
 
         Args:
-            key (str): The key to check.
+            key: the key to check
 
         Returns:
-            bool: True if the key exists, False otherwise.
+            True/False whether the key exists
         """
         return self._store_service.has_key(self.store_name, key)
 
@@ -94,26 +96,26 @@ def update_ttl(self, key: str, new_ttl: int) -> None:
         """Updates the TTL for a specific key.
 
         Args:
-            key (str): The key to update the TTL for.
-            new_ttl (int): The new TTL in seconds.
+            key: the key to update the TTL for
+            new_ttl: the new TTL in seconds
         """
         self._store_service.update_ttl(self.store_name, key, new_ttl)
 
     def get_ttl(self, key: str) -> Optional[int]:
         """Retrieves the TTL for a specific key.
 
         Args:
-            key (str): The key to get the TTL for.
+            key: the key to get the TTL for
 
         Returns:
-            Optional[int]: The TTL in seconds, or None if the key does not have a TTL.
+            the TTL in seconds, or None if the key does not have a TTL
         """
         return self._store_service.get_ttl(self.store_name, key)
 
     def list_keys(self) -> list[str]:
         """Lists all keys in the store.
 
         Returns:
-            list: A list of keys in the store.
+            a list of keys in the store
         """
         return self._store_service.list_keys(self.store_name)

tests/unittests/dataset_tests.py

@@ -26,6 +26,7 @@
 import fiftyone as fo
 import fiftyone.core.fields as fof
 import fiftyone.core.odm as foo
+from fiftyone.operators.store import ExecutionStoreService
 import fiftyone.utils.data as foud
 from fiftyone import ViewField as F
 
@@ -123,8 +124,6 @@ def test_load_dataset(self):
 
     @drop_datasets
     def test_delete_dataset(self):
-        from fiftyone.operators.store import ExecutionStoreService
-
         IGNORED_DATASET_NAMES = fo.list_datasets()
 
         def list_datasets():
@@ -140,27 +139,35 @@ def list_datasets():
         self.assertListEqual(list_datasets(), dataset_names)
 
         name = dataset_names.pop(0)
+
         dataset = datasets[name]
         dataset_id = dataset._doc.id
-        exec_store_svc = ExecutionStoreService(dataset_id=dataset_id)
+        svc = ExecutionStoreService(dataset_id=dataset_id)
         store_name = "test_store"
-        exec_store_svc.set_key(store_name, "foo", "bar")
-        keys = exec_store_svc.list_keys(store_name)
-        self.assertEqual(keys, ["foo"])
+        svc.set_key(store_name, "foo", "bar")
+        keys = svc.list_keys(store_name)
+
+        self.assertListEqual(keys, ["foo"])
+
         dataset.delete()
+
         self.assertListEqual(list_datasets(), dataset_names)
-        keys = exec_store_svc.list_keys(store_name)
-        self.assertEqual(keys, [])
+        keys = svc.list_keys(store_name)
+
+        self.assertListEqual(keys, [])
         with self.assertRaises(ValueError):
-            len(datasets[name])
+            len(dataset)
 
         name = dataset_names.pop(0)
+
         fo.delete_dataset(name)
+
         self.assertListEqual(list_datasets(), dataset_names)
         with self.assertRaises(ValueError):
             len(datasets[name])
 
         new_dataset = fo.Dataset(name)
+
         self.assertEqual(len(new_dataset), 0)
 
     @drop_datasets