feat: Added RedisSentinelAsyncResultBackend.

DABND19 · DABND19 · commit a9b7f9d9f20d · 2024-04-14T17:15:46.000+03:00
diff --git a/taskiq_redis/__init__.py b/taskiq_redis/__init__.py
@@ -2,6 +2,7 @@
 from taskiq_redis.redis_backend import (
     RedisAsyncClusterResultBackend,
     RedisAsyncResultBackend,
+    RedisAsyncSentinelResultBackend,
 )
 from taskiq_redis.redis_broker import ListQueueBroker, PubSubBroker
 from taskiq_redis.redis_cluster_broker import ListQueueClusterBroker
@@ -14,6 +15,7 @@
 __all__ = [
     "RedisAsyncClusterResultBackend",
     "RedisAsyncResultBackend",
+    "RedisAsyncSentinelResultBackend",
     "ListQueueBroker",
     "PubSubBroker",
     "ListQueueClusterBroker",
diff --git a/taskiq_redis/redis_backend.py b/taskiq_redis/redis_backend.py
@@ -1,7 +1,19 @@
 import pickle
-from typing import Any, Dict, Optional, TypeVar, Union
+import sys
+from contextlib import asynccontextmanager
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncIterator,
+    Dict,
+    List,
+    Optional,
+    Tuple,
+    TypeVar,
+    Union,
+)
 
-from redis.asyncio import BlockingConnectionPool, Redis
+from redis.asyncio import BlockingConnectionPool, Redis, Sentinel
 from redis.asyncio.cluster import RedisCluster
 from taskiq import AsyncResultBackend
 from taskiq.abc.result_backend import TaskiqResult
@@ -12,6 +24,16 @@
     ResultIsMissingError,
 )
 
+if sys.version_info >= (3, 10):
+    from typing import TypeAlias
+else:
+    from typing_extensions import TypeAlias
+
+if TYPE_CHECKING:
+    _Redis: TypeAlias = Redis[bytes]
+else:
+    _Redis: TypeAlias = Redis
+
 _ReturnType = TypeVar("_ReturnType")
 
 
@@ -267,3 +289,138 @@ async def get_result(
             taskiq_result.log = None
 
         return taskiq_result
+
+
+class RedisAsyncSentinelResultBackend(AsyncResultBackend[_ReturnType]):
+    """Async result based on redis sentinel."""
+
+    def __init__(
+        self,
+        sentinels: List[Tuple[str, int]],
+        master_name: str,
+        keep_results: bool = True,
+        result_ex_time: Optional[int] = None,
+        result_px_time: Optional[int] = None,
+        min_other_sentinels: int = 0,
+        sentinel_kwargs: Optional[Any] = None,
+        **connection_kwargs: Any,
+    ) -> None:
+        """
+        Constructs a new result backend.
+
+        :param sentinels: list of sentinel host and ports pairs.
+        :param master_name: sentinel master name.
+        :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.
+        :param max_connection_pool_size: maximum number of connections in pool.
+        :param connection_kwargs: additional arguments for redis BlockingConnectionPool.
+
+        :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.sentinel = Sentinel(
+            sentinels=sentinels,
+            min_other_sentinels=min_other_sentinels,
+            sentinel_kwargs=sentinel_kwargs,
+            **connection_kwargs,
+        )
+        self.master_name = master_name
+        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.",
+            )
+
+    @asynccontextmanager
+    async def _acquire_master_conn(self) -> AsyncIterator[_Redis]:
+        async with self.sentinel.master_for(self.master_name) as redis_conn:
+            yield redis_conn
+
+    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
+
+        async with self._acquire_master_conn() as redis:
+            await 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.
+        """
+        async with self._acquire_master_conn() as redis:
+            return bool(await 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.
+        """
+        async with self._acquire_master_conn() as redis:
+            if self.keep_results:
+                result_value = await redis.get(
+                    name=task_id,
+                )
+            else:
+                result_value = await 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
diff --git a/tests/test_result_backend.py b/tests/test_result_backend.py
@@ -1,10 +1,15 @@
 import asyncio
 import uuid
+from typing import List, Tuple
 
 import pytest
 from taskiq import TaskiqResult
 
-from taskiq_redis import RedisAsyncClusterResultBackend, RedisAsyncResultBackend
+from taskiq_redis import (
+    RedisAsyncClusterResultBackend,
+    RedisAsyncResultBackend,
+    RedisAsyncSentinelResultBackend,
+)
 from taskiq_redis.exceptions import ResultIsMissingError
 
 
@@ -288,3 +293,148 @@ async def test_keep_results_after_reading_cluster(redis_cluster_url: str) -> Non
     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_sentinel(
+    redis_sentinels: List[Tuple[str, int]],
+    redis_sentinel_master_name: str,
+) -> None:
+    """
+    Tests that results can be set without errors in cluster mode.
+
+    :param redis_sentinels: list of host and port pairs.
+    :param redis_sentinel_master_name: redis sentinel master name string.
+    """
+    result_backend = RedisAsyncSentinelResultBackend(  # type: ignore
+        sentinels=redis_sentinels,
+        master_name=redis_sentinel_master_name,
+    )
+    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_sentinel(
+    redis_sentinels: List[Tuple[str, int]],
+    redis_sentinel_master_name: str,
+) -> None:
+    """
+    Check if fetching value without logs works fine.
+
+    :param redis_sentinels: list of host and port pairs.
+    :param redis_sentinel_master_name: redis sentinel master name string.
+    """
+    result_backend = RedisAsyncSentinelResultBackend(  # type: ignore
+        sentinels=redis_sentinels,
+        master_name=redis_sentinel_master_name,
+    )
+    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_sentinel(
+    redis_sentinels: List[Tuple[str, int]],
+    redis_sentinel_master_name: str,
+) -> None:
+    """
+    Check if removing results after reading works fine.
+
+    :param redis_sentinels: list of host and port pairs.
+    :param redis_sentinel_master_name: redis sentinel master name string.
+    """
+    result_backend = RedisAsyncSentinelResultBackend(  # type: ignore
+        sentinels=redis_sentinels,
+        master_name=redis_sentinel_master_name,
+        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_sentinel(
+    redis_sentinels: List[Tuple[str, int]],
+    redis_sentinel_master_name: str,
+) -> None:
+    """
+    Check if keeping results after reading works fine.
+
+    :param redis_sentinels: list of host and port pairs.
+    :param redis_sentinel_master_name: redis sentinel master name string.
+    """
+    result_backend = RedisAsyncSentinelResultBackend(  # type: ignore
+        sentinels=redis_sentinels,
+        master_name=redis_sentinel_master_name,
+        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()
