added deferred_task_uid to store

Author: Andrei Neagu

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

@@ -4,6 +4,7 @@
 
 import redis.asyncio as aioredis
 from pydantic import NonNegativeInt
+from servicelib.deferred_tasks import TaskUID
 from servicelib.redis._client import RedisClientSDK
 from servicelib.redis._utils import handle_redis_returns_union_types
 from settings_library.redis import RedisDatabase, RedisSettings
@@ -170,9 +171,10 @@ async def delete(self, *keys: _DeleteScheduleDataKeys) -> None:
 
 class _StepDict(TypedDict):
     status: NotRequired[StepStatus]
+    deferred_task_uid: NotRequired[TaskUID]
 
 
-_DeleteStepKeys = Literal["status"]
+_DeleteStepKeys = Literal["status", "deferred_task_uid"]
 
 
 class StepStoreProxy:
@@ -201,6 +203,8 @@ def _get_hash_key(self) -> str:
 
     @overload
     async def get(self, key: Literal["status"]) -> StepStatus: ...
+    @overload
+    async def get(self, key: Literal["deferred_task_uid"]) -> TaskUID: ...
     async def get(self, key: str) -> Any:
         """raises KeyNotFoundInHashError if the key is not present in the hash"""
         hash_key = self._get_hash_key()
@@ -213,6 +217,8 @@ async def get(self, key: str) -> Any:
 
     @overload
     async def set(self, key: Literal["status"], value: StepStatus) -> None: ...
+    @overload
+    async def set(self, key: Literal["deferred_task_uid"], value: TaskUID) -> None: ...
     async def set(self, key: str, value: Any) -> None:
         await self._store.set(self._get_hash_key(), key, value)
 

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

@@ -8,6 +8,7 @@
 
 import pytest
 from faker import Faker
+from servicelib.deferred_tasks import TaskUID
 from servicelib.redis._utils import handle_redis_returns_union_types
 from settings_library.redis import RedisSettings
 from simcore_service_dynamic_scheduler.services.generic_scheduler._models import (
@@ -183,9 +184,11 @@ async def test_step_store_proxy_workflow(store: Store, schedule_id: ScheduleId):
     await _assert_keys_in_hash(store, hash_key, set())
 
     # set multiple
-    await proxy.set_multiple({"status": StepStatus.SUCCESS})
+    await proxy.set_multiple(
+        {"status": StepStatus.SUCCESS, "deferred_task_uid": TaskUID("mytask")}
+    )
     await _assert_keys(store, {hash_key})
-    await _assert_keys_in_hash(store, hash_key, {"status"})
+    await _assert_keys_in_hash(store, hash_key, {"status", "deferred_task_uid"})
 
     # remove all keys an even missing ones
     await proxy.delete("status")