extended store

Author: Andrei Neagu

services/dynamic-scheduler/src/simcore_service_dynamic_scheduler/services/generic_scheduler/_store.py

@@ -15,6 +15,7 @@
     ScheduleId,
     StepGroupName,
     StepName,
+    StepStatus,
 )
 
 _SCHEDULE_NAMESPACE: Final[str] = "SCH"
@@ -167,20 +168,27 @@ async def delete(self, *keys: _DeleteScheduleDataKeys) -> None:
         await self._store.delete(self._get_hash_key(), *keys)
 
 
+class _StepDict(TypedDict):
+    status: NotRequired[StepStatus]
+
+
+_DeleteStepKeys = Literal["status"]
+
+
 class StepStoreProxy:
     def __init__(
         self,
         *,
         store: Store,
         schedule_id: ScheduleId,
         operation_name: OperationName,
-        group: StepGroupName,
+        step_group_name: StepGroupName,
         step_name: StepName,
     ) -> None:
         self._store = store
         self._schedule_id = schedule_id
         self._operation_name = operation_name
-        self._group = group
+        self._group = step_group_name
         self._step_name = step_name
 
     def _get_hash_key(self) -> str:
@@ -190,3 +198,26 @@ def _get_hash_key(self) -> str:
             group=self._group,
             step_name=self._step_name,
         )
+
+    @overload
+    async def get(self, key: Literal["status"]) -> StepStatus: ...
+    async def get(self, key: str) -> Any:
+        """raises KeyNotFoundInHashError if the key is not present in the hash"""
+        hash_key = self._get_hash_key()
+        (result,) = await self._store.get(hash_key, key)
+        if result is None:
+            raise KeyNotFoundInHashError(
+                schedule_id=self._schedule_id, hash_key=hash_key
+            )
+        return result
+
+    @overload
+    async def set(self, key: Literal["status"], value: StepStatus) -> None: ...
+    async def set(self, key: str, value: Any) -> None:
+        await self._store.set(self._get_hash_key(), key, value)
+
+    async def set_multiple(self, values: _StepDict) -> None:
+        await self._store.set_multiple(self._get_hash_key(), updates=values)
+
+    async def delete(self, *keys: _DeleteStepKeys) -> None:
+        await self._store.delete(self._get_hash_key(), *keys)

services/dynamic-scheduler/tests/unit/service_generic_scheduler/test__store.py

@@ -12,9 +12,11 @@
 from settings_library.redis import RedisSettings
 from simcore_service_dynamic_scheduler.services.generic_scheduler._models import (
     ScheduleId,
+    StepStatus,
 )
 from simcore_service_dynamic_scheduler.services.generic_scheduler._store import (
     ScheduleDataStoreProxy,
+    StepStoreProxy,
     Store,
 )
 
@@ -154,3 +156,38 @@ async def test_schedule_data_store_proxy_workflow(
     await proxy.delete("operation_name", "is_creating", "group_index")
     await _assert_keys(store, set())
     await _assert_keys_in_hash(store, hash_key, set())
+
+
+async def test_step_store_proxy_workflow(store: Store, schedule_id: ScheduleId):
+    step_name = "MyStep"
+    proxy = StepStoreProxy(
+        store=store,
+        schedule_id=schedule_id,
+        operation_name="op1",
+        step_group_name="sg1",
+        step_name=step_name,
+    )
+    hash_key = f"SCH:{schedule_id}:STEPS:op1:sg1:{step_name}"
+
+    # set
+    await proxy.set("status", StepStatus.RUNNING)
+    await _assert_keys(store, {hash_key})
+    await _assert_keys_in_hash(store, hash_key, {"status"})
+
+    # get
+    assert await proxy.get("status") == StepStatus.RUNNING
+
+    # remove
+    await proxy.delete("status")
+    await _assert_keys(store, set())
+    await _assert_keys_in_hash(store, hash_key, set())
+
+    # set multiple
+    await proxy.set_multiple({"status": StepStatus.SUCCESS})
+    await _assert_keys(store, {hash_key})
+    await _assert_keys_in_hash(store, hash_key, {"status"})
+
+    # remove all keys an even missing ones
+    await proxy.delete("status")
+    await _assert_keys(store, set())
+    await _assert_keys_in_hash(store, hash_key, set())