Add cluster option to result backend

Author: stinovlas

.github/workflows/test.yml

@@ -42,12 +42,21 @@ jobs:
           --health-retries=30
         ports:
           - 6379:6379
+      redis-cluster:
+        image: bitnami/redis-cluster:6.2.5
+        env:
+          ALLOW_EMPTY_PASSWORD: "yes"
+          REDIS_NODES: "localhost"
+        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
+      - name: Set up single-node Redis cluster
+        run: redis-cli -h redis-cluster -p 7000 --cluster-yes CLUSTER ADDSLOTSRANGE 0 16383
       - name: Set up Python
         uses: actions/setup-python@v2
         with:

taskiq_redis/redis_backend.py

@@ -1,7 +1,7 @@
 import pickle
-from typing import Dict, Optional, TypeVar, Union
+from typing import Dict, Optional, Type, TypeVar, Union
 
-from redis.asyncio import ConnectionPool, Redis
+from redis.asyncio import Redis, RedisCluster
 from taskiq import AsyncResultBackend
 from taskiq.abc.result_backend import TaskiqResult
 
@@ -23,11 +23,15 @@ def __init__(
         keep_results: bool = True,
         result_ex_time: Optional[int] = None,
         result_px_time: Optional[int] = None,
+        *,
+        redis_cls: Union[Type[Redis], Type[RedisCluster], None] = None,
     ) -> None:
         """
         Constructs a new result backend.
 
         :param redis_url: url to redis.
+        :param redis_cls: async redis class, should be either redis.asyncio.Redis
+            or redis.asyncio.RedisCluster.
         :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.
@@ -37,7 +41,10 @@ def __init__(
         :raises ExpireTimeMustBeMoreThanZeroError: if result_ex_time
             and result_px_time are equal zero.
         """
-        self.redis_pool = ConnectionPool.from_url(redis_url)
+        if redis_cls is None:
+            redis_cls = Redis
+
+        self.redis = redis_cls.from_url(redis_url)
         self.keep_results = keep_results
         self.result_ex_time = result_ex_time
         self.result_px_time = result_px_time
@@ -58,11 +65,6 @@ def __init__(
                 "Choose either result_ex_time or result_px_time.",
             )
 
-    async def shutdown(self) -> None:
-        """Closes redis connection."""
-        await self.redis_pool.disconnect()
-        await super().shutdown()
-
     async def set_result(
         self,
         task_id: str,
@@ -86,8 +88,7 @@ async def set_result(
         elif self.result_px_time:
             redis_set_params["px"] = self.result_px_time
 
-        async with Redis(connection_pool=self.redis_pool) as redis:
-            await redis.set(**redis_set_params)  # type: ignore
+        await self.redis.set(**redis_set_params)  # type: ignore
 
     async def is_result_ready(self, task_id: str) -> bool:
         """
@@ -97,8 +98,7 @@ async def is_result_ready(self, task_id: str) -> bool:
 
         :returns: True if the result is ready else False.
         """
-        async with Redis(connection_pool=self.redis_pool) as redis:
-            return bool(await redis.exists(task_id))
+        return bool(await self.redis.exists(task_id))
 
     async def get_result(
         self,
@@ -113,15 +113,14 @@ async def get_result(
         :raises ResultIsMissingError: if there is no result when trying to get it.
         :return: task's return value.
         """
-        async with Redis(connection_pool=self.redis_pool) as redis:
-            if self.keep_results:
-                result_value = await redis.get(
-                    name=task_id,
-                )
-            else:
-                result_value = await redis.getdel(
-                    name=task_id,
-                )
+        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

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

@@ -1,6 +1,7 @@
 import uuid
 
 import pytest
+from redis.asyncio import RedisCluster
 from taskiq import TaskiqResult
 
 from taskiq_redis import RedisAsyncResultBackend
@@ -130,3 +131,37 @@ 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 = RedisAsyncResultBackend(  # type: ignore
+        redis_url=redis_cluster_url,
+        redis_cls=RedisCluster,
+    )
+    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()