From 5343ac76a67fb2f52074f880f0c315edad668fea Mon Sep 17 00:00:00 2001 From: Ritchie Martori Date: Fri, 20 Sep 2024 08:28:40 -0700 Subject: [PATCH 1/9] initial execution store --- fiftyone/factory/repo_factory.py | 19 +++ fiftyone/factory/repos/execution_store.py | 64 ++++++++ fiftyone/operators/store/__init__.py | 18 +++ fiftyone/operators/store/models.py | 64 ++++++++ fiftyone/operators/store/permissions.py | 47 ++++++ fiftyone/operators/store/store.py | 131 +++++++++++++++++ tests/unittests/execution_store_tests.py | 170 ++++++++++++++++++++++ 7 files changed, 513 insertions(+) create mode 100644 fiftyone/factory/repos/execution_store.py create mode 100644 fiftyone/operators/store/__init__.py create mode 100644 fiftyone/operators/store/models.py create mode 100644 fiftyone/operators/store/permissions.py create mode 100644 fiftyone/operators/store/store.py create mode 100644 tests/unittests/execution_store_tests.py diff --git a/fiftyone/factory/repo_factory.py b/fiftyone/factory/repo_factory.py index 2b031588da..314163544e 100644 --- a/fiftyone/factory/repo_factory.py +++ b/fiftyone/factory/repo_factory.py @@ -13,6 +13,10 @@ DelegatedOperationRepo, MongoDelegatedOperationRepo, ) +from fiftyone.factory.repos.execution_store import ( + ExecutionStoreRepo, + MongoExecutionStoreRepo, +) _db: Database = None @@ -44,3 +48,18 @@ def delegated_operation_repo() -> DelegatedOperationRepo: return RepositoryFactory.repos[ MongoDelegatedOperationRepo.COLLECTION_NAME ] + + @staticmethod + def execution_store_repo() -> ExecutionStoreRepo: + """Factory method for execution store repository.""" + if ( + MongoExecutionStoreRepo.COLLECTION_NAME + not in RepositoryFactory.repos + ): + RepositoryFactory.repos[ + MongoExecutionStoreRepo.COLLECTION_NAME + ] = MongoExecutionStoreRepo( + collection=_get_db()[MongoExecutionStoreRepo.COLLECTION_NAME] + ) + + return RepositoryFactory.repos[MongoExecutionStoreRepo.COLLECTION_NAME] diff --git a/fiftyone/factory/repos/execution_store.py b/fiftyone/factory/repos/execution_store.py new file mode 100644 index 0000000000..ab56ccc620 --- /dev/null +++ b/fiftyone/factory/repos/execution_store.py @@ -0,0 +1,64 @@ +""" +Execution store repository for handling storage and retrieval of stores. +""" + +from pymongo.collection import Collection +from fiftyone.operators.store.models import StoreDocument, KeyDocument + + +class ExecutionStoreRepo: + """Base class for execution store repositories.""" + + COLLECTION_NAME = "execution_store" + + def __init__(self, collection: Collection): + self._collection = collection + + def create_store(self, store_name, permissions=None): + """Creates a store in the execution store.""" + store_doc = StoreDocument( + store_name=store_name, permissions=permissions + ) + self._collection.insert_one(store_doc.dict()) + return store_doc + + def set_key(self, store_name, key, value, ttl=None): + """Sets a key in the specified store.""" + key_doc = KeyDocument(key=key, value=value, ttl=ttl) + self._collection.update_one( + {"store_name": store_name}, + {"$set": {f"keys.{key}": key_doc.dict()}}, + ) + return key_doc + + def get_key(self, store_name, key): + """Gets a key from the specified store.""" + store = self._collection.find_one({"store_name": store_name}) + if store and key in store.get("keys", {}): + return KeyDocument(**store["keys"][key]) + return None + + def update_ttl(self, store_name, key, ttl): + """Updates the TTL for a key.""" + self._collection.update_one( + {"store_name": store_name}, {"$set": {f"keys.{key}.ttl": ttl}} + ) + + def delete_key(self, store_name, key): + """Deletes a key from the store.""" + self._collection.update_one( + {"store_name": store_name}, {"$unset": {f"keys.{key}": ""}} + ) + + def delete_store(self, store_name): + """Deletes the entire store.""" + self._collection.delete_one({"store_name": store_name}) + + +class MongoExecutionStoreRepo(ExecutionStoreRepo): + """MongoDB implementation of execution store repository.""" + + COLLECTION_NAME = "execution_store" + + def __init__(self, collection: Collection): + super().__init__(collection) diff --git a/fiftyone/operators/store/__init__.py b/fiftyone/operators/store/__init__.py new file mode 100644 index 0000000000..34d7008577 --- /dev/null +++ b/fiftyone/operators/store/__init__.py @@ -0,0 +1,18 @@ +""" +FiftyOne execution store module. + +| Copyright 2017-2024, Voxel51, Inc. +| `voxel51.com `_ +| +""" + +from .store import ExecutionStoreService +from .models import StoreDocument, KeyDocument +from .permissions import StorePermissions + +__all__ = [ + "ExecutionStoreService", + "StoreDocument", + "KeyDocument", + "StorePermissions", +] diff --git a/fiftyone/operators/store/models.py b/fiftyone/operators/store/models.py new file mode 100644 index 0000000000..31a9ca0e87 --- /dev/null +++ b/fiftyone/operators/store/models.py @@ -0,0 +1,64 @@ +""" +Store and key models for the execution store. +""" + +from pydantic import BaseModel, Field +from typing import Optional, Dict, Any +import datetime + + +class KeyDocument(BaseModel): + """Model representing a key in the store.""" + + key: str + value: Any + ttl: Optional[int] = None # Time To Live in milliseconds + created_at: datetime.datetime = Field( + default_factory=datetime.datetime.utcnow + ) + updated_at: Optional[datetime.datetime] = None + + class Config: + schema_extra = { + "example": { + "key": "widgets_key", + "value": {"widget_1": "foo", "widget_2": "bar"}, + "ttl": 600000, # 10 minutes + } + } + + +class StoreDocument(BaseModel): + """Model representing a store in the execution store.""" + + store_name: str + keys: Dict[str, KeyDocument] = Field(default_factory=dict) + permissions: Optional[ + Dict[str, Any] + ] = None # Permissions can be a role/user/plugin map + created_at: datetime.datetime = Field( + default_factory=datetime.datetime.utcnow + ) + + class Config: + schema_extra = { + "example": { + "store_name": "widget_store", + "keys": { + "widgets_key": { + "key": "widgets_key", + "value": {"widget_1": "foo", "widget_2": "bar"}, + "ttl": 600000, + } + }, + "permissions": { + "roles": {"admin": ["read", "write"]}, + "users": {"user_1": ["read", "write"], "user_2": ["read"]}, + "groups": {"group_1": ["read", "write"]}, + "plugins": { + "plugin_1_uri": ["read"], + "plugin_2_uri": ["write"], + }, + }, + } + } diff --git a/fiftyone/operators/store/permissions.py b/fiftyone/operators/store/permissions.py new file mode 100644 index 0000000000..4601ec7bcb --- /dev/null +++ b/fiftyone/operators/store/permissions.py @@ -0,0 +1,47 @@ +""" +Permission models for execution store. +""" + +from pydantic import BaseModel +from typing import List, Optional, Dict + +from typing import Dict, List, Optional +from pydantic import BaseModel, Field + + +class StorePermissions(BaseModel): + """Model representing permissions for a store.""" + + roles: Optional[Dict[str, List[str]]] = Field(default_factory=dict) + users: Optional[Dict[str, List[str]]] = Field(default_factory=dict) + plugins: Optional[Dict[str, List[str]]] = Field(default_factory=dict) + + @staticmethod + def default(): + """Provides default permissions, allowing full access to the creator.""" + return StorePermissions( + roles={"admin": ["read", "write"]}, + users={}, + plugins={}, + ) + + def has_permission(self, entity, action): + """Checks if the given entity (role, user, or plugin) has the specified action permission. + + Args: + entity: The entity to check (can be a role, user, or plugin) + action: The action to check (e.g., 'read', 'write') + + Returns: + bool: True if the entity has permission, False otherwise + """ + return True # permission implementation TBD + + # class Config: + # schema_extra = { + # "example": { + # "roles": {"admin": ["read", "write"]}, + # "users": {"user_1": ["read", "write"], "user_2": ["read"]}, + # "plugins": {"plugin_1_uri": ["read"], "plugin_2_uri": ["write"]} + # } + # } diff --git a/fiftyone/operators/store/store.py b/fiftyone/operators/store/store.py new file mode 100644 index 0000000000..d8ae577975 --- /dev/null +++ b/fiftyone/operators/store/store.py @@ -0,0 +1,131 @@ +""" +FiftyOne execution store service. + +| Copyright 2017-2024, Voxel51, Inc. +| `voxel51.com `_ +| +""" + +import logging +import traceback +from fiftyone.factory.repo_factory import RepositoryFactory +from fiftyone.operators.store.models import StoreDocument, KeyDocument +from fiftyone.operators.store.permissions import StorePermissions + +logger = logging.getLogger(__name__) + + +class ExecutionStoreService(object): + """Service for managing execution store operations.""" + + def __init__(self, repo=None): + if repo is None: + repo = RepositoryFactory.execution_store_repo() + + self._repo = repo + + def create_store(self, store_name, permissions=None): + """Creates a new store with the specified name and permissions. + + Args: + store_name: the name of the store + permissions (None): an optional permissions dict + + Returns: + a :class:`fiftyone.store.models.StoreDocument` + """ + return self._repo.create_store( + store_name=store_name, + permissions=permissions or StorePermissions.default(), + ) + + def set_key(self, store_name, key, value, ttl=None): + """Sets the value of a key in the specified store. + + Args: + store_name: the name of the store + key: the key to set + value: the value to set + ttl (None): an optional TTL in milliseconds + + Returns: + a :class:`fiftyone.store.models.KeyDocument` + """ + return self._repo.set_key( + store_name=store_name, key=key, value=value, ttl=ttl + ) + + def get_key(self, store_name, key): + """Retrieves the value of a key from the specified store. + + Args: + store_name: the name of the store + key: the key to retrieve + + Returns: + a :class:`fiftyone.store.models.KeyDocument` + """ + return self._repo.get_key(store_name=store_name, key=key) + + def delete_key(self, store_name, key): + """Deletes the specified key from the store. + + Args: + store_name: the name of the store + key: the key to delete + + Returns: + a :class:`fiftyone.store.models.KeyDocument` + """ + return self._repo.delete_key(store_name=store_name, key=key) + + def update_ttl(self, store_name, key, new_ttl): + """Updates the TTL of the specified key in the store. + + Args: + store_name: the name of the store + key: the key to update the TTL for + new_ttl: the new TTL in milliseconds + + Returns: + a :class:`fiftyone.store.models.KeyDocument` + """ + return self._repo.update_ttl( + store_name=store_name, key=key, ttl=new_ttl + ) + + def set_permissions(self, store_name, permissions): + """Sets the permissions for the specified store. + + Args: + store_name: the name of the store + permissions: a permissions object + + Returns: + a :class:`fiftyone.store.models.StoreDocument` + """ + return self._repo.update_permissions( + store_name=store_name, permissions=permissions + ) + + def list_stores(self, search=None, **kwargs): + """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) + + def delete_store(self, store_name): + """Deletes the specified store. + + Args: + store_name: the name of the store + + Returns: + a :class:`fiftyone.store.models.StoreDocument` + """ + return self._repo.delete_store(store_name=store_name) diff --git a/tests/unittests/execution_store_tests.py b/tests/unittests/execution_store_tests.py new file mode 100644 index 0000000000..c324a1e497 --- /dev/null +++ b/tests/unittests/execution_store_tests.py @@ -0,0 +1,170 @@ +""" +FiftyOne execution store related unit tests. + +| Copyright 2017-2024, Voxel51, Inc. +| `voxel51.com `_ +| +""" + +import time +import unittest +from unittest import mock +from unittest.mock import patch + +import fiftyone +from bson import ObjectId + +from fiftyone.operators.store import ExecutionStoreService +from fiftyone.operators.store.models import StoreDocument, KeyDocument +from fiftyone.operators.store.permissions import StorePermissions + + +class MockStoreRepo: + def __init__(self): + self.stores = {} + + def create_store(self, store_name, permissions=None): + if isinstance(permissions, StorePermissions): + permissions = permissions.dict() + + store = StoreDocument(store_name=store_name, permissions=permissions) + self.stores[store_name] = store + return store + + def set_key(self, store_name, key, value, ttl=None): + store = self.stores.get(store_name) + if store: + key_doc = KeyDocument(key=key, value=value, ttl=ttl) + store.keys[key] = key_doc + return key_doc + return None + + def get_key(self, store_name, key): + store = self.stores.get(store_name) + if store and key in store.keys: + return store.keys[key] + return None + + def delete_key(self, store_name, key): + store = self.stores.get(store_name) + if store and key in store.keys: + del store.keys[key] + + def delete_store(self, store_name): + if store_name in self.stores: + del self.stores[store_name] + + def update_ttl(self, store_name, key, ttl): + store = self.stores.get(store_name) + if store and key in store.keys: + store.keys[key].ttl = ttl + return store.keys[key] + return None + + +class ExecutionStoreServiceTests(unittest.TestCase): + def setUp(self): + # Mock repository and service for testing + self.repo = MockStoreRepo() + self.svc = ExecutionStoreService(repo=self.repo) + + def test_create_store(self): + store_name = "test_store" + permissions = StorePermissions.default() + + store = self.svc.create_store( + store_name=store_name, permissions=permissions + ) + + self.assertIsNotNone(store) + self.assertEqual(store.store_name, store_name) + # Now compare the dictionary instead of the attribute + self.assertEqual(store.permissions["roles"], permissions.roles) + + def test_set_and_get_key(self): + store_name = "test_store" + self.svc.create_store(store_name=store_name) + + key = "test_key" + value = {"foo": "bar"} + ttl = 10000 + + key_doc = self.svc.set_key( + store_name=store_name, key=key, value=value, ttl=ttl + ) + self.assertIsNotNone(key_doc) + self.assertEqual(key_doc.key, key) + self.assertEqual(key_doc.value, value) + self.assertEqual(key_doc.ttl, ttl) + + # Get the key and validate the value + fetched_key_doc = self.svc.get_key(store_name=store_name, key=key) + self.assertIsNotNone(fetched_key_doc) + self.assertEqual(fetched_key_doc.key, key) + self.assertEqual(fetched_key_doc.value, value) + + def test_delete_key(self): + store_name = "test_store" + self.svc.create_store(store_name=store_name) + + key = "test_key" + value = {"foo": "bar"} + self.svc.set_key(store_name=store_name, key=key, value=value) + + # Ensure the key exists + key_doc = self.svc.get_key(store_name=store_name, key=key) + self.assertIsNotNone(key_doc) + + # Delete the key + self.svc.delete_key(store_name=store_name, key=key) + + # Ensure the key is deleted + key_doc = self.svc.get_key(store_name=store_name, key=key) + self.assertIsNone(key_doc) + + def test_update_ttl(self): + store_name = "test_store" + self.svc.create_store(store_name=store_name) + + key = "test_key" + value = {"foo": "bar"} + self.svc.set_key(store_name=store_name, key=key, value=value) + + # Update TTL + new_ttl = 5000 + self.svc.update_ttl(store_name=store_name, key=key, new_ttl=new_ttl) + + # Fetch key and check updated TTL + key_doc = self.svc.get_key(store_name=store_name, key=key) + self.assertIsNotNone(key_doc) + self.assertEqual(key_doc.ttl, new_ttl) + + def test_delete_store(self): + store_name = "test_store" + self.svc.create_store(store_name=store_name) + + # Ensure the store exists + store = self.repo.stores.get(store_name) + self.assertIsNotNone(store) + + # Delete the store + self.svc.delete_store(store_name=store_name) + + # Ensure the store is deleted + store = self.repo.stores.get(store_name) + self.assertIsNone(store) + + def test_store_permissions(self): + store_name = "test_store" + permissions = StorePermissions.default() + self.svc.create_store(store_name=store_name, permissions=permissions) + + # Check default permissions + store = self.repo.stores.get(store_name) + self.assertIsNotNone(store) + # Now compare the dictionary instead of the attribute + self.assertEqual(store.permissions["roles"], permissions.roles) + + +if __name__ == "__main__": + unittest.main() From eb6728a63d93ea358ea551f3d936477e8aa35996 Mon Sep 17 00:00:00 2001 From: Ritchie Martori Date: Fri, 20 Sep 2024 09:54:59 -0700 Subject: [PATCH 2/9] initial int tests for execution store --- fiftyone/factory/repos/execution_store.py | 51 +++++---- fiftyone/operators/executor.py | 14 +++ fiftyone/operators/store/__init__.py | 4 +- fiftyone/operators/store/models.py | 25 ++-- fiftyone/operators/store/service.py | 130 +++++++++++++++++++++ fiftyone/operators/store/store.py | 133 ++++++++-------------- tests/unittests/execution_store_tests.py | 129 ++++++++++++++++++++- 7 files changed, 361 insertions(+), 125 deletions(-) create mode 100644 fiftyone/operators/store/service.py diff --git a/fiftyone/factory/repos/execution_store.py b/fiftyone/factory/repos/execution_store.py index ab56ccc620..abe86ee956 100644 --- a/fiftyone/factory/repos/execution_store.py +++ b/fiftyone/factory/repos/execution_store.py @@ -6,6 +6,13 @@ from fiftyone.operators.store.models import StoreDocument, KeyDocument +def _where(store_name, key=None): + query = {"store_name": store_name} + if key is not None: + query["key"] = key + return query + + class ExecutionStoreRepo: """Base class for execution store repositories.""" @@ -14,7 +21,7 @@ class ExecutionStoreRepo: def __init__(self, collection: Collection): self._collection = collection - def create_store(self, store_name, permissions=None): + def create_store(self, store_name, permissions=None) -> StoreDocument: """Creates a store in the execution store.""" store_doc = StoreDocument( store_name=store_name, permissions=permissions @@ -22,37 +29,39 @@ def create_store(self, store_name, permissions=None): self._collection.insert_one(store_doc.dict()) return store_doc - def set_key(self, store_name, key, value, ttl=None): - """Sets a key in the specified store.""" - key_doc = KeyDocument(key=key, value=value, ttl=ttl) + def set_key(self, store_name, key, value, ttl=None) -> KeyDocument: + """Sets or updates a key in the specified store.""" + key_doc = KeyDocument( + store_name=store_name, key=key, value=value, ttl=ttl + ) + # Update or insert the key self._collection.update_one( - {"store_name": store_name}, - {"$set": {f"keys.{key}": key_doc.dict()}}, + _where(store_name, key), {"$set": key_doc.dict()}, upsert=True ) return key_doc - def get_key(self, store_name, key): + def get_key(self, store_name, key) -> KeyDocument: """Gets a key from the specified store.""" - store = self._collection.find_one({"store_name": store_name}) - if store and key in store.get("keys", {}): - return KeyDocument(**store["keys"][key]) - return None + raw_key_doc = self._collection.find_one(_where(store_name, key)) + key_doc = KeyDocument(**raw_key_doc) if raw_key_doc else None + return key_doc - def update_ttl(self, store_name, key, ttl): + def update_ttl(self, store_name, key, ttl) -> bool: """Updates the TTL for a key.""" - self._collection.update_one( - {"store_name": store_name}, {"$set": {f"keys.{key}.ttl": ttl}} + result = self._collection.update_one( + _where(store_name, key), {"$set": {"ttl": ttl}} ) + return result.modified_count > 0 - def delete_key(self, store_name, key): - """Deletes a key from the store.""" - self._collection.update_one( - {"store_name": store_name}, {"$unset": {f"keys.{key}": ""}} - ) + def delete_key(self, store_name, key) -> bool: + """Deletes the document that matches the store name and key.""" + result = self._collection.delete_one(_where(store_name, key)) + return result.deleted_count > 0 - def delete_store(self, store_name): + def delete_store(self, store_name) -> int: """Deletes the entire store.""" - self._collection.delete_one({"store_name": store_name}) + result = self._collection.delete_many(_where(store_name)) + return result.deleted_count class MongoExecutionStoreRepo(ExecutionStoreRepo): diff --git a/fiftyone/operators/executor.py b/fiftyone/operators/executor.py index 6a4b596c99..c0751e46f4 100644 --- a/fiftyone/operators/executor.py +++ b/fiftyone/operators/executor.py @@ -830,6 +830,20 @@ def set_progress(self, progress=None, label=None): else: self.log(f"Progress: {progress} - {label}") + # TODO resolve circular import so this can have a type + def create_store(self, store_name): + """Creates a new store with the specified name. + + Args: + store_name: the name of the store + + Returns: + a :class:`fiftyone.operators.store.ExecutionStore` + """ + from fiftyone.operators.store import ExecutionStore + + return ExecutionStore.create(store_name) + def serialize(self): """Serializes the execution context. diff --git a/fiftyone/operators/store/__init__.py b/fiftyone/operators/store/__init__.py index 34d7008577..0c33732c4f 100644 --- a/fiftyone/operators/store/__init__.py +++ b/fiftyone/operators/store/__init__.py @@ -6,7 +6,8 @@ | """ -from .store import ExecutionStoreService +from .service import ExecutionStoreService +from .store import ExecutionStore from .models import StoreDocument, KeyDocument from .permissions import StorePermissions @@ -15,4 +16,5 @@ "StoreDocument", "KeyDocument", "StorePermissions", + "ExecutionStore", ] diff --git a/fiftyone/operators/store/models.py b/fiftyone/operators/store/models.py index 31a9ca0e87..1ba3005abf 100644 --- a/fiftyone/operators/store/models.py +++ b/fiftyone/operators/store/models.py @@ -10,6 +10,7 @@ class KeyDocument(BaseModel): """Model representing a key in the store.""" + store_name: str key: str value: Any ttl: Optional[int] = None # Time To Live in milliseconds @@ -17,10 +18,14 @@ class KeyDocument(BaseModel): default_factory=datetime.datetime.utcnow ) updated_at: Optional[datetime.datetime] = None + permissions: Optional[ + Dict[str, Any] + ] = None # Permissions can be a role/user/plugin map class Config: schema_extra = { "example": { + "store_name": "widget_store", "key": "widgets_key", "value": {"widget_1": "foo", "widget_2": "bar"}, "ttl": 600000, # 10 minutes @@ -28,29 +33,17 @@ class Config: } -class StoreDocument(BaseModel): +class StoreDocument(KeyDocument): """Model representing a store in the execution store.""" - store_name: str - keys: Dict[str, KeyDocument] = Field(default_factory=dict) - permissions: Optional[ - Dict[str, Any] - ] = None # Permissions can be a role/user/plugin map - created_at: datetime.datetime = Field( - default_factory=datetime.datetime.utcnow - ) + key: str = "__store__" + value: Optional[Dict[str, Any]] = None class Config: schema_extra = { "example": { + "key": "__store__", "store_name": "widget_store", - "keys": { - "widgets_key": { - "key": "widgets_key", - "value": {"widget_1": "foo", "widget_2": "bar"}, - "ttl": 600000, - } - }, "permissions": { "roles": {"admin": ["read", "write"]}, "users": {"user_1": ["read", "write"], "user_2": ["read"]}, diff --git a/fiftyone/operators/store/service.py b/fiftyone/operators/store/service.py new file mode 100644 index 0000000000..9da28053a2 --- /dev/null +++ b/fiftyone/operators/store/service.py @@ -0,0 +1,130 @@ +""" +FiftyOne execution store service. + +| Copyright 2017-2024, Voxel51, Inc. +| `voxel51.com `_ +| +""" + +import logging +import traceback +from fiftyone.factory.repo_factory import RepositoryFactory +from fiftyone.operators.store.permissions import StorePermissions + +logger = logging.getLogger(__name__) + + +class ExecutionStoreService(object): + """Service for managing execution store operations.""" + + def __init__(self, repo=None): + if repo is None: + repo = RepositoryFactory.execution_store_repo() + + self._repo = repo + + def create_store(self, store_name, permissions=None): + """Creates a new store with the specified name and permissions. + + Args: + store_name: the name of the store + permissions (None): an optional permissions dict + + Returns: + a :class:`fiftyone.store.models.StoreDocument` + """ + return self._repo.create_store( + store_name=store_name, + permissions=permissions or StorePermissions.default(), + ) + + def set_key(self, store_name, key, value, ttl=None): + """Sets the value of a key in the specified store. + + Args: + store_name: the name of the store + key: the key to set + value: the value to set + ttl (None): an optional TTL in milliseconds + + Returns: + a :class:`fiftyone.store.models.KeyDocument` + """ + return self._repo.set_key( + store_name=store_name, key=key, value=value, ttl=ttl + ) + + def get_key(self, store_name, key): + """Retrieves the value of a key from the specified store. + + Args: + store_name: the name of the store + key: the key to retrieve + + Returns: + a :class:`fiftyone.store.models.KeyDocument` + """ + return self._repo.get_key(store_name=store_name, key=key) + + def delete_key(self, store_name, key): + """Deletes the specified key from the store. + + Args: + store_name: the name of the store + key: the key to delete + + Returns: + a :class:`fiftyone.store.models.KeyDocument` + """ + return self._repo.delete_key(store_name=store_name, key=key) + + def update_ttl(self, store_name, key, new_ttl): + """Updates the TTL of the specified key in the store. + + Args: + store_name: the name of the store + key: the key to update the TTL for + new_ttl: the new TTL in milliseconds + + Returns: + a :class:`fiftyone.store.models.KeyDocument` + """ + return self._repo.update_ttl( + store_name=store_name, key=key, ttl=new_ttl + ) + + def set_permissions(self, store_name, permissions): + """Sets the permissions for the specified store. + + Args: + store_name: the name of the store + permissions: a permissions object + + Returns: + a :class:`fiftyone.store.models.StoreDocument` + """ + return self._repo.update_permissions( + store_name=store_name, permissions=permissions + ) + + def list_stores(self, search=None, **kwargs): + """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) + + def delete_store(self, store_name): + """Deletes the specified store. + + Args: + store_name: the name of the store + + Returns: + a :class:`fiftyone.store.models.StoreDocument` + """ + return self._repo.delete_store(store_name=store_name) diff --git a/fiftyone/operators/store/store.py b/fiftyone/operators/store/store.py index d8ae577975..bb5077fe8c 100644 --- a/fiftyone/operators/store/store.py +++ b/fiftyone/operators/store/store.py @@ -1,5 +1,5 @@ """ -FiftyOne execution store service. +FiftyOne execution store. | Copyright 2017-2024, Voxel51, Inc. | `voxel51.com `_ @@ -7,125 +7,90 @@ """ import logging -import traceback from fiftyone.factory.repo_factory import RepositoryFactory -from fiftyone.operators.store.models import StoreDocument, KeyDocument -from fiftyone.operators.store.permissions import StorePermissions +from fiftyone.operators.store.service import ExecutionStoreService +from typing import Any, Optional logger = logging.getLogger(__name__) -class ExecutionStoreService(object): - """Service for managing execution store operations.""" +class ExecutionStore: + @staticmethod + def create(store_name: str) -> "ExecutionStore": + return ExecutionStore(store_name, ExecutionStoreService()) - def __init__(self, repo=None): - if repo is None: - repo = RepositoryFactory.execution_store_repo() - - self._repo = repo - - def create_store(self, store_name, permissions=None): - """Creates a new store with the specified name and permissions. - - Args: - store_name: the name of the store - permissions (None): an optional permissions dict - - Returns: - a :class:`fiftyone.store.models.StoreDocument` + def __init__(self, store_name: str, store_service: ExecutionStoreService): """ - return self._repo.create_store( - store_name=store_name, - permissions=permissions or StorePermissions.default(), - ) - - def set_key(self, store_name, key, value, ttl=None): - """Sets the value of a key in the specified store. - Args: - store_name: the name of the store - key: the key to set - value: the value to set - ttl (None): an optional TTL in milliseconds - - Returns: - a :class:`fiftyone.store.models.KeyDocument` + store_name (str): The name of the store. + store_service (ExecutionStoreService): The store service instance. """ - return self._repo.set_key( - store_name=store_name, key=key, value=value, ttl=ttl - ) + self.store_name: str = store_name + self._store_service: ExecutionStoreService = store_service - def get_key(self, store_name, key): - """Retrieves the value of a key from the specified store. + def get(self, key: str) -> Optional[Any]: + """Retrieves a value from the store by its key. Args: - store_name: the name of the store - key: the key to retrieve + key (str): The key to retrieve the value for. Returns: - a :class:`fiftyone.store.models.KeyDocument` + Optional[Any]: The value stored under the given key, or None if not found. """ - return self._repo.get_key(store_name=store_name, key=key) + key_doc = self._store_service.get_key(self.store_name, key) + if key_doc is None: + return None + return key_doc.value - def delete_key(self, store_name, key): - """Deletes the specified key from the store. + def set(self, key: str, value: Any, ttl: Optional[int] = None) -> None: + """Sets a value in the store with an optional TTL. Args: - store_name: the name of the store - key: the key to delete - - Returns: - a :class:`fiftyone.store.models.KeyDocument` + key (str): The key to store the value under. + value (Any): The value to store. + ttl (Optional[int], optional): The time-to-live in milliseconds. Defaults to None. """ - return self._repo.delete_key(store_name=store_name, key=key) + self._store_service.set_key(self.store_name, key, value, ttl) - def update_ttl(self, store_name, key, new_ttl): - """Updates the TTL of the specified key in the store. + def delete(self, key: str) -> None: + """Deletes a key from the store. Args: - store_name: the name of the store - key: the key to update the TTL for - new_ttl: the new TTL in milliseconds - - Returns: - a :class:`fiftyone.store.models.KeyDocument` + key (str): The key to delete. """ - return self._repo.update_ttl( - store_name=store_name, key=key, ttl=new_ttl - ) + self._store_service.delete_key(self.store_name, key) - def set_permissions(self, store_name, permissions): - """Sets the permissions for the specified store. + def has(self, key: str) -> bool: + """Checks if the store has a specific key. Args: - store_name: the name of the store - permissions: a permissions object + key (str): The key to check. Returns: - a :class:`fiftyone.store.models.StoreDocument` + bool: True if the key exists, False otherwise. """ - return self._repo.update_permissions( - store_name=store_name, permissions=permissions - ) + return self._store_service.has_key(self.store_name, key) - def list_stores(self, search=None, **kwargs): - """Lists all stores matching the given criteria. + def clear(self) -> None: + """Clears all the data in the store.""" + self._store_service.clear_store(self.store_name) - Args: - search (None): optional search term dict + def update_ttl(self, key: str, new_ttl: int) -> None: + """Updates the TTL for a specific key. - Returns: - a list of :class:`fiftyone.store.models.StoreDocument` + Args: + key (str): The key to update the TTL for. + new_ttl (int): The new TTL in milliseconds. """ - return self._repo.list_stores(search=search, **kwargs) + self._store_service.update_ttl(self.store_name, key, new_ttl) - def delete_store(self, store_name): - """Deletes the specified store. + def get_ttl(self, key: str) -> Optional[int]: + """Retrieves the TTL for a specific key. Args: - store_name: the name of the store + key (str): The key to get the TTL for. Returns: - a :class:`fiftyone.store.models.StoreDocument` + Optional[int]: The TTL in milliseconds, or None if the key does not have a TTL. """ - return self._repo.delete_store(store_name=store_name) + return self._store_service.get_ttl(self.store_name, key) diff --git a/tests/unittests/execution_store_tests.py b/tests/unittests/execution_store_tests.py index c324a1e497..1ad21aa972 100644 --- a/tests/unittests/execution_store_tests.py +++ b/tests/unittests/execution_store_tests.py @@ -9,14 +9,37 @@ import time import unittest from unittest import mock -from unittest.mock import patch +from unittest.mock import patch, MagicMock import fiftyone from bson import ObjectId +import fiftyone.operators as foo from fiftyone.operators.store import ExecutionStoreService from fiftyone.operators.store.models import StoreDocument, KeyDocument from fiftyone.operators.store.permissions import StorePermissions +from fiftyone.factory.repo_factory import ExecutionStoreRepo +from fiftyone.operators.operator import Operator +from fiftyone.operators.store import ExecutionStoreService + + +class MockOperator: + """Mock operator that interacts with ExecutionStore.""" + + def __init__(self, ctx: foo.ExecutionContext): + self.ctx = ctx + + def execute(self): + # Example logic of interacting with the store + store = self.ctx.create_store("widgets") + + # Set a value in the store + store.set("widget_1", {"name": "Widget One", "value": 100}, ttl=60000) + + # Get the value back from the store + result = store.get("widget_1") + + return result class MockStoreRepo: @@ -166,5 +189,105 @@ def test_store_permissions(self): self.assertEqual(store.permissions["roles"], permissions.roles) -if __name__ == "__main__": - unittest.main() +class TestOperatorWithExecutionStore(Operator): + def execute(self, ctx): + # Create a store and interact with it + store = ctx.create_store("test_store") + + # Perform some store operations + store.set("my_key", {"foo": "bar"}) + value = store.get("my_key") + store.delete("my_key") + + return value + + +class ExecutionStoreIntegrationTests(unittest.TestCase): + def setUp(self): + # Create a MagicMock for the MongoDB collection + self.mock_collection = MagicMock() + + # Instantiate the store repo and replace the _collection attribute + self.store_repo = ExecutionStoreRepo(self.mock_collection) + + # Create the store service with the mocked repo + self.store_service = ExecutionStoreService(self.store_repo) + + # Create an execution context + from fiftyone.operators import ExecutionContext + + self.ctx = ExecutionContext() + self.ctx.store_service = ( + self.store_service + ) # Inject the store service into the context + + # Create an instance of the operator + self.operator = TestOperatorWithExecutionStore() + + def test_operator_execute_with_store(self): + # Mock MongoDB collection methods + self.mock_collection.update_one.return_value = None + self.mock_collection.find_one.return_value = { + "key": "my_key", + "value": {"foo": "bar"}, + } + self.mock_collection.delete_one.return_value = None + + # Call the operator's execute method + result = self.operator.execute(self.ctx) + + # Verify that the store interactions were made correctly + self.mock_collection.update_one.assert_called_once() # Checking that set operation inserts data + self.mock_collection.find_one.assert_called_once_with( + {"store_name": "test_store", "key": "my_key"} + ) # Checking that get operation retrieves data + self.mock_collection.delete_one.assert_called_once_with( + {"store_name": "test_store", "key": "my_key"} + ) # Checking that delete operation removes data + + # Verify the correct value was returned from the store + self.assertEqual(result, {"key": "my_key", "value": {"foo": "bar"}}) + + def test_operator_execute_set_key(self): + # Mock the insert_one call for the set operation + self.mock_collection.insert_one.return_value = None + + # Call the operator's execute method + self.operator.execute(self.ctx) + + # Check that insert_one (set) was called with the correct arguments + self.mock_collection.insert_one.assert_called_once_with( + { + "store_name": "test_store", + "key": "my_key", + "value": {"foo": "bar"}, + } + ) + + def test_operator_execute_get_key(self): + # Mock the find_one call for the get operation + self.mock_collection.find_one.return_value = { + "key": "my_key", + "value": {"foo": "bar"}, + } + + # Call the operator's execute method + result = self.operator.execute(self.ctx) + + # Check that find_one (get) was called correctly and returned the expected value + self.mock_collection.find_one.assert_called_once_with( + {"store_name": "test_store", "key": "my_key"} + ) + self.assertEqual(result, {"key": "my_key", "value": {"foo": "bar"}}) + + def test_operator_execute_delete_key(self): + # Mock the delete_one call for the delete operation + self.mock_collection.delete_one.return_value = None + + # Call the operator's execute method + self.operator.execute(self.ctx) + + # Check that delete_one was called with the correct arguments + self.mock_collection.delete_one.assert_called_once_with( + {"store_name": "test_store", "key": "my_key"} + ) From 19857d3e018e8eb53af4a0c51e8c3b74e1484ecc Mon Sep 17 00:00:00 2001 From: Ritchie Martori Date: Sun, 22 Sep 2024 13:48:25 -0700 Subject: [PATCH 3/9] fixes for execution store --- fiftyone/factory/repo_factory.py | 1 + fiftyone/factory/repos/execution_store.py | 24 +++++++++++++++++++++-- fiftyone/operators/store/models.py | 18 ++++++++--------- fiftyone/operators/store/service.py | 8 ++++++++ fiftyone/operators/store/store.py | 8 ++++++++ 5 files changed, 47 insertions(+), 12 deletions(-) diff --git a/fiftyone/factory/repo_factory.py b/fiftyone/factory/repo_factory.py index 314163544e..05a875fce1 100644 --- a/fiftyone/factory/repo_factory.py +++ b/fiftyone/factory/repo_factory.py @@ -52,6 +52,7 @@ def delegated_operation_repo() -> DelegatedOperationRepo: @staticmethod def execution_store_repo() -> ExecutionStoreRepo: """Factory method for execution store repository.""" + print("Factory method for execution store repository.") if ( MongoExecutionStoreRepo.COLLECTION_NAME not in RepositoryFactory.repos diff --git a/fiftyone/factory/repos/execution_store.py b/fiftyone/factory/repos/execution_store.py index abe86ee956..0b2e5aad9c 100644 --- a/fiftyone/factory/repos/execution_store.py +++ b/fiftyone/factory/repos/execution_store.py @@ -31,8 +31,12 @@ def create_store(self, store_name, permissions=None) -> StoreDocument: 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) key_doc = KeyDocument( - store_name=store_name, key=key, value=value, ttl=ttl + store_name=store_name, + key=key, + value=value, + expires_at=expiration if ttl else None, ) # Update or insert the key self._collection.update_one( @@ -46,10 +50,17 @@ def get_key(self, store_name, key) -> KeyDocument: key_doc = KeyDocument(**raw_key_doc) if raw_key_doc else None return key_doc + def list_keys(self, store_name) -> list[str]: + """Lists all keys in the specified store.""" + keys = self._collection.find(_where(store_name)) + # TODO: redact non-key fields + return [key["key"] for key in keys] + def update_ttl(self, store_name, key, ttl) -> bool: """Updates the TTL for a key.""" + expiration = KeyDocument.get_expiration(ttl) result = self._collection.update_one( - _where(store_name, key), {"$set": {"ttl": ttl}} + _where(store_name, key), {"$set": {"expires_at": expiration}} ) return result.modified_count > 0 @@ -71,3 +82,12 @@ class MongoExecutionStoreRepo(ExecutionStoreRepo): def __init__(self, collection: Collection): super().__init__(collection) + self._create_indexes() + + def _create_indexes(self): + indices = self._collection.list_indexes() + expires_at_name = "expires_at" + if expires_at_name not in indices: + self._collection.create_index( + expires_at_name, name=expires_at_name, expireAfterSeconds=0 + ) diff --git a/fiftyone/operators/store/models.py b/fiftyone/operators/store/models.py index 1ba3005abf..e2502a98ba 100644 --- a/fiftyone/operators/store/models.py +++ b/fiftyone/operators/store/models.py @@ -13,24 +13,22 @@ class KeyDocument(BaseModel): store_name: str key: str value: Any - ttl: Optional[int] = None # Time To Live in milliseconds created_at: datetime.datetime = Field( default_factory=datetime.datetime.utcnow ) updated_at: Optional[datetime.datetime] = None + expires_at: Optional[datetime.datetime] = None permissions: Optional[ Dict[str, Any] ] = None # Permissions can be a role/user/plugin map - class Config: - schema_extra = { - "example": { - "store_name": "widget_store", - "key": "widgets_key", - "value": {"widget_1": "foo", "widget_2": "bar"}, - "ttl": 600000, # 10 minutes - } - } + @staticmethod + def get_expiration(ttl: Optional[int]) -> Optional[datetime.datetime]: + """Gets the expiration date for a key with the given TTL.""" + if ttl is None: + return None + + return datetime.datetime.now() + datetime.timedelta(milliseconds=ttl) class StoreDocument(KeyDocument): diff --git a/fiftyone/operators/store/service.py b/fiftyone/operators/store/service.py index 9da28053a2..0cebd197f7 100644 --- a/fiftyone/operators/store/service.py +++ b/fiftyone/operators/store/service.py @@ -128,3 +128,11 @@ def delete_store(self, store_name): a :class:`fiftyone.store.models.StoreDocument` """ return self._repo.delete_store(store_name=store_name) + + def list_keys(self, store_name): + """Lists all keys in the specified store. + + Args: + store_name: the name of the store + """ + return self._repo.list_keys(store_name) diff --git a/fiftyone/operators/store/store.py b/fiftyone/operators/store/store.py index bb5077fe8c..d0fc738f9f 100644 --- a/fiftyone/operators/store/store.py +++ b/fiftyone/operators/store/store.py @@ -94,3 +94,11 @@ def get_ttl(self, key: str) -> Optional[int]: Optional[int]: The TTL in milliseconds, 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. + """ + return self._store_service.list_keys(self.store_name) From affb3ecc1c157575d2d5d662df25cb35fc736415 Mon Sep 17 00:00:00 2001 From: Ritchie Martori Date: Sun, 22 Sep 2024 14:34:53 -0700 Subject: [PATCH 4/9] exec store cleanup --- fiftyone/factory/repo_factory.py | 1 - fiftyone/operators/store/__init__.py | 2 - fiftyone/operators/store/models.py | 17 --------- fiftyone/operators/store/permissions.py | 47 ------------------------ fiftyone/operators/store/service.py | 26 ++----------- tests/unittests/execution_store_tests.py | 24 ++---------- 6 files changed, 7 insertions(+), 110 deletions(-) delete mode 100644 fiftyone/operators/store/permissions.py diff --git a/fiftyone/factory/repo_factory.py b/fiftyone/factory/repo_factory.py index 05a875fce1..314163544e 100644 --- a/fiftyone/factory/repo_factory.py +++ b/fiftyone/factory/repo_factory.py @@ -52,7 +52,6 @@ def delegated_operation_repo() -> DelegatedOperationRepo: @staticmethod def execution_store_repo() -> ExecutionStoreRepo: """Factory method for execution store repository.""" - print("Factory method for execution store repository.") if ( MongoExecutionStoreRepo.COLLECTION_NAME not in RepositoryFactory.repos diff --git a/fiftyone/operators/store/__init__.py b/fiftyone/operators/store/__init__.py index 0c33732c4f..50f9f3b400 100644 --- a/fiftyone/operators/store/__init__.py +++ b/fiftyone/operators/store/__init__.py @@ -9,12 +9,10 @@ from .service import ExecutionStoreService from .store import ExecutionStore from .models import StoreDocument, KeyDocument -from .permissions import StorePermissions __all__ = [ "ExecutionStoreService", "StoreDocument", "KeyDocument", - "StorePermissions", "ExecutionStore", ] diff --git a/fiftyone/operators/store/models.py b/fiftyone/operators/store/models.py index e2502a98ba..dab61012d4 100644 --- a/fiftyone/operators/store/models.py +++ b/fiftyone/operators/store/models.py @@ -36,20 +36,3 @@ class StoreDocument(KeyDocument): key: str = "__store__" value: Optional[Dict[str, Any]] = None - - class Config: - schema_extra = { - "example": { - "key": "__store__", - "store_name": "widget_store", - "permissions": { - "roles": {"admin": ["read", "write"]}, - "users": {"user_1": ["read", "write"], "user_2": ["read"]}, - "groups": {"group_1": ["read", "write"]}, - "plugins": { - "plugin_1_uri": ["read"], - "plugin_2_uri": ["write"], - }, - }, - } - } diff --git a/fiftyone/operators/store/permissions.py b/fiftyone/operators/store/permissions.py deleted file mode 100644 index 4601ec7bcb..0000000000 --- a/fiftyone/operators/store/permissions.py +++ /dev/null @@ -1,47 +0,0 @@ -""" -Permission models for execution store. -""" - -from pydantic import BaseModel -from typing import List, Optional, Dict - -from typing import Dict, List, Optional -from pydantic import BaseModel, Field - - -class StorePermissions(BaseModel): - """Model representing permissions for a store.""" - - roles: Optional[Dict[str, List[str]]] = Field(default_factory=dict) - users: Optional[Dict[str, List[str]]] = Field(default_factory=dict) - plugins: Optional[Dict[str, List[str]]] = Field(default_factory=dict) - - @staticmethod - def default(): - """Provides default permissions, allowing full access to the creator.""" - return StorePermissions( - roles={"admin": ["read", "write"]}, - users={}, - plugins={}, - ) - - def has_permission(self, entity, action): - """Checks if the given entity (role, user, or plugin) has the specified action permission. - - Args: - entity: The entity to check (can be a role, user, or plugin) - action: The action to check (e.g., 'read', 'write') - - Returns: - bool: True if the entity has permission, False otherwise - """ - return True # permission implementation TBD - - # class Config: - # schema_extra = { - # "example": { - # "roles": {"admin": ["read", "write"]}, - # "users": {"user_1": ["read", "write"], "user_2": ["read"]}, - # "plugins": {"plugin_1_uri": ["read"], "plugin_2_uri": ["write"]} - # } - # } diff --git a/fiftyone/operators/store/service.py b/fiftyone/operators/store/service.py index 0cebd197f7..47dcf1ada3 100644 --- a/fiftyone/operators/store/service.py +++ b/fiftyone/operators/store/service.py @@ -7,9 +7,7 @@ """ import logging -import traceback from fiftyone.factory.repo_factory import RepositoryFactory -from fiftyone.operators.store.permissions import StorePermissions logger = logging.getLogger(__name__) @@ -23,20 +21,16 @@ def __init__(self, repo=None): self._repo = repo - def create_store(self, store_name, permissions=None): - """Creates a new store with the specified name and permissions. + def create_store(self, store_name): + """Creates a new store with the specified name. Args: store_name: the name of the store - permissions (None): an optional permissions dict Returns: a :class:`fiftyone.store.models.StoreDocument` """ - return self._repo.create_store( - store_name=store_name, - permissions=permissions or StorePermissions.default(), - ) + return self._repo.create_store(store_name=store_name) def set_key(self, store_name, key, value, ttl=None): """Sets the value of a key in the specified store. @@ -93,20 +87,6 @@ def update_ttl(self, store_name, key, new_ttl): store_name=store_name, key=key, ttl=new_ttl ) - def set_permissions(self, store_name, permissions): - """Sets the permissions for the specified store. - - Args: - store_name: the name of the store - permissions: a permissions object - - Returns: - a :class:`fiftyone.store.models.StoreDocument` - """ - return self._repo.update_permissions( - store_name=store_name, permissions=permissions - ) - def list_stores(self, search=None, **kwargs): """Lists all stores matching the given criteria. diff --git a/tests/unittests/execution_store_tests.py b/tests/unittests/execution_store_tests.py index 1ad21aa972..12eb3a4c69 100644 --- a/tests/unittests/execution_store_tests.py +++ b/tests/unittests/execution_store_tests.py @@ -11,13 +11,9 @@ from unittest import mock from unittest.mock import patch, MagicMock -import fiftyone -from bson import ObjectId - import fiftyone.operators as foo from fiftyone.operators.store import ExecutionStoreService from fiftyone.operators.store.models import StoreDocument, KeyDocument -from fiftyone.operators.store.permissions import StorePermissions from fiftyone.factory.repo_factory import ExecutionStoreRepo from fiftyone.operators.operator import Operator from fiftyone.operators.store import ExecutionStoreService @@ -46,11 +42,8 @@ class MockStoreRepo: def __init__(self): self.stores = {} - def create_store(self, store_name, permissions=None): - if isinstance(permissions, StorePermissions): - permissions = permissions.dict() - - store = StoreDocument(store_name=store_name, permissions=permissions) + def create_store(self, store_name): + store = StoreDocument(store_name=store_name) self.stores[store_name] = store return store @@ -93,16 +86,10 @@ def setUp(self): def test_create_store(self): store_name = "test_store" - permissions = StorePermissions.default() - - store = self.svc.create_store( - store_name=store_name, permissions=permissions - ) + store = self.svc.create_store(store_name=store_name) self.assertIsNotNone(store) self.assertEqual(store.store_name, store_name) - # Now compare the dictionary instead of the attribute - self.assertEqual(store.permissions["roles"], permissions.roles) def test_set_and_get_key(self): store_name = "test_store" @@ -179,14 +166,11 @@ def test_delete_store(self): def test_store_permissions(self): store_name = "test_store" - permissions = StorePermissions.default() - self.svc.create_store(store_name=store_name, permissions=permissions) + self.svc.create_store(store_name=store_name) # Check default permissions store = self.repo.stores.get(store_name) self.assertIsNotNone(store) - # Now compare the dictionary instead of the attribute - self.assertEqual(store.permissions["roles"], permissions.roles) class TestOperatorWithExecutionStore(Operator): From ae74b8b8da9f56f91c5749f70bdd7cbdc48d7f21 Mon Sep 17 00:00:00 2001 From: Ritchie Martori Date: Sun, 22 Sep 2024 16:22:02 -0700 Subject: [PATCH 5/9] better exec store tests --- fiftyone/operators/store/models.py | 7 +- fiftyone/operators/store/service.py | 4 +- fiftyone/operators/store/store.py | 7 +- tests/unittests/execution_store_tests.py | 400 ++++++++++------------- 4 files changed, 187 insertions(+), 231 deletions(-) diff --git a/fiftyone/operators/store/models.py b/fiftyone/operators/store/models.py index dab61012d4..e05767fe22 100644 --- a/fiftyone/operators/store/models.py +++ b/fiftyone/operators/store/models.py @@ -14,13 +14,10 @@ class KeyDocument(BaseModel): key: str value: Any created_at: datetime.datetime = Field( - default_factory=datetime.datetime.utcnow + default_factory=datetime.datetime.now ) updated_at: Optional[datetime.datetime] = None expires_at: Optional[datetime.datetime] = None - permissions: Optional[ - Dict[str, Any] - ] = None # Permissions can be a role/user/plugin map @staticmethod def get_expiration(ttl: Optional[int]) -> Optional[datetime.datetime]: @@ -28,7 +25,7 @@ def get_expiration(ttl: Optional[int]) -> Optional[datetime.datetime]: if ttl is None: return None - return datetime.datetime.now() + datetime.timedelta(milliseconds=ttl) + return datetime.datetime.now() + datetime.timedelta(seconds=ttl) class StoreDocument(KeyDocument): diff --git a/fiftyone/operators/store/service.py b/fiftyone/operators/store/service.py index 47dcf1ada3..f01e3b874f 100644 --- a/fiftyone/operators/store/service.py +++ b/fiftyone/operators/store/service.py @@ -39,7 +39,7 @@ def set_key(self, store_name, key, value, ttl=None): store_name: the name of the store key: the key to set value: the value to set - ttl (None): an optional TTL in milliseconds + ttl (None): an optional TTL in seconds Returns: a :class:`fiftyone.store.models.KeyDocument` @@ -68,7 +68,7 @@ def delete_key(self, store_name, key): key: the key to delete Returns: - a :class:`fiftyone.store.models.KeyDocument` + `True` if the key was deleted, `False` otherwise """ return self._repo.delete_key(store_name=store_name, key=key) diff --git a/fiftyone/operators/store/store.py b/fiftyone/operators/store/store.py index d0fc738f9f..0fee055ea6 100644 --- a/fiftyone/operators/store/store.py +++ b/fiftyone/operators/store/store.py @@ -48,7 +48,7 @@ def set(self, key: str, value: Any, ttl: Optional[int] = None) -> None: 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 milliseconds. Defaults to None. + ttl (Optional[int], optional): The time-to-live in seconds. Defaults to None. """ self._store_service.set_key(self.store_name, key, value, ttl) @@ -57,8 +57,11 @@ def delete(self, key: str) -> None: Args: key (str): The key to delete. + + Returns: + bool: True if the key was deleted, False otherwise. """ - self._store_service.delete_key(self.store_name, key) + return self._store_service.delete_key(self.store_name, key) def has(self, key: str) -> bool: """Checks if the store has a specific key. diff --git a/tests/unittests/execution_store_tests.py b/tests/unittests/execution_store_tests.py index 12eb3a4c69..87d1b2f982 100644 --- a/tests/unittests/execution_store_tests.py +++ b/tests/unittests/execution_store_tests.py @@ -9,269 +9,225 @@ import time import unittest from unittest import mock -from unittest.mock import patch, MagicMock +from unittest.mock import patch, MagicMock, ANY, Mock +import datetime import fiftyone.operators as foo from fiftyone.operators.store import ExecutionStoreService from fiftyone.operators.store.models import StoreDocument, KeyDocument from fiftyone.factory.repo_factory import ExecutionStoreRepo -from fiftyone.operators.operator import Operator +from fiftyone.operators.store import ExecutionStore from fiftyone.operators.store import ExecutionStoreService +EPSILON = 0.1 -class MockOperator: - """Mock operator that interacts with ExecutionStore.""" - def __init__(self, ctx: foo.ExecutionContext): - self.ctx = ctx +class IsDateTime: + def __eq__(self, other): + return isinstance(other, datetime.datetime) - def execute(self): - # Example logic of interacting with the store - store = self.ctx.create_store("widgets") - # Set a value in the store - store.set("widget_1", {"name": "Widget One", "value": 100}, ttl=60000) +def assert_delta_seconds_approx(time_delta, seconds, epsilon=EPSILON): + assert abs(time_delta.total_seconds() - seconds) < epsilon - # Get the value back from the store - result = store.get("widget_1") - return result +class TeskKeyDocument(unittest.TestCase): + def test_get_expiration(self): + ttl = 1 + now = datetime.datetime.now() + expiration = KeyDocument.get_expiration(ttl) + time_delta = expiration - now + assert_delta_seconds_approx(time_delta, ttl) + assert isinstance(expiration, datetime.datetime) + def test_get_expiration_none(self): + ttl = None + expiration = KeyDocument.get_expiration(ttl) + assert expiration is None -class MockStoreRepo: - def __init__(self): - self.stores = {} - def create_store(self, store_name): - store = StoreDocument(store_name=store_name) - self.stores[store_name] = store - return store - - def set_key(self, store_name, key, value, ttl=None): - store = self.stores.get(store_name) - if store: - key_doc = KeyDocument(key=key, value=value, ttl=ttl) - store.keys[key] = key_doc - return key_doc - return None - - def get_key(self, store_name, key): - store = self.stores.get(store_name) - if store and key in store.keys: - return store.keys[key] - return None - - def delete_key(self, store_name, key): - store = self.stores.get(store_name) - if store and key in store.keys: - del store.keys[key] - - def delete_store(self, store_name): - if store_name in self.stores: - del self.stores[store_name] - - def update_ttl(self, store_name, key, ttl): - store = self.stores.get(store_name) - if store and key in store.keys: - store.keys[key].ttl = ttl - return store.keys[key] - return None +class ExecutionStoreServiceIntegrationTests(unittest.TestCase): + def setUp(self) -> None: + self.mock_collection = MagicMock() + self.store_repo = ExecutionStoreRepo(self.mock_collection) + self.store_service = ExecutionStoreService(self.store_repo) + def test_set_key(self): + self.store_repo.set_key( + "widgets", + "widget_1", + {"name": "Widget One", "value": 100}, + ttl=60000, + ) + self.mock_collection.update_one.assert_called_once() + self.mock_collection.update_one.assert_called_with( + {"store_name": "widgets", "key": "widget_1"}, + { + "$set": { + "store_name": "widgets", + "key": "widget_1", + "value": {"name": "Widget One", "value": 100}, + "created_at": IsDateTime(), + "updated_at": None, + "expires_at": IsDateTime(), + } + }, + upsert=True, + ) -class ExecutionStoreServiceTests(unittest.TestCase): - def setUp(self): - # Mock repository and service for testing - self.repo = MockStoreRepo() - self.svc = ExecutionStoreService(repo=self.repo) + def test_get_key(self): + self.mock_collection.find_one.return_value = { + "store_name": "widgets", + "key": "widget_1", + "value": {"name": "Widget One", "value": 100}, + "created_at": time.time(), + "updated_at": time.time(), + "expires_at": time.time() + 60000, + } + self.store_service.get_key(store_name="widgets", key="widget_1") + self.mock_collection.find_one.assert_called_once() + self.mock_collection.find_one.assert_called_with( + {"store_name": "widgets", "key": "widget_1"} + ) def test_create_store(self): - store_name = "test_store" - store = self.svc.create_store(store_name=store_name) - - self.assertIsNotNone(store) - self.assertEqual(store.store_name, store_name) - - def test_set_and_get_key(self): - store_name = "test_store" - self.svc.create_store(store_name=store_name) - - key = "test_key" - value = {"foo": "bar"} - ttl = 10000 - - key_doc = self.svc.set_key( - store_name=store_name, key=key, value=value, ttl=ttl + self.store_repo.create_store("widgets") + self.mock_collection.insert_one.assert_called_once() + self.mock_collection.insert_one.assert_called_with( + { + "store_name": "widgets", + "key": "__store__", + "value": None, + "created_at": IsDateTime(), + "updated_at": None, + "expires_at": None, + } ) - self.assertIsNotNone(key_doc) - self.assertEqual(key_doc.key, key) - self.assertEqual(key_doc.value, value) - self.assertEqual(key_doc.ttl, ttl) - - # Get the key and validate the value - fetched_key_doc = self.svc.get_key(store_name=store_name, key=key) - self.assertIsNotNone(fetched_key_doc) - self.assertEqual(fetched_key_doc.key, key) - self.assertEqual(fetched_key_doc.value, value) def test_delete_key(self): - store_name = "test_store" - self.svc.create_store(store_name=store_name) - - key = "test_key" - value = {"foo": "bar"} - self.svc.set_key(store_name=store_name, key=key, value=value) - - # Ensure the key exists - key_doc = self.svc.get_key(store_name=store_name, key=key) - self.assertIsNotNone(key_doc) - - # Delete the key - self.svc.delete_key(store_name=store_name, key=key) - - # Ensure the key is deleted - key_doc = self.svc.get_key(store_name=store_name, key=key) - self.assertIsNone(key_doc) + self.mock_collection.delete_one.return_value = Mock(deleted_count=1) + self.store_repo.delete_key("widgets", "widget_1") + self.mock_collection.delete_one.assert_called_once() + self.mock_collection.delete_one.assert_called_with( + {"store_name": "widgets", "key": "widget_1"} + ) def test_update_ttl(self): - store_name = "test_store" - self.svc.create_store(store_name=store_name) - - key = "test_key" - value = {"foo": "bar"} - self.svc.set_key(store_name=store_name, key=key, value=value) - - # Update TTL - new_ttl = 5000 - self.svc.update_ttl(store_name=store_name, key=key, new_ttl=new_ttl) + self.mock_collection.update_one.return_value = Mock(modified_count=1) + ttl_seconds = 60000 + expected_expiration = KeyDocument.get_expiration(ttl_seconds) + self.store_repo.update_ttl("widgets", "widget_1", ttl_seconds) + self.mock_collection.update_one.assert_called_once() - # Fetch key and check updated TTL - key_doc = self.svc.get_key(store_name=store_name, key=key) - self.assertIsNotNone(key_doc) - self.assertEqual(key_doc.ttl, new_ttl) + actual_call = self.mock_collection.update_one.call_args + actual_expires_at = actual_call[0][1]["$set"]["expires_at"] + time_delta = expected_expiration - actual_expires_at + assert_delta_seconds_approx(time_delta, 0, epsilon=0.0001) def test_delete_store(self): - store_name = "test_store" - self.svc.create_store(store_name=store_name) - - # Ensure the store exists - store = self.repo.stores.get(store_name) - self.assertIsNotNone(store) - - # Delete the store - self.svc.delete_store(store_name=store_name) - - # Ensure the store is deleted - store = self.repo.stores.get(store_name) - self.assertIsNone(store) - - def test_store_permissions(self): - store_name = "test_store" - self.svc.create_store(store_name=store_name) - - # Check default permissions - store = self.repo.stores.get(store_name) - self.assertIsNotNone(store) - - -class TestOperatorWithExecutionStore(Operator): - def execute(self, ctx): - # Create a store and interact with it - store = ctx.create_store("test_store") - - # Perform some store operations - store.set("my_key", {"foo": "bar"}) - value = store.get("my_key") - store.delete("my_key") + self.mock_collection.delete_many.return_value = Mock(deleted_count=1) + self.store_repo.delete_store("widgets") + self.mock_collection.delete_many.assert_called_once() + self.mock_collection.delete_many.assert_called_with( + {"store_name": "widgets"} + ) - return value + def test_list_keys(self): + self.mock_collection.find.return_value = [ + {"store_name": "widgets", "key": "widget_1"}, + {"store_name": "widgets", "key": "widget_2"}, + ] + keys = self.store_repo.list_keys("widgets") + assert keys == ["widget_1", "widget_2"] + self.mock_collection.find.assert_called_once() + self.mock_collection.find.assert_called_with({"store_name": "widgets"}) -class ExecutionStoreIntegrationTests(unittest.TestCase): - def setUp(self): - # Create a MagicMock for the MongoDB collection +class TestExecutionStoreIntegration(unittest.TestCase): + def setUp(self) -> None: self.mock_collection = MagicMock() - - # Instantiate the store repo and replace the _collection attribute self.store_repo = ExecutionStoreRepo(self.mock_collection) - - # Create the store service with the mocked repo self.store_service = ExecutionStoreService(self.store_repo) + self.store = ExecutionStore("mock_store", self.store_service) - # Create an execution context - from fiftyone.operators import ExecutionContext - - self.ctx = ExecutionContext() - self.ctx.store_service = ( - self.store_service - ) # Inject the store service into the context - - # Create an instance of the operator - self.operator = TestOperatorWithExecutionStore() - - def test_operator_execute_with_store(self): - # Mock MongoDB collection methods - self.mock_collection.update_one.return_value = None - self.mock_collection.find_one.return_value = { - "key": "my_key", - "value": {"foo": "bar"}, - } - self.mock_collection.delete_one.return_value = None - - # Call the operator's execute method - result = self.operator.execute(self.ctx) - - # Verify that the store interactions were made correctly - self.mock_collection.update_one.assert_called_once() # Checking that set operation inserts data - self.mock_collection.find_one.assert_called_once_with( - {"store_name": "test_store", "key": "my_key"} - ) # Checking that get operation retrieves data - self.mock_collection.delete_one.assert_called_once_with( - {"store_name": "test_store", "key": "my_key"} - ) # Checking that delete operation removes data - - # Verify the correct value was returned from the store - self.assertEqual(result, {"key": "my_key", "value": {"foo": "bar"}}) - - def test_operator_execute_set_key(self): - # Mock the insert_one call for the set operation - self.mock_collection.insert_one.return_value = None - - # Call the operator's execute method - self.operator.execute(self.ctx) - - # Check that insert_one (set) was called with the correct arguments - self.mock_collection.insert_one.assert_called_once_with( + def test_set(self): + self.store.set( + "widget_1", {"name": "Widget One", "value": 100}, ttl=60000 + ) + self.mock_collection.update_one.assert_called_once() + self.mock_collection.update_one.assert_called_with( + {"store_name": "mock_store", "key": "widget_1"}, { - "store_name": "test_store", - "key": "my_key", - "value": {"foo": "bar"}, - } + "$set": { + "store_name": "mock_store", + "key": "widget_1", + "value": {"name": "Widget One", "value": 100}, + "created_at": IsDateTime(), + "updated_at": None, + "expires_at": IsDateTime(), + } + }, + upsert=True, ) - def test_operator_execute_get_key(self): - # Mock the find_one call for the get operation + # def test_update(self): + # self.mock_collection.find_one.return_value = { + # "store_name": "mock_store", + # "key": "widget_1", + # "value": {"name": "Widget One", "value": 100}, + # "created_at": time.time(), + # "updated_at": time.time(), + # "expires_at": time.time() + 60000 + # } + # self.store.update_key("widget_1", {"name": "Widget One", "value": 200}) + # self.mock_collection.update_one.assert_called_once() + # self.mock_collection.update_one.assert_called_with( + # {"store_name": "mock_store", "key": "widget_1"}, + # { + # "$set": { + # "store_name": "mock_store", + # "key": "widget_1", + # "value": {"name": "Widget One", "value": 200}, + # "created_at": IsDateTime(), + # "updated_at": IsDateTime(), + # "expires_at": IsDateTime() + # } + # } + # ) + + def test_get(self): self.mock_collection.find_one.return_value = { - "key": "my_key", - "value": {"foo": "bar"}, + "store_name": "mock_store", + "key": "widget_1", + "value": {"name": "Widget One", "value": 100}, + "created_at": time.time(), + "updated_at": time.time(), + "expires_at": time.time() + 60000, } - - # Call the operator's execute method - result = self.operator.execute(self.ctx) - - # Check that find_one (get) was called correctly and returned the expected value - self.mock_collection.find_one.assert_called_once_with( - {"store_name": "test_store", "key": "my_key"} + value = self.store.get("widget_1") + assert value == {"name": "Widget One", "value": 100} + self.mock_collection.find_one.assert_called_once() + self.mock_collection.find_one.assert_called_with( + {"store_name": "mock_store", "key": "widget_1"} ) - self.assertEqual(result, {"key": "my_key", "value": {"foo": "bar"}}) - - def test_operator_execute_delete_key(self): - # Mock the delete_one call for the delete operation - self.mock_collection.delete_one.return_value = None - # Call the operator's execute method - self.operator.execute(self.ctx) + def test_list_keys(self): + self.mock_collection.find.return_value = [ + {"store_name": "mock_store", "key": "widget_1"}, + {"store_name": "mock_store", "key": "widget_2"}, + ] + keys = self.store.list_keys() + assert keys == ["widget_1", "widget_2"] + self.mock_collection.find.assert_called_once() + self.mock_collection.find.assert_called_with( + {"store_name": "mock_store"} + ) - # Check that delete_one was called with the correct arguments - self.mock_collection.delete_one.assert_called_once_with( - {"store_name": "test_store", "key": "my_key"} + def test_delete(self): + self.mock_collection.delete_one.return_value = Mock(deleted_count=1) + deleted = self.store.delete("widget_1") + assert deleted + self.mock_collection.delete_one.assert_called_once() + self.mock_collection.delete_one.assert_called_with( + {"store_name": "mock_store", "key": "widget_1"} ) From dcb0ecbe69834ac5b02331d825c6fbd2d95087df Mon Sep 17 00:00:00 2001 From: Ritchie Martori Date: Sun, 22 Sep 2024 16:24:11 -0700 Subject: [PATCH 6/9] add missing clear store test --- fiftyone/operators/store/store.py | 2 +- tests/unittests/execution_store_tests.py | 7 +++++++ 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/fiftyone/operators/store/store.py b/fiftyone/operators/store/store.py index 0fee055ea6..55ae6a15d1 100644 --- a/fiftyone/operators/store/store.py +++ b/fiftyone/operators/store/store.py @@ -76,7 +76,7 @@ def has(self, key: str) -> bool: def clear(self) -> None: """Clears all the data in the store.""" - self._store_service.clear_store(self.store_name) + self._store_service.delete_store(self.store_name) def update_ttl(self, key: str, new_ttl: int) -> None: """Updates the TTL for a specific key. diff --git a/tests/unittests/execution_store_tests.py b/tests/unittests/execution_store_tests.py index 87d1b2f982..e0cc1bc495 100644 --- a/tests/unittests/execution_store_tests.py +++ b/tests/unittests/execution_store_tests.py @@ -231,3 +231,10 @@ def test_delete(self): self.mock_collection.delete_one.assert_called_with( {"store_name": "mock_store", "key": "widget_1"} ) + + def test_clear(self): + self.store.clear() + self.mock_collection.delete_many.assert_called_once() + self.mock_collection.delete_many.assert_called_with( + {"store_name": "mock_store"} + ) From bccaa0c587567ef0c9cf0e2b1d6edfa8e4501e17 Mon Sep 17 00:00:00 2001 From: Ritchie Martori Date: Sun, 22 Sep 2024 16:43:06 -0700 Subject: [PATCH 7/9] 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: From 5e123628a38348abf6e0469f71d76915786a1115 Mon Sep 17 00:00:00 2001 From: Ritchie Martori Date: Sun, 22 Sep 2024 17:32:45 -0700 Subject: [PATCH 8/9] exec store cleanup and fixes --- fiftyone/factory/repos/execution_store.py | 50 ++++++++++++++++++----- tests/unittests/execution_store_tests.py | 47 +++++++-------------- 2 files changed, 54 insertions(+), 43 deletions(-) diff --git a/fiftyone/factory/repos/execution_store.py b/fiftyone/factory/repos/execution_store.py index 7f2c4aa117..69fa78b077 100644 --- a/fiftyone/factory/repos/execution_store.py +++ b/fiftyone/factory/repos/execution_store.py @@ -2,6 +2,7 @@ Execution store repository. """ +import datetime from pymongo.collection import Collection from fiftyone.operators.store.models import StoreDocument, KeyDocument @@ -35,18 +36,36 @@ def list_stores(self) -> list[str]: 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.""" + """Sets or updates a key in the specified store""" + now = datetime.datetime.now() expiration = KeyDocument.get_expiration(ttl) key_doc = KeyDocument( - store_name=store_name, - key=key, - value=value, - expires_at=expiration if ttl else None, + store_name=store_name, key=key, value=value, updated_at=now ) - # Update or insert the key - self._collection.update_one( - _where(store_name, key), {"$set": key_doc.dict()}, upsert=True + + # Prepare the update operations + update_fields = { + "$set": key_doc.dict( + exclude={"created_at", "expires_at", "store_name", "key"} + ), + "$setOnInsert": { + "store_name": store_name, + "key": key, + "created_at": now, + "expires_at": expiration if ttl else None, + }, + } + + # Perform the upsert operation + result = self._collection.update_one( + _where(store_name, key), update_fields, upsert=True ) + + if result.upserted_id: + key_doc.created_at = now + else: + key_doc.updated_at = now + return key_doc def get_key(self, store_name, key) -> KeyDocument: @@ -57,8 +76,7 @@ def get_key(self, store_name, key) -> KeyDocument: def list_keys(self, store_name) -> list[str]: """Lists all keys in the specified store.""" - keys = self._collection.find(_where(store_name)) - # TODO: redact non-key fields + keys = self._collection.find(_where(store_name), {"key": 1}) return [key["key"] for key in keys] def update_ttl(self, store_name, key, ttl) -> bool: @@ -92,7 +110,19 @@ def __init__(self, collection: Collection): def _create_indexes(self): indices = self._collection.list_indexes() expires_at_name = "expires_at" + store_name_name = "store_name" + key_name = "key" + full_key_name = "store_name_and_key" if expires_at_name not in indices: self._collection.create_index( expires_at_name, name=expires_at_name, expireAfterSeconds=0 ) + if full_key_name not in indices: + self._collection.create_index( + [(store_name_name, 1), (key_name, 1)], + name=full_key_name, + unique=True, + ) + for name in [store_name_name, key_name]: + if name not in indices: + self._collection.create_index(name, name=name) diff --git a/tests/unittests/execution_store_tests.py b/tests/unittests/execution_store_tests.py index d5b145ba63..470d399026 100644 --- a/tests/unittests/execution_store_tests.py +++ b/tests/unittests/execution_store_tests.py @@ -64,13 +64,15 @@ def test_set_key(self): {"store_name": "widgets", "key": "widget_1"}, { "$set": { + "value": {"name": "Widget One", "value": 100}, + "updated_at": IsDateTime(), + }, + "$setOnInsert": { "store_name": "widgets", "key": "widget_1", - "value": {"name": "Widget One", "value": 100}, "created_at": IsDateTime(), - "updated_at": None, "expires_at": IsDateTime(), - } + }, }, upsert=True, ) @@ -140,7 +142,9 @@ def test_list_keys(self): keys = self.store_repo.list_keys("widgets") assert keys == ["widget_1", "widget_2"] self.mock_collection.find.assert_called_once() - self.mock_collection.find.assert_called_with({"store_name": "widgets"}) + self.mock_collection.find.assert_called_with( + {"store_name": "widgets"}, {"key": 1} + ) def test_list_stores(self): self.mock_collection.distinct.return_value = ["widgets", "gadgets"] @@ -166,42 +170,19 @@ def test_set(self): {"store_name": "mock_store", "key": "widget_1"}, { "$set": { + "updated_at": IsDateTime(), + "value": {"name": "Widget One", "value": 100}, + }, + "$setOnInsert": { "store_name": "mock_store", "key": "widget_1", - "value": {"name": "Widget One", "value": 100}, "created_at": IsDateTime(), - "updated_at": None, "expires_at": IsDateTime(), - } + }, }, upsert=True, ) - # def test_update(self): - # self.mock_collection.find_one.return_value = { - # "store_name": "mock_store", - # "key": "widget_1", - # "value": {"name": "Widget One", "value": 100}, - # "created_at": time.time(), - # "updated_at": time.time(), - # "expires_at": time.time() + 60000 - # } - # self.store.update_key("widget_1", {"name": "Widget One", "value": 200}) - # self.mock_collection.update_one.assert_called_once() - # self.mock_collection.update_one.assert_called_with( - # {"store_name": "mock_store", "key": "widget_1"}, - # { - # "$set": { - # "store_name": "mock_store", - # "key": "widget_1", - # "value": {"name": "Widget One", "value": 200}, - # "created_at": IsDateTime(), - # "updated_at": IsDateTime(), - # "expires_at": IsDateTime() - # } - # } - # ) - def test_get(self): self.mock_collection.find_one.return_value = { "store_name": "mock_store", @@ -227,7 +208,7 @@ def test_list_keys(self): assert keys == ["widget_1", "widget_2"] self.mock_collection.find.assert_called_once() self.mock_collection.find.assert_called_with( - {"store_name": "mock_store"} + {"store_name": "mock_store"}, {"key": 1} ) def test_delete(self): From a3627b411ed143b94cd10e893e3db09694d546ab Mon Sep 17 00:00:00 2001 From: Br2850 Date: Mon, 30 Sep 2024 15:08:16 -0400 Subject: [PATCH 9/9] add spinner variant --- .../src/components/Loading/LoadingSpinner.tsx | 29 +++++++++++++++++++ .../SchemaIO/components/LoadingView.tsx | 9 ++++-- 2 files changed, 36 insertions(+), 2 deletions(-) create mode 100644 app/packages/components/src/components/Loading/LoadingSpinner.tsx diff --git a/app/packages/components/src/components/Loading/LoadingSpinner.tsx b/app/packages/components/src/components/Loading/LoadingSpinner.tsx new file mode 100644 index 0000000000..b14e99daf3 --- /dev/null +++ b/app/packages/components/src/components/Loading/LoadingSpinner.tsx @@ -0,0 +1,29 @@ +import React from "react"; + +import { CircularProgress } from "@mui/material"; + +const LoadingSpinner = ({ + color = "base", + size = "medium", +}: { + color?: string; + size?: string; +}) => { + const COLORS: { [key: string]: string } = { + base: "#FFC59B", + primary: "primary", + secondary: "secondary", + error: "error", + warning: "warning", + info: "info", + success: "success", + }; + const SIZES: { [key: string]: string } = { + small: "1rem", + medium: "2rem", + large: "3rem", + }; + return ; +}; + +export default LoadingSpinner; diff --git a/app/packages/core/src/plugins/SchemaIO/components/LoadingView.tsx b/app/packages/core/src/plugins/SchemaIO/components/LoadingView.tsx index 719c6f3ec6..bf5bf4090e 100644 --- a/app/packages/core/src/plugins/SchemaIO/components/LoadingView.tsx +++ b/app/packages/core/src/plugins/SchemaIO/components/LoadingView.tsx @@ -1,4 +1,5 @@ import LoadingDots from "@fiftyone/components/src/components/Loading/LoadingDots"; +import LoadingSpinner from "@fiftyone/components/src/components/Loading/LoadingSpinner"; import { Box } from "@mui/material"; import React from "react"; import { getComponentProps } from "../utils"; @@ -6,11 +7,15 @@ import { getComponentProps } from "../utils"; export default function LoadingView(props) { const { schema } = props; const { view = {} } = schema; - const { label = "Loading" } = view; + const { text = "Loading", variant, color, size } = view; return ( - + {variant === "spinner" ? ( + + ) : ( + + )} ); }