feat: Added RedisSentinelScheduleSource.

Author: DABND19

taskiq_redis/__init__.py

@@ -10,6 +10,7 @@
 from taskiq_redis.schedule_source import (
     RedisClusterScheduleSource,
     RedisScheduleSource,
+    RedisSentinelScheduleSource,
 )
 
 __all__ = [
@@ -22,4 +23,5 @@
     "ListQueueSentinelBroker",
     "RedisScheduleSource",
     "RedisClusterScheduleSource",
+    "RedisSentinelScheduleSource",
 ]

taskiq_redis/schedule_source.py

@@ -1,13 +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,
+    ConnectionPool,
+    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_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]
+else:
+    _Redis: TypeAlias = Redis
+
 
 class RedisScheduleSource(ScheduleSource):
     """
@@ -174,3 +192,97 @@ 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)
+
+
+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
+            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)

tests/test_schedule_source.py

@@ -1,11 +1,16 @@
 import asyncio
 import datetime as dt
 import uuid
+from typing import List, Tuple
 
 import pytest
 from taskiq import ScheduledTask
 
-from taskiq_redis import RedisClusterScheduleSource, RedisScheduleSource
+from taskiq_redis import (
+    RedisClusterScheduleSource,
+    RedisScheduleSource,
+    RedisSentinelScheduleSource,
+)
 
 
 @pytest.mark.anyio
@@ -220,3 +225,140 @@ async def test_cluster_buffer(redis_cluster_url: str) -> None:
     assert schedule1 in schedules
     assert schedule2 in schedules
     await source.shutdown()
+
+
+@pytest.mark.anyio
+async def test_sentinel_set_schedule(
+    redis_sentinels: List[Tuple[str, int]],
+    redis_sentinel_master_name: str,
+) -> None:
+    prefix = uuid.uuid4().hex
+    source = RedisSentinelScheduleSource(
+        sentinels=redis_sentinels,
+        master_name=redis_sentinel_master_name,
+        prefix=prefix,
+    )
+    schedule = ScheduledTask(
+        task_name="test_task",
+        labels={},
+        args=[],
+        kwargs={},
+        cron="* * * * *",
+    )
+    await source.add_schedule(schedule)
+    schedules = await source.get_schedules()
+    assert schedules == [schedule]
+    await source.shutdown()
+
+
+@pytest.mark.anyio
+async def test_sentinel_delete_schedule(
+    redis_sentinels: List[Tuple[str, int]],
+    redis_sentinel_master_name: str,
+) -> None:
+    prefix = uuid.uuid4().hex
+    source = RedisSentinelScheduleSource(
+        sentinels=redis_sentinels,
+        master_name=redis_sentinel_master_name,
+        prefix=prefix,
+    )
+    schedule = ScheduledTask(
+        task_name="test_task",
+        labels={},
+        args=[],
+        kwargs={},
+        cron="* * * * *",
+    )
+    await source.add_schedule(schedule)
+    schedules = await source.get_schedules()
+    assert schedules == [schedule]
+    await source.delete_schedule(schedule.schedule_id)
+    schedules = await source.get_schedules()
+    # Schedules are empty.
+    assert not schedules
+    await source.shutdown()
+
+
+@pytest.mark.anyio
+async def test_sentinel_post_run_cron(
+    redis_sentinels: List[Tuple[str, int]],
+    redis_sentinel_master_name: str,
+) -> None:
+    prefix = uuid.uuid4().hex
+    source = RedisSentinelScheduleSource(
+        sentinels=redis_sentinels,
+        master_name=redis_sentinel_master_name,
+        prefix=prefix,
+    )
+    schedule = ScheduledTask(
+        task_name="test_task",
+        labels={},
+        args=[],
+        kwargs={},
+        cron="* * * * *",
+    )
+    await source.add_schedule(schedule)
+    assert await source.get_schedules() == [schedule]
+    await source.post_send(schedule)
+    assert await source.get_schedules() == [schedule]
+    await source.shutdown()
+
+
+@pytest.mark.anyio
+async def test_sentinel_post_run_time(
+    redis_sentinels: List[Tuple[str, int]],
+    redis_sentinel_master_name: str,
+) -> None:
+    prefix = uuid.uuid4().hex
+    source = RedisSentinelScheduleSource(
+        sentinels=redis_sentinels,
+        master_name=redis_sentinel_master_name,
+        prefix=prefix,
+    )
+    schedule = ScheduledTask(
+        task_name="test_task",
+        labels={},
+        args=[],
+        kwargs={},
+        time=dt.datetime(2000, 1, 1),
+    )
+    await source.add_schedule(schedule)
+    assert await source.get_schedules() == [schedule]
+    await source.post_send(schedule)
+    assert await source.get_schedules() == []
+    await source.shutdown()
+
+
+@pytest.mark.anyio
+async def test_sentinel_buffer(
+    redis_sentinels: List[Tuple[str, int]],
+    redis_sentinel_master_name: str,
+) -> None:
+    prefix = uuid.uuid4().hex
+    source = RedisSentinelScheduleSource(
+        sentinels=redis_sentinels,
+        master_name=redis_sentinel_master_name,
+        prefix=prefix,
+        buffer_size=1,
+    )
+    schedule1 = ScheduledTask(
+        task_name="test_task1",
+        labels={},
+        args=[],
+        kwargs={},
+        cron="* * * * *",
+    )
+    schedule2 = ScheduledTask(
+        task_name="test_task2",
+        labels={},
+        args=[],
+        kwargs={},
+        cron="* * * * *",
+    )
+    await source.add_schedule(schedule1)
+    await source.add_schedule(schedule2)
+    schedules = await source.get_schedules()
+    assert len(schedules) == 2
+    assert schedule1 in schedules
+    assert schedule2 in schedules
+    await source.shutdown()