diff --git a/fiftyone/core/dataset.py b/fiftyone/core/dataset.py index 7b65f9487d..df0e2ecce9 100644 --- a/fiftyone/core/dataset.py +++ b/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() diff --git a/fiftyone/core/odm/database.py b/fiftyone/core/odm/database.py index f6820a96ca..f54b291f63 100644 --- a/fiftyone/core/odm/database.py +++ b/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 diff --git a/fiftyone/operators/store/__init__.py b/fiftyone/operators/store/__init__.py index 50f9f3b400..a586088529 100644 --- a/fiftyone/operators/store/__init__.py +++ b/fiftyone/operators/store/__init__.py @@ -5,14 +5,16 @@ | `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) ] diff --git a/fiftyone/operators/store/service.py b/fiftyone/operators/store/service.py index 305e35514e..2e67690af7 100644 --- a/fiftyone/operators/store/service.py +++ b/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() diff --git a/fiftyone/operators/store/store.py b/fiftyone/operators/store/store.py index be87e7a110..cb53f696fb 100644 --- a/fiftyone/operators/store/store.py +++ b/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,11 +26,19 @@ 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() @@ -44,10 +46,10 @@ 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,9 +60,9 @@ 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) @@ -68,10 +70,10 @@ 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) @@ -79,10 +81,10 @@ 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,8 +96,8 @@ 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) @@ -103,10 +105,10 @@ 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) @@ -114,6 +116,6 @@ 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) diff --git a/tests/unittests/dataset_tests.py b/tests/unittests/dataset_tests.py index b64c5bd394..9ef306e471 100644 --- a/tests/unittests/dataset_tests.py +++ b/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