Add scheduler pre_send exception handling (#254)

Author: stinovlas

docs/examples/extending/schedule_source.py

@@ -38,6 +38,9 @@ async def pre_send(self, task: "ScheduledTask") -> None:
         """
         Actions to execute before task will be sent to broker.
 
+        This method may raise ScheduledTaskCancelledError.
+        This cancels the task execution.
+
         :param task: task that will be sent
         """
 

taskiq/abc/schedule_source.py

@@ -60,6 +60,9 @@ def pre_send(  # noqa: B027
         """
         Actions to execute before task will be sent to broker.
 
+        This method may raise ScheduledTaskCancelledError.
+        This cancels the task execution.
+
         :param task: task that will be sent
         """
 

taskiq/exceptions.py

@@ -40,3 +40,7 @@ class NoResultError(TaskiqError):
 
 class TaskRejectedError(TaskiqError):
     """Task was rejected."""
+
+
+class ScheduledTaskCancelledError(TaskiqError):
+    """Scheduled task was cancelled and not sent to the queue."""

taskiq/scheduler/scheduler.py

@@ -1,5 +1,7 @@
+from logging import getLogger
 from typing import TYPE_CHECKING, List
 
+from taskiq.exceptions import ScheduledTaskCancelledError
 from taskiq.kicker import AsyncKicker
 from taskiq.scheduler.scheduled_task import ScheduledTask
 from taskiq.utils import maybe_awaitable
@@ -8,6 +10,8 @@
     from taskiq.abc.broker import AsyncBroker
     from taskiq.abc.schedule_source import ScheduleSource
 
+logger = getLogger(__name__)
+
 
 class TaskiqScheduler:
     """Scheduler class."""
@@ -36,12 +40,16 @@ async def on_ready(self, source: "ScheduleSource", task: ScheduledTask) -> None:
         It's triggered on proper time depending on `task.cron` or `task.time` attribute.
         :param task: task to send
         """
-        await maybe_awaitable(source.pre_send(task))
-        await AsyncKicker(task.task_name, self.broker, task.labels).kiq(
-            *task.args,
-            **task.kwargs,
-        )
-        await maybe_awaitable(source.post_send(task))
+        try:
+            await maybe_awaitable(source.pre_send(task))
+        except ScheduledTaskCancelledError:
+            logger.info("Scheduled task %s has been cancelled.", task.task_name)
+        else:
+            await AsyncKicker(task.task_name, self.broker, task.labels).kiq(
+                *task.args,
+                **task.kwargs,
+            )
+            await maybe_awaitable(source.post_send(task))
 
     async def shutdown(self) -> None:
         """Shutdown the scheduler process."""

tests/scheduler/test_scheduler.py

@@ -0,0 +1,38 @@
+from typing import Any, Coroutine, List, Union
+
+import pytest
+
+from taskiq.abc.schedule_source import ScheduleSource
+from taskiq.brokers.inmemory_broker import InMemoryBroker
+from taskiq.exceptions import ScheduledTaskCancelledError
+from taskiq.scheduler.scheduled_task import ScheduledTask
+from taskiq.scheduler.scheduler import TaskiqScheduler
+
+
+class CancellingScheduleSource(ScheduleSource):
+    async def get_schedules(self) -> List["ScheduledTask"]:
+        """Return schedules list."""
+        return []
+
+    def pre_send(
+        self,
+        task: "ScheduledTask",
+    ) -> Union[None, Coroutine[Any, Any, None]]:
+        """Raise cancelled error."""
+        raise ScheduledTaskCancelledError
+
+
+@pytest.mark.anyio
+async def test_scheduled_task_cancelled() -> None:
+    broker = InMemoryBroker()
+    source = CancellingScheduleSource()
+    scheduler = TaskiqScheduler(broker=broker, sources=[source])
+    task = ScheduledTask(
+        task_name="ping:pong",
+        labels={},
+        args=[],
+        kwargs={},
+        cron="* * * * *",
+    )
+
+    await scheduler.on_ready(source, task)  # error is caught