Added --fail-on-err scheduler arg.

Author: s3rius

taskiq/cli/scheduler/args.py

@@ -20,6 +20,7 @@ class SchedulerArgs:
     skip_first_run: bool = False
     update_interval: int | None = None
     loop_interval: int | None = None
+    fail_on_err: bool = False
 
     @classmethod
     def from_cli(cls, args: Sequence[str] | None = None) -> "SchedulerArgs":
@@ -111,6 +112,12 @@ def from_cli(cls, args: Sequence[str] | None = None) -> "SchedulerArgs":
                 "If not specified, scheduler will run once a second."
             ),
         )
+        parser.add_argument(
+            "--fail-on-err",
+            action="store_false",
+            dest="fail_on_err",
+            help="Fail scheduler on any error during schedules collection",
+        )
 
         namespace = parser.parse_args(args)
         # If there are any patterns specified, remove default.

taskiq/cli/scheduler/run.py

@@ -36,7 +36,10 @@ def to_tz_aware(time: datetime) -> datetime:
     return time
 
 
-async def get_schedules(source: ScheduleSource) -> list[ScheduledTask]:
+async def get_schedules(
+    source: ScheduleSource,
+    fail_on_err: bool,
+) -> list[ScheduledTask]:
     """
     Get schedules from source.
 
@@ -48,6 +51,13 @@ async def get_schedules(source: ScheduleSource) -> list[ScheduledTask]:
     try:
         return await source.get_schedules()
     except Exception as exc:
+        if fail_on_err:
+            logger.warning(
+                "Cannot update schedules with source: %s",
+                source,
+                exc_info=True,
+            )
+            raise exc
         logger.warning(
             "Cannot update schedules with source: %s",
             source,
@@ -58,6 +68,7 @@ async def get_schedules(source: ScheduleSource) -> list[ScheduledTask]:
 
 async def get_all_schedules(
     scheduler: TaskiqScheduler,
+    fail_on_err: bool = False,
 ) -> list[tuple[ScheduleSource, list[ScheduledTask]]]:
     """
     Task to update all schedules.
@@ -71,7 +82,7 @@ async def get_all_schedules(
     """
     logger.debug("Started schedule update.")
     schedules: list[list[ScheduledTask]] = await asyncio.gather(
-        *[get_schedules(source) for source in scheduler.sources],
+        *[get_schedules(source, fail_on_err) for source in scheduler.sources],
     )
     return list(zip(scheduler.sources, schedules, strict=True))
 
@@ -186,8 +197,10 @@ def __init__(
         scheduler: TaskiqScheduler,
         *,
         event_loop: asyncio.AbstractEventLoop | None = None,
+        fail_on_err: bool = False,
     ) -> None:
         self.scheduler = scheduler
+        self.fail_on_err = fail_on_err
         self._event_loop = event_loop or asyncio.get_event_loop()
 
         # Variables for control the last run of schedules.
@@ -310,7 +323,7 @@ async def run(
 
         running_schedules: dict[ScheduleId, asyncio.Task[Any]] = {}
 
-        self.scheduled_tasks = await get_all_schedules(self.scheduler)
+        self.scheduled_tasks = await get_all_schedules(self.scheduler, self.fail_on_err)
         self.scheduled_tasks_updated_at = datetime.now(tz=timezone.utc)
 
         if skip_first_run:
@@ -405,7 +418,7 @@ async def run_scheduler(args: SchedulerArgs) -> None:
     await scheduler.startup()
     logger.info("Startup completed.")
 
-    scheduler_loop = SchedulerLoop(scheduler)
+    scheduler_loop = SchedulerLoop(scheduler, fail_on_err=args.fail_on_err)
     try:
         await scheduler_loop.run(
             update_interval=update_interval,

tests/cli/scheduler/test_updater.py

@@ -1,5 +1,7 @@
 from datetime import datetime
 
+import pytest
+
 from taskiq import InMemoryBroker, ScheduleSource
 from taskiq.cli.scheduler.run import get_all_schedules
 from taskiq.scheduler.scheduled_task import ScheduledTask
@@ -80,3 +82,25 @@ async def test_get_schedules_error() -> None:
         (source1, source1.schedules),
         (source2, []),
     ]
+
+
+async def test_get_schedules_fail_on_err() -> None:
+    """Tests that if source returned an error, empty list will be returned."""
+    source1 = DummySource(
+        [
+            ScheduledTask(
+                task_name="a",
+                labels={},
+                args=[],
+                kwargs={},
+                time=datetime.now(),
+            ),
+        ],
+    )
+    source2 = DummySource(Exception("test"))
+
+    with pytest.raises(Exception, match="test"):
+        await get_all_schedules(
+            TaskiqScheduler(InMemoryBroker(), [source1, source2]),
+            fail_on_err=True,
+        )