merge latest execution store

fiftyone/factory/repos/execution_store.py

@@ -24,10 +24,8 @@ def __init__(self, collection: Collection):
 
     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
-        )
-        self._collection.insert_one(store_doc.dict())
+        store_doc = StoreDocument(store_name=store_name, value=permissions)
+        self._collection.insert_one(store_doc.to_mongo_dict())
         return store_doc
 
     def list_stores(self) -> list[str]:
@@ -40,14 +38,21 @@ def set_key(self, store_name, key, value, ttl=None) -> KeyDocument:
         now = datetime.datetime.now()
         expiration = KeyDocument.get_expiration(ttl)
         key_doc = KeyDocument(
-            store_name=store_name, key=key, value=value, updated_at=now
+            store_name=store_name,
+            key=key,
+            value=value,
+            updated_at=now,
+            expires_at=expiration,
         )
 
         # Prepare the update operations
         update_fields = {
-            "$set": key_doc.dict(
-                exclude={"created_at", "expires_at", "store_name", "key"}
-            ),
+            "$set": {
+                k: v
+                for k, v in key_doc.to_mongo_dict().items()
+                if k
+                not in {"_id", "created_at", "expires_at", "store_name", "key"}
+            },
             "$setOnInsert": {
                 "store_name": store_name,
                 "key": key,
@@ -101,8 +106,6 @@ def delete_store(self, store_name) -> int:
 class MongoExecutionStoreRepo(ExecutionStoreRepo):
     """MongoDB implementation of execution store repository."""
 
-    COLLECTION_NAME = "execution_store"
-
     def __init__(self, collection: Collection):
         super().__init__(collection)
         self._create_indexes()

fiftyone/operators/store/models.py

@@ -1,21 +1,19 @@
-"""
-Store and key models for the execution store.
-"""
-
-from pydantic import BaseModel, Field
+from dataclasses import dataclass, field, asdict
 from typing import Optional, Dict, Any
 import datetime
 
 
-class KeyDocument(BaseModel):
+@dataclass
+class KeyDocument:
     """Model representing a key in the store."""
 
     store_name: str
     key: str
     value: Any
-    created_at: datetime.datetime = Field(
+    created_at: datetime.datetime = field(
         default_factory=datetime.datetime.now
     )
+    _id: Optional[Any] = None
     updated_at: Optional[datetime.datetime] = None
     expires_at: Optional[datetime.datetime] = None
 
@@ -24,10 +22,17 @@ 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(seconds=ttl)
 
+    def to_mongo_dict(self, exclude_id: bool = True) -> Dict[str, Any]:
+        """Serializes the document to a MongoDB dictionary."""
+        data = asdict(self)
+        if exclude_id:
+            data.pop("_id", None)
+        return data
+
 
+@dataclass
 class StoreDocument(KeyDocument):
     """Model representing a Store."""
 

fiftyone/operators/store/service.py

@@ -7,21 +7,24 @@
 """
 
 import logging
+from typing import Optional, List
 from fiftyone.factory.repo_factory import RepositoryFactory
+from fiftyone.operators.store.models import StoreDocument, KeyDocument
+from fiftyone.factory.repos.execution_store import ExecutionStoreRepo
 
 logger = logging.getLogger(__name__)
 
 
-class ExecutionStoreService(object):
+class ExecutionStoreService:
     """Service for managing execution store operations."""
 
-    def __init__(self, repo=None):
+    def __init__(self, repo: Optional[ExecutionStoreRepo] = None):
         if repo is None:
             repo = RepositoryFactory.execution_store_repo()
 
-        self._repo = repo
+        self._repo: ExecutionStoreRepo = repo
 
-    def create_store(self, store_name):
+    def create_store(self, store_name: str) -> StoreDocument:
         """Creates a new store with the specified name.
 
         Args:
@@ -32,7 +35,9 @@ def create_store(self, store_name):
         """
         return self._repo.create_store(store_name=store_name)
 
-    def set_key(self, store_name, key, value, ttl=None):
+    def set_key(
+        self, store_name: str, key: str, value: str, ttl: Optional[int] = None
+    ) -> KeyDocument:
         """Sets the value of a key in the specified store.
 
         Args:
@@ -48,7 +53,7 @@ def set_key(self, store_name, key, value, ttl=None):
             store_name=store_name, key=key, value=value, ttl=ttl
         )
 
-    def get_key(self, store_name, key):
+    def get_key(self, store_name: str, key: str) -> KeyDocument:
         """Retrieves the value of a key from the specified store.
 
         Args:
@@ -60,7 +65,7 @@ def get_key(self, store_name, key):
         """
         return self._repo.get_key(store_name=store_name, key=key)
 
-    def delete_key(self, store_name, key):
+    def delete_key(self, store_name: str, key: str) -> bool:
         """Deletes the specified key from the store.
 
         Args:
@@ -72,13 +77,15 @@ def delete_key(self, store_name, key):
         """
         return self._repo.delete_key(store_name=store_name, key=key)
 
-    def update_ttl(self, store_name, key, new_ttl):
+    def update_ttl(
+        self, store_name: str, key: str, new_ttl: int
+    ) -> KeyDocument:
         """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
+            new_ttl: the new TTL in seconds
 
         Returns:
             a :class:`fiftyone.store.models.KeyDocument`
@@ -87,15 +94,15 @@ def update_ttl(self, store_name, key, new_ttl):
             store_name=store_name, key=key, ttl=new_ttl
         )
 
-    def list_stores(self):
+    def list_stores(self) -> List[StoreDocument]:
         """Lists all stores matching the given criteria.
 
         Returns:
             a list of :class:`fiftyone.store.models.StoreDocument`
         """
         return self._repo.list_stores()
 
-    def delete_store(self, store_name):
+    def delete_store(self, store_name: str) -> StoreDocument:
         """Deletes the specified store.
 
         Args:
@@ -106,10 +113,13 @@ def delete_store(self, store_name):
         """
         return self._repo.delete_store(store_name=store_name)
 
-    def list_keys(self, store_name):
+    def list_keys(self, store_name: str) -> List[str]:
         """Lists all keys in the specified store.
 
         Args:
             store_name: the name of the store
+
+        Returns:
+            a list of keys in the store
         """
         return self._repo.list_keys(store_name)

fiftyone/operators/store/store.py

@@ -7,7 +7,6 @@
 """
 
 import logging
-from fiftyone.factory.repo_factory import RepositoryFactory
 from fiftyone.operators.store.service import ExecutionStoreService
 from typing import Any, Optional
 
@@ -60,7 +59,7 @@ def set(self, key: str, value: Any, ttl: Optional[int] = None) -> None:
         """
         self._store_service.set_key(self.store_name, key, value, ttl)
 
-    def delete(self, key: str) -> None:
+    def delete(self, key: str) -> bool:
         """Deletes a key from the store.
 
         Args:
@@ -91,7 +90,7 @@ def update_ttl(self, key: str, new_ttl: int) -> None:
 
         Args:
             key (str): The key to update the TTL for.
-            new_ttl (int): The new TTL in milliseconds.
+            new_ttl (int): The new TTL in seconds.
         """
         self._store_service.update_ttl(self.store_name, key, new_ttl)
 
@@ -102,7 +101,7 @@ def get_ttl(self, key: str) -> Optional[int]:
             key (str): The key to get the TTL for.
 
         Returns:
-            Optional[int]: The TTL in milliseconds, or None if the key does not have a TTL.
+            Optional[int]: The TTL in seconds, or None if the key does not have a TTL.
         """
         return self._store_service.get_ttl(self.store_name, key)
 

tests/unittests/execution_store_tests.py

@@ -8,16 +8,13 @@
 
 import time
 import unittest
-from unittest import mock
 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.operators.store.models import KeyDocument
 from fiftyone.factory.repo_factory import ExecutionStoreRepo
 from fiftyone.operators.store import ExecutionStore
-from fiftyone.operators.store import ExecutionStoreService
 
 EPSILON = 0.1