Add cluster option to result backend

Author: stinovlas

.github/workflows/test.yml

@@ -42,12 +42,30 @@ jobs:
           --health-retries=30
         ports:
           - 6379:6379
+      redis-cluster:
+        image: bitnami/redis-cluster:6.2.5
+        env:
+          ALLOW_EMPTY_PASSWORD: "yes"
+          REDIS_NODES: "localhost"
+        options: >-
+          --health-cmd="redis-cli ping"
+          --health-interval=5s
+          --health-timeout=5s
+          --health-retries=30
+        ports:
+          - 7000:6379
     strategy:
       matrix:
         py_version: ["3.8", "3.9", "3.10", "3.11", "3.12"]
     runs-on: "ubuntu-latest"
     steps:
       - uses: actions/checkout@v4
+      - uses: shogo82148/actions-setup-redis@v1
+        with:
+          redis-version: "6.x"
+          auto-start: false
+      - name: Set up single-node Redis cluster
+        run: redis-cli -h localhost -p 7000 --cluster-yes CLUSTER ADDSLOTSRANGE 0 16383
       - name: Set up Python
         uses: actions/setup-python@v2
         with:

taskiq_redis/__init__.py

@@ -1,9 +1,13 @@
 """Package for redis integration."""
-from taskiq_redis.redis_backend import RedisAsyncResultBackend
+from taskiq_redis.redis_backend import (
+    RedisAsyncClusterResultBackend,
+    RedisAsyncResultBackend,
+)
 from taskiq_redis.redis_broker import ListQueueBroker, PubSubBroker
 from taskiq_redis.schedule_source import RedisScheduleSource
 
 __all__ = [
+    "RedisAsyncClusterResultBackend",
     "RedisAsyncResultBackend",
     "ListQueueBroker",
     "PubSubBroker",

taskiq_redis/redis_backend.py

@@ -1,7 +1,7 @@
 import pickle
 from typing import Dict, Optional, TypeVar, Union
 
-from redis.asyncio import ConnectionPool, Redis
+from redis.asyncio import ConnectionPool, Redis, RedisCluster
 from taskiq import AsyncResultBackend
 from taskiq.abc.result_backend import TaskiqResult
 
@@ -134,3 +134,117 @@ async def get_result(
             taskiq_result.log = None
 
         return taskiq_result
+
+
+class RedisAsyncClusterResultBackend(AsyncResultBackend[_ReturnType]):
+    """Async result backend based on redis cluster."""
+
+    def __init__(
+        self,
+        redis_url: str,
+        keep_results: bool = True,
+        result_ex_time: Optional[int] = None,
+        result_px_time: Optional[int] = None,
+    ) -> None:
+        """
+        Constructs a new result backend.
+
+        :param redis_url: url to redis cluster.
+        :param keep_results: flag to not remove results from Redis after reading.
+        :param result_ex_time: expire time in seconds for result.
+        :param result_px_time: expire time in milliseconds for result.
+
+        :raises DuplicateExpireTimeSelectedError: if result_ex_time
+            and result_px_time are selected.
+        :raises ExpireTimeMustBeMoreThanZeroError: if result_ex_time
+            and result_px_time are equal zero.
+        """
+        self.redis = RedisCluster.from_url(redis_url)
+        self.keep_results = keep_results
+        self.result_ex_time = result_ex_time
+        self.result_px_time = result_px_time
+
+        unavailable_conditions = any(
+            (
+                self.result_ex_time is not None and self.result_ex_time <= 0,
+                self.result_px_time is not None and self.result_px_time <= 0,
+            ),
+        )
+        if unavailable_conditions:
+            raise ExpireTimeMustBeMoreThanZeroError(
+                "You must select one expire time param and it must be more than zero.",
+            )
+
+        if self.result_ex_time and self.result_px_time:
+            raise DuplicateExpireTimeSelectedError(
+                "Choose either result_ex_time or result_px_time.",
+            )
+
+    async def set_result(
+        self,
+        task_id: str,
+        result: TaskiqResult[_ReturnType],
+    ) -> None:
+        """
+        Sets task result in redis.
+
+        Dumps TaskiqResult instance into the bytes and writes
+        it to redis.
+
+        :param task_id: ID of the task.
+        :param result: TaskiqResult instance.
+        """
+        redis_set_params: Dict[str, Union[str, bytes, int]] = {
+            "name": task_id,
+            "value": pickle.dumps(result),
+        }
+        if self.result_ex_time:
+            redis_set_params["ex"] = self.result_ex_time
+        elif self.result_px_time:
+            redis_set_params["px"] = self.result_px_time
+
+        await self.redis.set(**redis_set_params)  # type: ignore
+
+    async def is_result_ready(self, task_id: str) -> bool:
+        """
+        Returns whether the result is ready.
+
+        :param task_id: ID of the task.
+
+        :returns: True if the result is ready else False.
+        """
+        return bool(await self.redis.exists(task_id))
+
+    async def get_result(
+        self,
+        task_id: str,
+        with_logs: bool = False,
+    ) -> TaskiqResult[_ReturnType]:
+        """
+        Gets result from the task.
+
+        :param task_id: task's id.
+        :param with_logs: if True it will download task's logs.
+        :raises ResultIsMissingError: if there is no result when trying to get it.
+        :return: task's return value.
+        """
+        if self.keep_results:
+            result_value = await self.redis.get(
+                name=task_id,
+            )
+        else:
+            result_value = await self.redis.getdel(
+                name=task_id,
+            )
+
+        if result_value is None:
+            raise ResultIsMissingError
+
+        taskiq_result: TaskiqResult[_ReturnType] = pickle.loads(  # noqa: S301
+            result_value,
+        )
+
+        if not with_logs:
+            taskiq_result.log = None
+
+        return taskiq_result

tests/conftest.py

@@ -26,3 +26,17 @@ def redis_url() -> str:
     :return: URL string.
     """
     return os.environ.get("TEST_REDIS_URL", "redis://localhost")
+
+
+@pytest.fixture
+def redis_cluster_url() -> str:
+    """
+    URL to connect to redis cluster.
+
+    It tries to get it from environ,
+    and return default one if the variable is
+    not set.
+
+    :return: URL string.
+    """
+    return os.environ.get("TEST_REDIS_CLUSTER_URL", "redis://localhost:7000")

tests/test_result_backend.py

@@ -3,7 +3,8 @@
 import pytest
 from taskiq import TaskiqResult
 
-from taskiq_redis import RedisAsyncResultBackend
+from taskiq_redis import (RedisAsyncClusterResultBackend,
+                          RedisAsyncResultBackend)
 from taskiq_redis.exceptions import ResultIsMissingError
 
 
@@ -130,3 +131,128 @@ async def test_keep_results_after_reading(redis_url: str) -> None:
     res2 = await result_backend.get_result(task_id=task_id)
     assert res1 == res2
     await result_backend.shutdown()
+
+
+@pytest.mark.anyio
+async def test_set_result_success_cluster(redis_cluster_url: str) -> None:
+    """
+    Tests that results can be set without errors in cluster mode.
+
+    :param redis_url: redis URL.
+    """
+    result_backend = RedisAsyncClusterResultBackend(  # type: ignore
+        redis_url=redis_cluster_url,
+    )
+    task_id = uuid.uuid4().hex
+    result: "TaskiqResult[int]" = TaskiqResult(
+        is_err=True,
+        log="My Log",
+        return_value=11,
+        execution_time=112.2,
+    )
+    await result_backend.set_result(
+        task_id=task_id,
+        result=result,
+    )
+
+    fetched_result = await result_backend.get_result(
+        task_id=task_id,
+        with_logs=True,
+    )
+    assert fetched_result.log == "My Log"
+    assert fetched_result.return_value == 11
+    assert fetched_result.execution_time == 112.2
+    assert fetched_result.is_err
+    await result_backend.shutdown()
+
+
+@pytest.mark.anyio
+async def test_fetch_without_logs_cluster(redis_cluster_url: str) -> None:
+    """
+    Check if fetching value without logs works fine.
+
+    :param redis_url: redis URL.
+    """
+    result_backend = RedisAsyncClusterResultBackend(  # type: ignore
+        redis_url=redis_cluster_url,
+    )
+    task_id = uuid.uuid4().hex
+    result: "TaskiqResult[int]" = TaskiqResult(
+        is_err=True,
+        log="My Log",
+        return_value=11,
+        execution_time=112.2,
+    )
+    await result_backend.set_result(
+        task_id=task_id,
+        result=result,
+    )
+
+    fetched_result = await result_backend.get_result(
+        task_id=task_id,
+        with_logs=False,
+    )
+    assert fetched_result.log is None
+    assert fetched_result.return_value == 11
+    assert fetched_result.execution_time == 112.2
+    assert fetched_result.is_err
+    await result_backend.shutdown()
+
+
+@pytest.mark.anyio
+async def test_remove_results_after_reading_cluster(redis_cluster_url: str) -> None:
+    """
+    Check if removing results after reading works fine.
+
+    :param redis_url: redis URL.
+    """
+    result_backend = RedisAsyncClusterResultBackend(  # type: ignore
+        redis_url=redis_cluster_url,
+        keep_results=False,
+    )
+    task_id = uuid.uuid4().hex
+    result: "TaskiqResult[int]" = TaskiqResult(
+        is_err=True,
+        log="My Log",
+        return_value=11,
+        execution_time=112.2,
+    )
+    await result_backend.set_result(
+        task_id=task_id,
+        result=result,
+    )
+
+    await result_backend.get_result(task_id=task_id)
+    with pytest.raises(ResultIsMissingError):
+        await result_backend.get_result(task_id=task_id)
+
+    await result_backend.shutdown()
+
+
+@pytest.mark.anyio
+async def test_keep_results_after_reading_cluster(redis_cluster_url: str) -> None:
+    """
+    Check if keeping results after reading works fine.
+
+    :param redis_url: redis URL.
+    """
+    result_backend = RedisAsyncClusterResultBackend(  # type: ignore
+        redis_url=redis_cluster_url,
+        keep_results=True,
+    )
+    task_id = uuid.uuid4().hex
+    result: "TaskiqResult[int]" = TaskiqResult(
+        is_err=True,
+        log="My Log",
+        return_value=11,
+        execution_time=112.2,
+    )
+    await result_backend.set_result(
+        task_id=task_id,
+        result=result,
+    )
+
+    res1 = await result_backend.get_result(task_id=task_id)
+    res2 = await result_backend.get_result(task_id=task_id)
+    assert res1 == res2
+    await result_backend.shutdown()