taskiq-redis 0.5.3__py3-none-any.whl → 0.5.5__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/schedule_source.py

@@ -1,9 +1,9 @@
-import dataclasses
 from typing import Any, List, Optional
 
-from redis.asyncio import ConnectionPool, Redis
+from redis.asyncio import ConnectionPool, Redis, RedisCluster
 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_redis.serializer import PickleSerializer
@@ -60,7 +60,7 @@ class RedisScheduleSource(ScheduleSource):
         async with Redis(connection_pool=self.connection_pool) as redis:
             await redis.set(
                 f"{self.prefix}:{schedule.schedule_id}",
-                self.serializer.dumpb(dataclasses.asdict(schedule)),
+                self.serializer.dumpb(model_dump(schedule)),
             )
 
     async def get_schedules(self) -> List[ScheduledTask]:
@@ -82,7 +82,7 @@ class RedisScheduleSource(ScheduleSource):
             if buffer:
                 schedules.extend(await redis.mget(buffer))
         return [
-            ScheduledTask(**self.serializer.loadb(schedule))
+            model_validate(ScheduledTask, self.serializer.loadb(schedule))
             for schedule in schedules
             if schedule
         ]
@@ -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.3.dist-info → taskiq_redis-0.5.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: taskiq-redis
-Version: 0.5.3
+Version: 0.5.5
 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.1,<1)
+Requires-Dist: taskiq (>=0.10.3,<1)
 Project-URL: Repository, https://github.com/taskiq-python/taskiq-redis
 Description-Content-Type: text/markdown
 

{taskiq_redis-0.5.3.dist-info → taskiq_redis-0.5.5.dist-info}/RECORD

@@ -1,11 +1,11 @@
-taskiq_redis/__init__.py,sha256=sPv4uLLPSPc8zXK42puvsSMDY5HwbCePdJheYylyApg,527
+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=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=7AdY7394wm0UEPCn7ovT4vpKjSLXTw7EUgdU5rGJEaY,3374
+taskiq_redis/schedule_source.py,sha256=kFJP4418JwhqzhOMoP98EIaLNYGFvQWLKtlNUU0EvsY,6166
 taskiq_redis/serializer.py,sha256=x-1ExYoD_EnDiM53lyvI99MdTpNj_pORMIaCL07-6nU,416
-taskiq_redis-0.5.3.dist-info/METADATA,sha256=LbMWjnYvn4mPuTyyYWgp0-3ExZnqmkbnPQQPYnyOgbE,3588
-taskiq_redis-0.5.3.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
-taskiq_redis-0.5.3.dist-info/RECORD,,
+taskiq_redis-0.5.5.dist-info/METADATA,sha256=5GSOGw4oUTeMojhu9IpIz6ZeEkqbIFrdqcMUbIaQbv4,3588
+taskiq_redis-0.5.5.dist-info/WHEEL,sha256=FMvqSimYX_P7y0a7UY-_Mc83r5zkBZsCYPm7Lr0Bsq4,88
+taskiq_redis-0.5.5.dist-info/RECORD,,