added base scheduler store

Author: Andrei Neagu

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

@@ -0,0 +1,6 @@
+from typing import TypeAlias
+
+ScheduleId: TypeAlias = str
+OperationName: TypeAlias = str
+StepGroup: TypeAlias = str
+StepName: TypeAlias = str

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

@@ -0,0 +1,168 @@
+import base64
+import pickle
+from typing import Any, Final, Literal, NotRequired, TypedDict, overload
+
+import redis.asyncio as aioredis
+from pydantic import NonNegativeInt
+from servicelib.redis._client import RedisClientSDK
+from servicelib.redis._utils import handle_redis_returns_union_types
+from settings_library.redis import RedisDatabase, RedisSettings
+
+from ._models import OperationName, ScheduleId, StepGroup, StepName
+
+_SCHEDULE_NAMESPACE: Final[str] = "SCH"
+_STEPS_KEY: Final[str] = "STEPS"
+# Figure out hwo to store data in Redis as flat as possible
+
+
+def _get_scheduler_data_hash_key(*, schedule_id: ScheduleId) -> str:
+    # SCHEDULE_NAMESPACE:SCHEDULE_ID:KEY
+    # Example:
+    # - SCH:00000000-0000-0000-0000-000000000000
+    return f"{_SCHEDULE_NAMESPACE}:{schedule_id}"
+
+
+def _get_step_hash_key(
+    *,
+    schedule_id: ScheduleId,
+    operation_name: OperationName,
+    group: StepGroup,
+    step_name: StepName,
+) -> str:
+    # SCHEDULE_NAMESPACE:SCHEDULE_ID:STEPS:OPERATION_NAME:GROUP_INDEX:STEP_NAME:KEY
+    # - SCHEDULE_NAMESPACE: something short to identify tgis
+    # - SCHEDULE_ID: the unique scheudle_id assigned
+    # - STEPS: the constant "STEPS"
+    # - OPERATION_NAME form the vairble's name during registration
+    # - GROUP_INDEX
+    #   -> "{index}(S|P)[R]": S=single or P=parallel and optinally, "R" if steps should be repeated forever
+    # - STEP_NAME form it's class
+    # Example:
+    # - SCH:00000000-0000-0000-0000-000000000000:STEPS:START_SERVICE:0S:BS1
+    return f"{_SCHEDULE_NAMESPACE}:{schedule_id}:{_STEPS_KEY}:{operation_name}:{group}:{step_name}"
+
+
+def _dumps(obj: Any) -> str:
+    return base64.b85encode(pickle.dumps(obj)).decode("utf-8")
+
+
+def _loads(obj_str: str) -> Any:
+    return pickle.loads(base64.b85decode(obj_str))  # noqa: S301
+
+
+class Store:
+    def __init__(self, redis_settings: RedisSettings) -> None:
+        self.redis_settings = redis_settings
+
+        self._client: RedisClientSDK | None = None
+
+    async def setup(self) -> None:
+        self._client = RedisClientSDK(
+            self.redis_settings.build_redis_dsn(RedisDatabase.DYNAMIC_SERVICES),
+            client_name=__name__,
+        )
+        await self._client.setup()
+
+    async def shutdown(self) -> None:
+        if self._client:
+            await self._client.shutdown()
+
+    @property
+    def _redis(self) -> aioredis.Redis:
+        assert self._client  # nosec
+        return self._client.redis
+
+    async def set_multiple(self, hash_key: str, updates: dict[str, Any]) -> None:
+        """saves multiple key-value pairs in a hash"""
+        await handle_redis_returns_union_types(
+            self._redis.hset(
+                hash_key, mapping={k: _dumps(v) for k, v in updates.items()}
+            )
+        )
+
+    async def set(self, hash_key: str, key: str, value: Any) -> None:
+        """saves a single key-value pair in a hash"""
+        await self.set_multiple(hash_key, {key: value})
+
+    async def get(self, hash_key: str, *keys: str) -> tuple[Any, ...]:
+        """retrieves one or more keys from a hash"""
+        result: list[str | None] = await handle_redis_returns_union_types(
+            self._redis.hmget(hash_key, list(keys))
+        )
+        return tuple(_loads(x) if x else None for x in result)
+
+    async def delete(self, hash_key: str, *keys: str) -> None:
+        """removes one or more keys form a hash"""
+        await handle_redis_returns_union_types(self._redis.hdel(hash_key, *keys))
+
+    async def remove(self, hash_key: str) -> None:
+        """removes the entire hash"""
+        await handle_redis_returns_union_types(self._redis.delete(hash_key))
+
+
+class _UpdateScheduleDataDict(TypedDict):
+    operation_name: NotRequired[OperationName]
+    group_index: NotRequired[NonNegativeInt]
+    is_creating: NotRequired[bool]
+
+
+_DeleteScheduleDataKeys = Literal["operation_name", "group_index", "is_creating"]
+
+
+class ScheduleDataStoreProxy:
+    def __init__(self, *, store: Store, schedule_id: ScheduleId) -> None:
+        self._store = store
+        self._schedule_id = schedule_id
+
+    def _get_hash_key(self) -> str:
+        return _get_scheduler_data_hash_key(schedule_id=self._schedule_id)
+
+    @overload
+    async def get(self, key: Literal["operation_name"]) -> str: ...
+    @overload
+    async def get(self, key: Literal["group_index"]) -> int: ...
+    @overload
+    async def get(self, key: Literal["is_creating"]) -> bool: ...
+    async def get(self, key: str) -> Any:
+        (result,) = await self._store.get(self._get_hash_key(), key)
+        return result
+
+    @overload
+    async def set(self, key: Literal["operation_name"], value: str) -> None: ...
+    @overload
+    async def set(self, key: Literal["group_index"], value: int) -> None: ...
+    @overload
+    async def set(self, key: Literal["is_creating"], *, value: bool) -> 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: _UpdateScheduleDataDict) -> None:
+        await self._store.set_multiple(self._get_hash_key(), updates=values)
+
+    async def delete(self, *keys: _DeleteScheduleDataKeys) -> None:
+        await self._store.delete(self._get_hash_key(), *keys)
+
+
+class StepStoreProxy:
+    def __init__(
+        self,
+        *,
+        store: Store,
+        schedule_id: ScheduleId,
+        operation_name: OperationName,
+        group: StepGroup,
+        step_name: StepName,
+    ) -> None:
+        self._store = store
+        self._schedule_id = schedule_id
+        self._operation_name = operation_name
+        self._group = group
+        self._step_name = step_name
+
+    def _get_hash_key(self) -> str:
+        return _get_step_hash_key(
+            schedule_id=self._schedule_id,
+            operation_name=self._operation_name,
+            group=self._group,
+            step_name=self._step_name,
+        )

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

