taskiq-redis 0.5.6__tar.gz → 1.0.0__tar.gz

taskiq_redis-1.0.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022-2024 Pavel Kirilin
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

{taskiq_redis-0.5.6 → taskiq_redis-1.0.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: taskiq-redis
-Version: 0.5.6
+Version: 1.0.0
 Summary: Redis integration for taskiq
 Home-page: https://github.com/taskiq-python/taskiq-redis
 Keywords: taskiq,tasks,distributed,async,redis,result_backend
@@ -16,7 +16,7 @@ Classifier: Programming Language :: Python :: 3.12
 Classifier: Programming Language :: Python :: 3 :: Only
 Classifier: Programming Language :: Python :: 3.8
 Requires-Dist: redis (>=5,<6)
-Requires-Dist: taskiq (>=0.10.3,<1)
+Requires-Dist: taskiq (>=0.11.1,<1)
 Project-URL: Repository, https://github.com/taskiq-python/taskiq-redis
 Description-Content-Type: text/markdown
 

{taskiq_redis-0.5.6 → taskiq_redis-1.0.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "taskiq-redis"
-version = "0.5.6"
+version = "1.0.0"
 description = "Redis integration for taskiq"
 authors = ["taskiq-team <taskiq@norely.com>"]
 readme = "README.md"
@@ -26,7 +26,7 @@ keywords = [
 
 [tool.poetry.dependencies]
 python = "^3.8.1"
-taskiq = ">=0.10.3,<1"
+taskiq = ">=0.11.1,<1"
 redis = "^5"
 
 [tool.poetry.group.dev.dependencies]
@@ -40,7 +40,7 @@ fakeredis = "^2"
 pre-commit = "^2.20.0"
 pytest-xdist = { version = "^2.5.0", extras = ["psutil"] }
 ruff = "^0.1.0"
-types-redis = "^4.6.0.7"
+types-redis = "^4.6.0.20240425"
 
 [tool.mypy]
 strict = true

{taskiq_redis-0.5.6 → taskiq_redis-1.0.0}/taskiq_redis/__init__.py

@@ -2,20 +2,30 @@
 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
+from taskiq_redis.redis_sentinel_broker import (
+    ListQueueSentinelBroker,
+    PubSubSentinelBroker,
+)
 from taskiq_redis.schedule_source import (
     RedisClusterScheduleSource,
     RedisScheduleSource,
+    RedisSentinelScheduleSource,
 )
 
 __all__ = [
     "RedisAsyncClusterResultBackend",
     "RedisAsyncResultBackend",
+    "RedisAsyncSentinelResultBackend",
     "ListQueueBroker",
     "PubSubBroker",
     "ListQueueClusterBroker",
+    "ListQueueSentinelBroker",
+    "PubSubSentinelBroker",
     "RedisScheduleSource",
     "RedisClusterScheduleSource",
+    "RedisSentinelScheduleSource",
 ]

{taskiq_redis-0.5.6 → taskiq_redis-1.0.0}/taskiq_redis/redis_backend.py

@@ -1,10 +1,25 @@
-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 redis.asyncio.connection import Connection
 from taskiq import AsyncResultBackend
 from taskiq.abc.result_backend import TaskiqResult
+from taskiq.abc.serializer import TaskiqSerializer
+from taskiq.compat import model_dump, model_validate
+from taskiq.serializers import PickleSerializer
 
 from taskiq_redis.exceptions import (
     DuplicateExpireTimeSelectedError,
@@ -12,6 +27,18 @@ from taskiq_redis.exceptions import (
     ResultIsMissingError,
 )
 
+if sys.version_info >= (3, 10):
+    from typing import TypeAlias
+else:
+    from typing_extensions import TypeAlias
+
+if TYPE_CHECKING:
+    _Redis: TypeAlias = Redis[bytes]
+    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool[Connection]
+else:
+    _Redis: TypeAlias = Redis
+    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool
+
 _ReturnType = TypeVar("_ReturnType")
 
 
@@ -25,6 +52,7 @@ class RedisAsyncResultBackend(AsyncResultBackend[_ReturnType]):
         result_ex_time: Optional[int] = None,
         result_px_time: Optional[int] = None,
         max_connection_pool_size: Optional[int] = None,
+        serializer: Optional[TaskiqSerializer] = None,
         **connection_kwargs: Any,
     ) -> None:
         """
@@ -42,11 +70,12 @@ class RedisAsyncResultBackend(AsyncResultBackend[_ReturnType]):
         :raises ExpireTimeMustBeMoreThanZeroError: if result_ex_time
             and result_px_time are equal zero.
         """
-        self.redis_pool = BlockingConnectionPool.from_url(
+        self.redis_pool: _BlockingConnectionPool = BlockingConnectionPool.from_url(
             url=redis_url,
             max_connections=max_connection_pool_size,
             **connection_kwargs,
         )
+        self.serializer = serializer or PickleSerializer()
         self.keep_results = keep_results
         self.result_ex_time = result_ex_time
         self.result_px_time = result_px_time
@@ -86,9 +115,9 @@ class RedisAsyncResultBackend(AsyncResultBackend[_ReturnType]):
         :param task_id: ID of the task.
         :param result: TaskiqResult instance.
         """
-        redis_set_params: Dict[str, Union[str, bytes, int]] = {
+        redis_set_params: Dict[str, Union[str, int, bytes]] = {
             "name": task_id,
-            "value": pickle.dumps(result),
+            "value": self.serializer.dumpb(model_dump(result)),
         }
         if self.result_ex_time:
             redis_set_params["ex"] = self.result_ex_time
@@ -135,8 +164,9 @@ class RedisAsyncResultBackend(AsyncResultBackend[_ReturnType]):
         if result_value is None:
             raise ResultIsMissingError
 
-        taskiq_result: TaskiqResult[_ReturnType] = pickle.loads(  # noqa: S301
-            result_value,
+        taskiq_result = model_validate(
+            TaskiqResult[_ReturnType],
+            self.serializer.loadb(result_value),
         )
 
         if not with_logs:
@@ -154,6 +184,7 @@ class RedisAsyncClusterResultBackend(AsyncResultBackend[_ReturnType]):
         keep_results: bool = True,
         result_ex_time: Optional[int] = None,
         result_px_time: Optional[int] = None,
+        serializer: Optional[TaskiqSerializer] = None,
         **connection_kwargs: Any,
     ) -> None:
         """
@@ -174,6 +205,7 @@ class RedisAsyncClusterResultBackend(AsyncResultBackend[_ReturnType]):
             redis_url,
             **connection_kwargs,
         )
+        self.serializer = serializer or PickleSerializer()
         self.keep_results = keep_results
         self.result_ex_time = result_ex_time
         self.result_px_time = result_px_time
@@ -215,7 +247,7 @@ class RedisAsyncClusterResultBackend(AsyncResultBackend[_ReturnType]):
         """
         redis_set_params: Dict[str, Union[str, bytes, int]] = {
             "name": task_id,
-            "value": pickle.dumps(result),
+            "value": self.serializer.dumpb(model_dump(result)),
         }
         if self.result_ex_time:
             redis_set_params["ex"] = self.result_ex_time
@@ -259,11 +291,155 @@ class RedisAsyncClusterResultBackend(AsyncResultBackend[_ReturnType]):
         if result_value is None:
             raise ResultIsMissingError
 
-        taskiq_result: TaskiqResult[_ReturnType] = pickle.loads(  # noqa: S301
-            result_value,
+        taskiq_result: TaskiqResult[_ReturnType] = model_validate(
+            TaskiqResult[_ReturnType],
+            self.serializer.loadb(result_value),
+        )
+
+        if not with_logs:
+            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,
+        serializer: Optional[TaskiqSerializer] = 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.serializer = serializer or PickleSerializer()
+        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": self.serializer.dumpb(model_dump(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 = model_validate(
+            TaskiqResult[_ReturnType],
+            self.serializer.loadb(result_value),
         )
 
         if not with_logs:
             taskiq_result.log = None
 
         return taskiq_result
+
+    async def shutdown(self) -> None:
+        """Shutdown sentinel connections."""
+        for sentinel in self.sentinel.sentinels:
+            await sentinel.aclose()  # type: ignore[attr-defined]

{taskiq_redis-0.5.6 → taskiq_redis-1.0.0}/taskiq_redis/redis_broker.py

@@ -1,7 +1,8 @@
+import sys
 from logging import getLogger
-from typing import Any, AsyncGenerator, Callable, Optional, TypeVar
+from typing import TYPE_CHECKING, Any, AsyncGenerator, Callable, Optional, TypeVar
 
-from redis.asyncio import BlockingConnectionPool, ConnectionPool, Redis
+from redis.asyncio import BlockingConnectionPool, Connection, Redis
 from taskiq.abc.broker import AsyncBroker
 from taskiq.abc.result_backend import AsyncResultBackend
 from taskiq.message import BrokerMessage
@@ -10,6 +11,16 @@ _T = TypeVar("_T")
 
 logger = getLogger("taskiq.redis_broker")
 
+if sys.version_info >= (3, 10):
+    from typing import TypeAlias
+else:
+    from typing_extensions import TypeAlias
+
+if TYPE_CHECKING:
+    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool[Connection]
+else:
+    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool
+
 
 class BaseRedisBroker(AsyncBroker):
     """Base broker that works with Redis."""
@@ -40,7 +51,7 @@ class BaseRedisBroker(AsyncBroker):
             task_id_generator=task_id_generator,
         )
 
-        self.connection_pool: ConnectionPool = BlockingConnectionPool.from_url(
+        self.connection_pool: _BlockingConnectionPool = BlockingConnectionPool.from_url(
             url=url,
             max_connections=max_connection_pool_size,
             **connection_kwargs,

taskiq_redis-1.0.0/taskiq_redis/redis_sentinel_broker.py

@@ -0,0 +1,132 @@
+import sys
+from contextlib import asynccontextmanager
+from logging import getLogger
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncGenerator,
+    AsyncIterator,
+    Callable,
+    List,
+    Optional,
+    Tuple,
+    TypeVar,
+)
+
+from redis.asyncio import Redis, Sentinel
+from taskiq import AsyncResultBackend, BrokerMessage
+from taskiq.abc.broker import AsyncBroker
+
+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
+
+_T = TypeVar("_T")
+
+logger = getLogger("taskiq.redis_sentinel_broker")
+
+
+class BaseSentinelBroker(AsyncBroker):
+    """Base broker that works with Sentinel."""
+
+    def __init__(
+        self,
+        sentinels: List[Tuple[str, int]],
+        master_name: str,
+        result_backend: Optional[AsyncResultBackend[_T]] = None,
+        task_id_generator: Optional[Callable[[], str]] = None,
+        queue_name: str = "taskiq",
+        min_other_sentinels: int = 0,
+        sentinel_kwargs: Optional[Any] = None,
+        **connection_kwargs: Any,
+    ) -> None:
+        super().__init__(
+            result_backend=result_backend,
+            task_id_generator=task_id_generator,
+        )
+
+        self.sentinel = Sentinel(
+            sentinels=sentinels,
+            min_other_sentinels=min_other_sentinels,
+            sentinel_kwargs=sentinel_kwargs,
+            **connection_kwargs,
+        )
+        self.master_name = master_name
+        self.queue_name = queue_name
+
+    @asynccontextmanager
+    async def _acquire_master_conn(self) -> AsyncIterator[_Redis]:
+        async with self.sentinel.master_for(self.master_name) as redis_conn:
+            yield redis_conn
+
+
+class PubSubSentinelBroker(BaseSentinelBroker):
+    """Broker that works with Sentinel and broadcasts tasks to all workers."""
+
+    async def kick(self, message: BrokerMessage) -> None:
+        """
+        Publish message over PUBSUB channel.
+
+        :param message: message to send.
+        """
+        queue_name = message.labels.get("queue_name") or self.queue_name
+        async with self._acquire_master_conn() as redis_conn:
+            await redis_conn.publish(queue_name, message.message)
+
+    async def listen(self) -> AsyncGenerator[bytes, None]:
+        """
+        Listen redis queue for new messages.
+
+        This function listens to the pubsub channel
+        and yields all messages with proper types.
+
+        :yields: broker messages.
+        """
+        async with self._acquire_master_conn() as redis_conn:
+            redis_pubsub_channel = redis_conn.pubsub()
+            await redis_pubsub_channel.subscribe(self.queue_name)
+            async for message in redis_pubsub_channel.listen():
+                if not message:
+                    continue
+                if message["type"] != "message":
+                    logger.debug("Received non-message from redis: %s", message)
+                    continue
+                yield message["data"]
+
+
+class ListQueueSentinelBroker(BaseSentinelBroker):
+    """Broker that works with Sentinel and distributes tasks between workers."""
+
+    async def kick(self, message: BrokerMessage) -> None:
+        """
+        Put a message in a list.
+
+        This method appends a message to the list of all messages.
+
+        :param message: message to append.
+        """
+        queue_name = message.labels.get("queue_name") or self.queue_name
+        async with self._acquire_master_conn() as redis_conn:
+            await redis_conn.lpush(queue_name, message.message)
+
+    async def listen(self) -> AsyncGenerator[bytes, None]:
+        """
+        Listen redis queue for new messages.
+
+        This function listens to the queue
+        and yields new messages if they have BrokerMessage type.
+
+        :yields: broker messages.
+        """
+        redis_brpop_data_position = 1
+        async with self._acquire_master_conn() as redis_conn:
+            while True:
+                yield (await redis_conn.brpop(self.queue_name))[
+                    redis_brpop_data_position
+                ]

{taskiq_redis-0.5.6 → taskiq_redis-1.0.0}/taskiq_redis/schedule_source.py

@@ -1,12 +1,31 @@
-from typing import Any, List, Optional
+import sys
+from contextlib import asynccontextmanager
+from typing import TYPE_CHECKING, Any, AsyncIterator, List, Optional, Tuple
 
-from redis.asyncio import BlockingConnectionPool, ConnectionPool, Redis, RedisCluster
+from redis.asyncio import (
+    BlockingConnectionPool,
+    Connection,
+    Redis,
+    RedisCluster,
+    Sentinel,
+)
 from taskiq import ScheduleSource
 from taskiq.abc.serializer import TaskiqSerializer
 from taskiq.compat import model_dump, model_validate
 from taskiq.scheduler.scheduled_task import ScheduledTask
+from taskiq.serializers import PickleSerializer
 
-from taskiq_redis.serializer import PickleSerializer
+if sys.version_info >= (3, 10):
+    from typing import TypeAlias
+else:
+    from typing_extensions import TypeAlias
+
+if TYPE_CHECKING:
+    _Redis: TypeAlias = Redis[bytes]
+    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool[Connection]
+else:
+    _Redis: TypeAlias = Redis
+    _BlockingConnectionPool: TypeAlias = BlockingConnectionPool
 
 
 class RedisScheduleSource(ScheduleSource):
@@ -35,7 +54,7 @@ class RedisScheduleSource(ScheduleSource):
         **connection_kwargs: Any,
     ) -> None:
         self.prefix = prefix
-        self.connection_pool: ConnectionPool = BlockingConnectionPool.from_url(
+        self.connection_pool: _BlockingConnectionPool = BlockingConnectionPool.from_url(
             url=url,
             max_connections=max_connection_pool_size,
             **connection_kwargs,
@@ -117,7 +136,6 @@ class RedisClusterScheduleSource(ScheduleSource):
         self,
         url: str,
         prefix: str = "schedule",
-        buffer_size: int = 50,
         serializer: Optional[TaskiqSerializer] = None,
         **connection_kwargs: Any,
     ) -> None:
@@ -126,7 +144,6 @@ class RedisClusterScheduleSource(ScheduleSource):
             url,
             **connection_kwargs,
         )
-        self.buffer_size = buffer_size
         if serializer is None:
             serializer = PickleSerializer()
         self.serializer = serializer
@@ -156,14 +173,107 @@ class RedisClusterScheduleSource(ScheduleSource):
         :return: list of schedules.
         """
         schedules = []
-        buffer = []
         async for key in self.redis.scan_iter(f"{self.prefix}:*"):  # type: ignore[attr-defined]
-            buffer.append(key)
-            if len(buffer) >= self.buffer_size:
-                schedules.extend(await self.redis.mget(buffer))  # type: ignore[attr-defined]
-                buffer = []
-        if buffer:
-            schedules.extend(await self.redis.mget(buffer))  # type: ignore[attr-defined]
+            raw_schedule = await self.redis.get(key)  # type: ignore[attr-defined]
+            parsed_schedule = model_validate(
+                ScheduledTask,
+                self.serializer.loadb(raw_schedule),
+            )
+            schedules.append(parsed_schedule)
+        return schedules
+
+    async def post_send(self, task: ScheduledTask) -> None:
+        """Delete a task after it's completed."""
+        if task.time is not None:
+            await self.delete_schedule(task.schedule_id)
+
+    async def shutdown(self) -> None:
+        """Shut down the schedule source."""
+        await self.redis.aclose()  # type: ignore[attr-defined]
+
+
+class RedisSentinelScheduleSource(ScheduleSource):
+    """
+    Source of schedules for redis cluster.
+
+    This class allows you to store schedules in redis.
+    Also it supports dynamic schedules.
+
+    :param sentinels: list of sentinel host and ports pairs.
+    :param master_name: sentinel master name.
+    :param prefix: prefix for redis schedule keys.
+    :param buffer_size: buffer size for redis scan.
+        This is how many keys will be fetched at once.
+    :param max_connection_pool_size: maximum number of connections in pool.
+    :param serializer: serializer for data.
+    :param connection_kwargs: additional arguments for RedisCluster.
+    """
+
+    def __init__(
+        self,
+        sentinels: List[Tuple[str, int]],
+        master_name: str,
+        prefix: str = "schedule",
+        buffer_size: int = 50,
+        serializer: Optional[TaskiqSerializer] = None,
+        min_other_sentinels: int = 0,
+        sentinel_kwargs: Optional[Any] = None,
+        **connection_kwargs: Any,
+    ) -> None:
+        self.prefix = prefix
+        self.sentinel = Sentinel(
+            sentinels=sentinels,
+            min_other_sentinels=min_other_sentinels,
+            sentinel_kwargs=sentinel_kwargs,
+            **connection_kwargs,
+        )
+        self.master_name = master_name
+        self.buffer_size = buffer_size
+        if serializer is None:
+            serializer = PickleSerializer()
+        self.serializer = serializer
+
+    @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 delete_schedule(self, schedule_id: str) -> None:
+        """Remove schedule by id."""
+        async with self._acquire_master_conn() as redis:
+            await redis.delete(f"{self.prefix}:{schedule_id}")
+
+    async def add_schedule(self, schedule: ScheduledTask) -> None:
+        """
+        Add schedule to redis.
+
+        :param schedule: schedule to add.
+        :param schedule_id: schedule id.
+        """
+        async with self._acquire_master_conn() as redis:
+            await redis.set(
+                f"{self.prefix}:{schedule.schedule_id}",
+                self.serializer.dumpb(model_dump(schedule)),
+            )
+
+    async def get_schedules(self) -> List[ScheduledTask]:
+        """
+        Get all schedules from redis.
+
+        This method is used by scheduler to get all schedules.
+
+        :return: list of schedules.
+        """
+        schedules = []
+        async with self._acquire_master_conn() as redis:
+            buffer = []
+            async for key in redis.scan_iter(f"{self.prefix}:*"):
+                buffer.append(key)
+                if len(buffer) >= self.buffer_size:
+                    schedules.extend(await redis.mget(buffer))
+                    buffer = []
+            if buffer:
+                schedules.extend(await redis.mget(buffer))
         return [
             model_validate(ScheduledTask, self.serializer.loadb(schedule))
             for schedule in schedules
@@ -174,3 +284,8 @@ class RedisClusterScheduleSource(ScheduleSource):
         """Delete a task after it's completed."""
         if task.time is not None:
             await self.delete_schedule(task.schedule_id)
+
+    async def shutdown(self) -> None:
+        """Shut down the schedule source."""
+        for sentinel in self.sentinel.sentinels:
+            await sentinel.aclose()  # type: ignore[attr-defined]

taskiq_redis-0.5.6/taskiq_redis/serializer.py

@@ -1,16 +0,0 @@
-import pickle
-from typing import Any
-
-from taskiq.abc.serializer import TaskiqSerializer
-
-
-class PickleSerializer(TaskiqSerializer):
-    """Serializer that uses pickle."""
-
-    def dumpb(self, value: Any) -> bytes:
-        """Dumps value to bytes."""
-        return pickle.dumps(value)
-
-    def loadb(self, value: bytes) -> Any:
-        """Loads value from bytes."""
-        return pickle.loads(value)  # noqa: S301