taskiq-redis 0.5.3__tar.gz → 0.5.5__tar.gz

{taskiq_redis-0.5.3 → taskiq_redis-0.5.5}/PKG-INFO

@@ -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 → taskiq_redis-0.5.5}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "taskiq-redis"
-version = "0.5.3"
+version = "0.5.5"
 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.1,<1"
+taskiq = ">=0.10.3,<1"
 redis = "^5"
 
 [tool.poetry.group.dev.dependencies]

{taskiq_redis-0.5.3 → taskiq_redis-0.5.5}/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-0.5.3 → taskiq_redis-0.5.5}/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)