@@ -0,0 +1,156 @@
+# pylint: disable=protected-access
+# pylint: disable=redefined-outer-name
+# pylint: disable=unused-argument
+
+from collections.abc import AsyncIterable
+from enum import Enum
+from typing import Any, Self
+
+import pytest
+from faker import Faker
+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 (
+    ScheduleId,
+)
+from simcore_service_dynamic_scheduler.services.generic_scheduler._store import (
+    ScheduleDataStoreProxy,
+    Store,
+)
+
+
+@pytest.fixture
+async def store(use_in_memory_redis: RedisSettings) -> AsyncIterable[Store]:
+    store = Store(use_in_memory_redis)
+    await store.setup()
+    yield store
+    await store._redis.client().flushall()  # noqa: SLF001
+    await store.shutdown()
+
+
+async def _assert_keys(store: Store, expected_keys: set[str]) -> None:
+    keys = set(await store._redis.keys())  # noqa: SLF001
+    assert keys == expected_keys
+
+
+async def _assert_keys_in_hash(
+    store: Store, hash_key: str, expected_keys: set[str]
+) -> None:
+    keys = set(
+        await handle_redis_returns_union_types(
+            store._redis.hkeys(hash_key)  # noqa: SLF001
+        )
+    )
+    assert keys == expected_keys
+
+
+async def test_store_workflow(store: Store):
+    # save single value
+    await store.set("hash1", "key1", "value1")
+    await _assert_keys(store, {"hash1"})
+    await _assert_keys_in_hash(store, "hash1", {"key1"})
+    assert await store.get("hash1", "key1") == ("value1",)
+    assert await store.get("hash1", "key1", "key1") == ("value1", "value1")
+    assert await store.get("hash1", "missing1", "missing2") == (None, None)
+
+    # remove last key in hash
+    await store.delete("hash1", "key1")
+    await _assert_keys(store, set())
+    await _assert_keys_in_hash(store, "hash1", set())
+    assert await store.get("hash1", "key1") == (None,)
+
+    # save multiple values
+    await store.set_multiple("hash2", {"key1": "value1", "key2": 2, "key3": True})
+    await _assert_keys(store, {"hash2"})
+    await _assert_keys_in_hash(store, "hash2", {"key1", "key2", "key3"})
+    assert await store.get("hash2", "key1", "key2", "key3") == ("value1", 2, True)
+
+    # delete a few keys form hash
+    await store.delete("hash2", "key1", "key3", "missing1", "missing2", "missing3")
+    await _assert_keys(store, {"hash2"})
+    await _assert_keys_in_hash(store, "hash2", {"key2"})
+    assert await store.get("hash2", "key1", "key2", "key3") == (None, 2, None)
+
+    # remove hash completely
+    await store.remove("hash2")
+    await _assert_keys(store, set())
+    await _assert_keys_in_hash(store, "hash2", set())
+    assert await store.get("hash2", "key1", "key2", "key3") == (None, None, None)
+
+
+class _BS(str, Enum):
+    A = "A"
+    B = "B"
+    C = "C"
+
+
+class _CustomObj:
+    def __init__(self, a: int, b: str) -> None:
+        self.a = a
+        self.b = b
+
+    def __eq__(self, other: Self) -> bool:
+        if not isinstance(other, _CustomObj):
+            return False
+        return self.a == other.a and self.b == other.b
+
+
+@pytest.mark.parametrize(
+    "value",
+    [
+        1,
+        "a",
+        b"some_bytes",
+        3.14,
+        True,
+        None,
+        {"dict": "with_data"},
+        {"a", "set"},
+        {"some": _BS.A, "enum": _BS.B},
+        _CustomObj(1, "a"),
+    ],
+)
+async def test_store_supporse_multiple_python_base_types(store: Store, value: Any):
+    await store.set("hash1", "key1", value)
+    assert (await store.get("hash1", "key1")) == (value,)
+
+
+@pytest.fixture
+def schedule_id(faker: Faker) -> ScheduleId:
+    return faker.uuid4()
+
+
+async def test_schedule_data_store_proxy_workflow(
+    store: Store, schedule_id: ScheduleId
+):
+    proxy = ScheduleDataStoreProxy(store=store, schedule_id=schedule_id)
+    hash_key = f"SCH:{schedule_id}"
+
+    # set
+    await proxy.set("operation_name", "op1")
+    await proxy.set("group_index", 1)
+    await proxy.set("is_creating", value=True)
+    await _assert_keys(store, {hash_key})
+    await _assert_keys_in_hash(
+        store, hash_key, {"operation_name", "group_index", "is_creating"}
+    )
+
+    # get
+    assert await proxy.get("operation_name") == "op1"
+    assert await proxy.get("group_index") == 1
+    assert await proxy.get("is_creating") is True
+
+    # remove
+    await proxy.delete("operation_name", "is_creating", "group_index")
+    await _assert_keys(store, set())
+    await _assert_keys_in_hash(store, hash_key, set())
+
+    # set multiple
+    await proxy.set_multiple({"group_index": 2, "is_creating": False})
+    await _assert_keys(store, {hash_key})
+    await _assert_keys_in_hash(store, hash_key, {"group_index", "is_creating"})
+
+    # remove all keys an even missing ones
+    await proxy.delete("operation_name", "is_creating", "group_index")
+    await _assert_keys(store, set())
+    await _assert_keys_in_hash(store, hash_key, set())