From 78fa4b9f1d8bb7c3128cbc073d41e142bc5ae34c Mon Sep 17 00:00:00 2001
From: Ritchie Martori <ritchie@voxel51.com>
Date: Sun, 22 Sep 2024 16:43:06 -0700
Subject: [PATCH] add store listing

---
 fiftyone/factory/repos/execution_store.py | 7 ++++++-
 fiftyone/operators/store/models.py        | 2 +-
 fiftyone/operators/store/service.py       | 7 ++-----
 fiftyone/operators/store/store.py         | 8 ++++++++
 tests/unittests/execution_store_tests.py  | 7 +++++++
 5 files changed, 24 insertions(+), 7 deletions(-)

diff --git a/fiftyone/factory/repos/execution_store.py b/fiftyone/factory/repos/execution_store.py
index 0b2e5aad9c..7f2c4aa117 100644
--- a/fiftyone/factory/repos/execution_store.py
+++ b/fiftyone/factory/repos/execution_store.py
@@ -1,5 +1,5 @@
 """
-Execution store repository for handling storage and retrieval of stores.
+Execution store repository.
 """
 
 from pymongo.collection import Collection
@@ -29,6 +29,11 @@ def create_store(self, store_name, permissions=None) -> StoreDocument:
         self._collection.insert_one(store_doc.dict())
         return store_doc
 
+    def list_stores(self) -> list[str]:
+        """Lists all stores in the execution store."""
+        # ensure that only store_name is returned, and only unique values
+        return self._collection.distinct("store_name")
+
     def set_key(self, store_name, key, value, ttl=None) -> KeyDocument:
         """Sets or updates a key in the specified store."""
         expiration = KeyDocument.get_expiration(ttl)
diff --git a/fiftyone/operators/store/models.py b/fiftyone/operators/store/models.py
index e05767fe22..5fd6c89e21 100644
--- a/fiftyone/operators/store/models.py
+++ b/fiftyone/operators/store/models.py
@@ -29,7 +29,7 @@ def get_expiration(ttl: Optional[int]) -> Optional[datetime.datetime]:
 
 
 class StoreDocument(KeyDocument):
-    """Model representing a store in the execution store."""
+    """Model representing a Store."""
 
     key: str = "__store__"
     value: Optional[Dict[str, Any]] = None
diff --git a/fiftyone/operators/store/service.py b/fiftyone/operators/store/service.py
index f01e3b874f..dc90039985 100644
--- a/fiftyone/operators/store/service.py
+++ b/fiftyone/operators/store/service.py
@@ -87,16 +87,13 @@ def update_ttl(self, store_name, key, new_ttl):
             store_name=store_name, key=key, ttl=new_ttl
         )
 
-    def list_stores(self, search=None, **kwargs):
+    def list_stores(self):
         """Lists all stores matching the given criteria.
 
-        Args:
-            search (None): optional search term dict
-
         Returns:
             a list of :class:`fiftyone.store.models.StoreDocument`
         """
-        return self._repo.list_stores(search=search, **kwargs)
+        return self._repo.list_stores()
 
     def delete_store(self, store_name):
         """Deletes the specified store.
diff --git a/fiftyone/operators/store/store.py b/fiftyone/operators/store/store.py
index 55ae6a15d1..15e0ab3a01 100644
--- a/fiftyone/operators/store/store.py
+++ b/fiftyone/operators/store/store.py
@@ -28,6 +28,14 @@ def __init__(self, store_name: str, store_service: ExecutionStoreService):
         self.store_name: str = store_name
         self._store_service: ExecutionStoreService = store_service
 
+    def list_all_stores(self) -> list[str]:
+        """Lists all stores in the execution store.
+
+        Returns:
+            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.
 
diff --git a/tests/unittests/execution_store_tests.py b/tests/unittests/execution_store_tests.py
index e0cc1bc495..d5b145ba63 100644
--- a/tests/unittests/execution_store_tests.py
+++ b/tests/unittests/execution_store_tests.py
@@ -142,6 +142,13 @@ def test_list_keys(self):
         self.mock_collection.find.assert_called_once()
         self.mock_collection.find.assert_called_with({"store_name": "widgets"})
 
+    def test_list_stores(self):
+        self.mock_collection.distinct.return_value = ["widgets", "gadgets"]
+        stores = self.store_repo.list_stores()
+        assert stores == ["widgets", "gadgets"]
+        self.mock_collection.distinct.assert_called_once()
+        self.mock_collection.distinct.assert_called_with("store_name")
+
 
 class TestExecutionStoreIntegration(unittest.TestCase):
     def setUp(self) -> None: