taskiq-redis 0.5.4__py3-none-any.whl → 0.5.6__py3-none-any.whl

taskiq_redis/__init__.py

@@ -5,7 +5,10 @@ from taskiq_redis.redis_backend import (
 )
 from taskiq_redis.redis_broker import ListQueueBroker, PubSubBroker
 from taskiq_redis.redis_cluster_broker import ListQueueClusterBroker
-from taskiq_redis.schedule_source import RedisScheduleSource
+from taskiq_redis.schedule_source import (
+    RedisClusterScheduleSource,
+    RedisScheduleSource,
+)
 
 __all__ = [
     "RedisAsyncClusterResultBackend",
@@ -14,4 +17,5 @@ __all__ = [
     "PubSubBroker",
     "ListQueueClusterBroker",
     "RedisScheduleSource",
+    "RedisClusterScheduleSource",
 ]

taskiq_redis/redis_backend.py

@@ -1,7 +1,7 @@
 import pickle
-from typing import Dict, Optional, TypeVar, Union
+from typing import Any, Dict, Optional, TypeVar, Union
 
-from redis.asyncio import ConnectionPool, Redis
+from redis.asyncio import BlockingConnectionPool, Redis
 from redis.asyncio.cluster import RedisCluster
 from taskiq import AsyncResultBackend
 from taskiq.abc.result_backend import TaskiqResult
@@ -24,6 +24,8 @@ class RedisAsyncResultBackend(AsyncResultBackend[_ReturnType]):
         keep_results: bool = True,
         result_ex_time: Optional[int] = None,
         result_px_time: Optional[int] = None,
+        max_connection_pool_size: Optional[int] = None,
+        **connection_kwargs: Any,
     ) -> None:
         """
         Constructs a new result backend.
@@ -32,13 +34,19 @@ class RedisAsyncResultBackend(AsyncResultBackend[_ReturnType]):
         :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.redis_pool = ConnectionPool.from_url(redis_url)
+        self.redis_pool = BlockingConnectionPool.from_url(
+            url=redis_url,
+            max_connections=max_connection_pool_size,
+            **connection_kwargs,
+        )
         self.keep_results = keep_results
         self.result_ex_time = result_ex_time
         self.result_px_time = result_px_time
@@ -146,6 +154,7 @@ class RedisAsyncClusterResultBackend(AsyncResultBackend[_ReturnType]):
         keep_results: bool = True,
         result_ex_time: Optional[int] = None,
         result_px_time: Optional[int] = None,
+        **connection_kwargs: Any,
     ) -> None:
         """
         Constructs a new result backend.
@@ -154,13 +163,17 @@ class RedisAsyncClusterResultBackend(AsyncResultBackend[_ReturnType]):
         :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 connection_kwargs: additional arguments for RedisCluster.
 
         :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[bytes] = RedisCluster.from_url(redis_url)
+        self.redis: RedisCluster[bytes] = RedisCluster.from_url(
+            redis_url,
+            **connection_kwargs,
+        )
         self.keep_results = keep_results
         self.result_ex_time = result_ex_time
         self.result_px_time = result_px_time

taskiq_redis/redis_broker.py

@@ -1,7 +1,7 @@
 from logging import getLogger
 from typing import Any, AsyncGenerator, Callable, Optional, TypeVar
 
-from redis.asyncio import ConnectionPool, Redis
+from redis.asyncio import BlockingConnectionPool, ConnectionPool, Redis
 from taskiq.abc.broker import AsyncBroker
 from taskiq.abc.result_backend import AsyncResultBackend
 from taskiq.message import BrokerMessage
@@ -31,14 +31,16 @@ class BaseRedisBroker(AsyncBroker):
         :param result_backend: custom result backend.
         :param queue_name: name for a list in redis.
         :param max_connection_pool_size: maximum number of connections in pool.
-        :param connection_kwargs: additional arguments for aio-redis ConnectionPool.
+            Each worker opens its own connection. Therefore this value has to be
+            at least number of workers + 1.
+        :param connection_kwargs: additional arguments for redis BlockingConnectionPool.
         """
         super().__init__(
             result_backend=result_backend,
             task_id_generator=task_id_generator,
         )
 
-        self.connection_pool: ConnectionPool = ConnectionPool.from_url(
+        self.connection_pool: ConnectionPool = BlockingConnectionPool.from_url(
             url=url,
             max_connections=max_connection_pool_size,
             **connection_kwargs,
@@ -60,8 +62,9 @@ class PubSubBroker(BaseRedisBroker):
 
         :param message: message to send.
         """
+        queue_name = message.labels.get("queue_name") or self.queue_name
         async with Redis(connection_pool=self.connection_pool) as redis_conn:
-            await redis_conn.publish(self.queue_name, message.message)
+            await redis_conn.publish(queue_name, message.message)
 
     async def listen(self) -> AsyncGenerator[bytes, None]:
         """
@@ -95,8 +98,9 @@ class ListQueueBroker(BaseRedisBroker):
 
         :param message: message to append.
         """
+        queue_name = message.labels.get("queue_name") or self.queue_name
         async with Redis(connection_pool=self.connection_pool) as redis_conn:
-            await redis_conn.lpush(self.queue_name, message.message)
+            await redis_conn.lpush(queue_name, message.message)
 
     async def listen(self) -> AsyncGenerator[bytes, None]:
         """

taskiq_redis/schedule_source.py

@@ -1,6 +1,6 @@
 from typing import Any, List, Optional
 
-from redis.asyncio import ConnectionPool, Redis
+from redis.asyncio import BlockingConnectionPool, ConnectionPool, Redis, RedisCluster
 from taskiq import ScheduleSource
 from taskiq.abc.serializer import TaskiqSerializer
 from taskiq.compat import model_dump, model_validate
@@ -22,7 +22,7 @@ class RedisScheduleSource(ScheduleSource):
         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 aio-redis ConnectionPool.
+    :param connection_kwargs: additional arguments for redis BlockingConnectionPool.
     """
 
     def __init__(
@@ -35,7 +35,7 @@ class RedisScheduleSource(ScheduleSource):
         **connection_kwargs: Any,
     ) -> None:
         self.prefix = prefix
-        self.connection_pool: ConnectionPool = ConnectionPool.from_url(
+        self.connection_pool: ConnectionPool = BlockingConnectionPool.from_url(
             url=url,
             max_connections=max_connection_pool_size,
             **connection_kwargs,
@@ -95,3 +95,82 @@ class RedisScheduleSource(ScheduleSource):
     async def shutdown(self) -> None:
         """Shut down the schedule source."""
         await self.connection_pool.disconnect()
+
+
+class RedisClusterScheduleSource(ScheduleSource):
+    """
+    Source of schedules for redis cluster.
+
+    This class allows you to store schedules in redis.
+    Also it supports dynamic schedules.
+
+    :param url: url to redis cluster.
+    :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,
+        url: str,
+        prefix: str = "schedule",
+        buffer_size: int = 50,
+        serializer: Optional[TaskiqSerializer] = None,
+        **connection_kwargs: Any,
+    ) -> None:
+        self.prefix = prefix
+        self.redis: RedisCluster[bytes] = RedisCluster.from_url(
+            url,
+            **connection_kwargs,
+        )
+        self.buffer_size = buffer_size
+        if serializer is None:
+            serializer = PickleSerializer()
+        self.serializer = serializer
+
+    async def delete_schedule(self, schedule_id: str) -> None:
+        """Remove schedule by id."""
+        await self.redis.delete(f"{self.prefix}:{schedule_id}")  # type: ignore[attr-defined]
+
+    async def add_schedule(self, schedule: ScheduledTask) -> None:
+        """
+        Add schedule to redis.
+
+        :param schedule: schedule to add.
+        :param schedule_id: schedule id.
+        """
+        await self.redis.set(  # type: ignore[attr-defined]
+            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 = []
+        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]
+        return [
+            model_validate(ScheduledTask, self.serializer.loadb(schedule))
+            for schedule in schedules
+            if schedule
+        ]
+
+    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)

{taskiq_redis-0.5.4.dist-info → taskiq_redis-0.5.6.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: taskiq-redis
-Version: 0.5.4
+Version: 0.5.6
 Summary: Redis integration for taskiq
 Home-page: https://github.com/taskiq-python/taskiq-redis
 Keywords: taskiq,tasks,distributed,async,redis,result_backend
@@ -93,6 +93,9 @@ Brokers parameters:
 * `result_backend` - custom result backend.
 * `queue_name` - name of the pub/sub channel in redis.
 * `max_connection_pool_size` - maximum number of connections in pool.
+* Any other keyword arguments are passed to `redis.asyncio.BlockingConnectionPool`.
+  Notably, you can use `timeout` to set custom timeout in seconds for reconnects
+  (or set it to `None` to try reconnects indefinitely).
 
 ## RedisAsyncResultBackend configuration
 
@@ -101,6 +104,9 @@ RedisAsyncResultBackend parameters:
 * `keep_results` - flag to not remove results from Redis after reading.
 * `result_ex_time` - expire time in seconds (by default - not specified)
 * `result_px_time` - expire time in milliseconds (by default - not specified)
+* Any other keyword arguments are passed to `redis.asyncio.BlockingConnectionPool`.
+  Notably, you can use `timeout` to set custom timeout in seconds for reconnects
+  (or set it to `None` to try reconnects indefinitely).
 > IMPORTANT: **It is highly recommended to use expire time ​​in RedisAsyncResultBackend**
 > If you want to add expiration, either `result_ex_time` or `result_px_time` must be set.
 >```python

taskiq_redis-0.5.6.dist-info/RECORD

@@ -0,0 +1,11 @@
+taskiq_redis/__init__.py,sha256=fMdXYxulcaKur66UUlmqAQf_q24jT5UHDYsMYP6J4fw,602
+taskiq_redis/exceptions.py,sha256=eS4bfZVAjyMsnFs3IF74uYwO1KZOlrYxhxgPqD49ztU,561
+taskiq_redis/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+taskiq_redis/redis_backend.py,sha256=BwzFWXLHqpQEhZ675dvt2ueOfB7xjJOAGTSHZSyoR7A,8914
+taskiq_redis/redis_broker.py,sha256=b5oOKXP-uuqffGnNhUsT4HgTPmBiBdAfpUOd5V0VfFc,4254
+taskiq_redis/redis_cluster_broker.py,sha256=CgPKkoEHZ1moNM-VNmzPQdjjNOrhiVUCNV-7FrUgqTo,2121
+taskiq_redis/schedule_source.py,sha256=uznI6wbrdSbD-hIAF3xDGTVXD99SnfWYAuYsTQUIL8E,6202
+taskiq_redis/serializer.py,sha256=x-1ExYoD_EnDiM53lyvI99MdTpNj_pORMIaCL07-6nU,416
+taskiq_redis-0.5.6.dist-info/METADATA,sha256=zM6LiFFui-OFZqAVJnQ-qdh-U7_wKrTBnHYNMJuar4M,4030
+taskiq_redis-0.5.6.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+taskiq_redis-0.5.6.dist-info/RECORD,,

{taskiq_redis-0.5.4.dist-info → taskiq_redis-0.5.6.dist-info}/WHEEL

@@ -1,4 +1,4 @@
 Wheel-Version: 1.0
-Generator: poetry-core 1.8.1
+Generator: poetry-core 1.9.0
 Root-Is-Purelib: true
 Tag: py3-none-any

taskiq_redis-0.5.4.dist-info/RECORD

@@ -1,11 +0,0 @@
-taskiq_redis/__init__.py,sha256=sPv4uLLPSPc8zXK42puvsSMDY5HwbCePdJheYylyApg,527
-taskiq_redis/exceptions.py,sha256=eS4bfZVAjyMsnFs3IF74uYwO1KZOlrYxhxgPqD49ztU,561
-taskiq_redis/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-taskiq_redis/redis_backend.py,sha256=Q_pJ1Bz-NpTyFT68UswBvsNYWkLtnSPF1x8QbOnbWI0,8357
-taskiq_redis/redis_broker.py,sha256=qQLWWvY-NacVXkgDGVCe2fyWYPkjZiOnggB0hpPStqw,3957
-taskiq_redis/redis_cluster_broker.py,sha256=CgPKkoEHZ1moNM-VNmzPQdjjNOrhiVUCNV-7FrUgqTo,2121
-taskiq_redis/schedule_source.py,sha256=SWMGC3limXRgyxZS5BawnF3bHGgJpjljByIUL5nC2BY,3414
-taskiq_redis/serializer.py,sha256=x-1ExYoD_EnDiM53lyvI99MdTpNj_pORMIaCL07-6nU,416
-taskiq_redis-0.5.4.dist-info/METADATA,sha256=XKrSxYM3F5RkGwX5mb-HGxertnZWeBpf1f_9c_yaYPQ,3588
-taskiq_redis-0.5.4.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
-taskiq_redis-0.5.4.dist-info/RECORD,,