From 0b661ea44a56876d6795fe00e49c72b998687176 Mon Sep 17 00:00:00 2001 From: spikeninja Date: Mon, 12 Oct 2026 16:59:34 +0200 Subject: [PATCH] feat: add cron validation --- docs/available-components/schedule-sources.md | 4 +- taskiq/__init__.py | 3 +- taskiq/cli/scheduler/run.py | 10 +- taskiq/decor.py | 6 +- taskiq/kicker.py | 12 +- taskiq/scheduler/scheduled_task/__init__.py | 3 +- taskiq/scheduler/scheduled_task/cron.py | 204 ++++++++++++++++++ taskiq/scheduler/scheduled_task/cron_spec.py | 16 +- taskiq/scheduler/scheduled_task/task.py | 3 +- tests/scheduler/test_cron.py | 89 ++++++++ 10 files changed, 328 insertions(+), 22 deletions(-) create mode 100644 taskiq/scheduler/scheduled_task/cron.py create mode 100644 tests/scheduler/test_cron.py diff --git a/docs/available-components/schedule-sources.md b/docs/available-components/schedule-sources.md index 814fc9c5..ac7e0254 100644 --- a/docs/available-components/schedule-sources.md +++ b/docs/available-components/schedule-sources.md @@ -36,7 +36,7 @@ The format of the schedule label is the following: @broker.task( schedule=[ { - "cron": "* * * * *", # type: str | None, either cron, interval or time should be specified. + "cron": "* * * * *", # type: str | Cron | None, either cron, interval or time should be specified. "cron_offset": None, # type: str | timedelta | None, can be omitted. "interval": None, # type: int | timedelta | None, either cron, interval or time should be specified. "time": None, # type: datetime | None, either cron, interval or time should be specified. @@ -53,7 +53,7 @@ async def my_task(): Parameters: -- `cron` - crontab string when to run the task. +- `cron` - crontab string (5 fields; names like `MON`/`JAN` are allowed) when to run the task. Invalid values are rejected when the schedule is created. - `cron_offset` - timezone offset for cron values. Explained [here](../guide/scheduling-tasks.md#working-with-timezones) - `interval` - Interval to run periodic tasks. Must be at least 1 second and cannot have fractional seconds. - `time` - specific time when to send the task. diff --git a/taskiq/__init__.py b/taskiq/__init__.py index 1a764bc4..da38c9d1 100644 --- a/taskiq/__init__.py +++ b/taskiq/__init__.py @@ -32,7 +32,7 @@ SmartRetryMiddleware, ) from taskiq.result import TaskiqResult -from taskiq.scheduler.scheduled_task import ScheduledTask +from taskiq.scheduler.scheduled_task import Cron, ScheduledTask from taskiq.scheduler.scheduler import TaskiqScheduler from taskiq.state import TaskiqState from taskiq.task import AsyncTaskiqTask @@ -48,6 +48,7 @@ "AsyncTaskiqTask", "BrokerMessage", "Context", + "Cron", "InMemoryBroker", "NoResultError", "PrometheusMiddleware", diff --git a/taskiq/cli/scheduler/run.py b/taskiq/cli/scheduler/run.py index a0872cd7..b9182a92 100644 --- a/taskiq/cli/scheduler/run.py +++ b/taskiq/cli/scheduler/run.py @@ -11,7 +11,7 @@ from taskiq.abc.schedule_source import ScheduleSource from taskiq.cli.scheduler.args import SchedulerArgs from taskiq.cli.utils import import_object, import_tasks -from taskiq.scheduler.scheduled_task import ScheduledTask +from taskiq.scheduler.scheduled_task import Cron, CronError, ScheduledTask from taskiq.scheduler.scheduler import TaskiqScheduler logger = getLogger(__name__) @@ -77,7 +77,7 @@ class CronValueError(Exception): def is_cron_task_now( - cron_value: str, + cron_value: str | Cron, now: datetime, offset: str | timedelta | None = None, last_run: datetime | None = None, @@ -108,8 +108,10 @@ def is_cron_task_now( raise CronValueError(e) from e try: - return pycron.is_now(cron_value, now) - except ValueError as e: + if isinstance(cron_value, str): + cron_value = Cron.parse(cron_value) + return pycron.is_now(cron_value.to_pycron(), now) + except (CronError, ValueError) as e: raise CronValueError(e) from e diff --git a/taskiq/decor.py b/taskiq/decor.py index 8f90541c..dd4e300d 100644 --- a/taskiq/decor.py +++ b/taskiq/decor.py @@ -20,7 +20,7 @@ if TYPE_CHECKING: # pragma: no cover from taskiq.abc.broker import AsyncBroker from taskiq.abc.schedule_source import ScheduleSource - from taskiq.scheduler.scheduled_task import CronSpec + from taskiq.scheduler.scheduled_task import Cron, CronSpec _T = TypeVar("_T") _FuncParams = ParamSpec("_FuncParams") @@ -139,7 +139,7 @@ async def kiq( async def schedule_by_cron( self, source: "ScheduleSource", - cron: Union[str, "CronSpec"], + cron: Union[str, "CronSpec", "Cron"], *args: _FuncParams.args, **kwargs: _FuncParams.kwargs, ) -> CreatedSchedule[_ReturnType]: @@ -150,7 +150,7 @@ async def schedule_by_cron( which is capable of dynamically adding new schedules. :param source: schedule source. - :param cron: cron string or a CronSpec instance. + :param cron: cron string, Cron or CronSpec instance. :param args: function's arguments. :param kwargs: function's key word arguments. :return: schedule id. diff --git a/taskiq/kicker.py b/taskiq/kicker.py index 07f4e51a..6750dddf 100644 --- a/taskiq/kicker.py +++ b/taskiq/kicker.py @@ -20,7 +20,7 @@ from taskiq.labels import prepare_label from taskiq.message import TaskiqMessage from taskiq.scheduler.created_schedule import CreatedSchedule -from taskiq.scheduler.scheduled_task import CronSpec, ScheduledTask +from taskiq.scheduler.scheduled_task import Cron, CronSpec, ScheduledTask from taskiq.task import AsyncTaskiqTask from taskiq.utils import maybe_awaitable @@ -180,7 +180,7 @@ async def kiq( async def schedule_by_cron( self, source: "ScheduleSource", - cron: Union[str, "CronSpec"], + cron: Union[str, "CronSpec", "Cron"], *args: _FuncParams.args, **kwargs: _FuncParams.kwargs, ) -> CreatedSchedule[_ReturnType]: @@ -188,7 +188,7 @@ async def schedule_by_cron( Function to schedule task with cron. :param source: schedule source. - :param cron: cron expression. + :param cron: cron string, Cron or CronSpec instance. :param args: function's args. :param kwargs: function's kwargs. @@ -200,10 +200,10 @@ async def schedule_by_cron( message = self._prepare_message(*args, **kwargs) cron_offset = None if isinstance(cron, CronSpec): - cron_str = cron.to_cron() + cron_value = cron.to_cron() cron_offset = cron.offset else: - cron_str = cron + cron_value = Cron.parse(cron) if isinstance(cron, str) else cron scheduled = ScheduledTask( schedule_id=schedule_id, task_name=message.task_name, @@ -211,7 +211,7 @@ async def schedule_by_cron( args=message.args, kwargs=message.kwargs, task_id=self.custom_task_id, - cron=cron_str, + cron=cron_value, cron_offset=cron_offset, ) await source.add_schedule(scheduled) diff --git a/taskiq/scheduler/scheduled_task/__init__.py b/taskiq/scheduler/scheduled_task/__init__.py index b597bd37..44e6d9ff 100644 --- a/taskiq/scheduler/scheduled_task/__init__.py +++ b/taskiq/scheduler/scheduled_task/__init__.py @@ -1,4 +1,5 @@ +from .cron import Cron, CronError from .cron_spec import CronSpec from .task import ScheduledTask -__all__ = ["CronSpec", "ScheduledTask"] +__all__ = ["Cron", "CronError", "CronSpec", "ScheduledTask"] diff --git a/taskiq/scheduler/scheduled_task/cron.py b/taskiq/scheduler/scheduled_task/cron.py new file mode 100644 index 00000000..0f8cae72 --- /dev/null +++ b/taskiq/scheduler/scheduled_task/cron.py @@ -0,0 +1,204 @@ +from collections.abc import Mapping +from dataclasses import dataclass, field +from enum import Enum +from typing import Annotated, Any + +from pydantic import PlainSerializer, PlainValidator, WithJsonSchema + + +class CronError(ValueError): + """Raised on invalid cron expression.""" + + +class WeekDays(str, Enum): + """Weekday names.""" + + SUN = "SUN" + MON = "MON" + TUE = "TUE" + WED = "WED" + THU = "THU" + FRI = "FRI" + SAT = "SAT" + + +class Months(str, Enum): + """Month names.""" + + JAN = "JAN" + FEB = "FEB" + MAR = "MAR" + APR = "APR" + MAY = "MAY" + JUN = "JUN" + JUL = "JUL" + AUG = "AUG" + SEP = "SEP" + OCT = "OCT" + NOV = "NOV" + DEC = "DEC" + + +@dataclass(frozen=True) +class FieldSpec: + """Allowed values of a cron field.""" + + name: str + low: int + high: int + aliases: Mapping[str, int] = field(default_factory=dict) + + +MONTH_NAMES: dict[str, int] = {n.value: i for i, n in enumerate(Months, start=1)} +WEEKDAY_NAMES: dict[str, int] = {n.value: i for i, n in enumerate(WeekDays)} + +MINUTE = FieldSpec("minute", 0, 59) +HOUR = FieldSpec("hour", 0, 23) +DAY = FieldSpec("day", 1, 31) +MONTH = FieldSpec("month", 1, 12, MONTH_NAMES) +WEEKDAY = FieldSpec("weekday", 0, 7, WEEKDAY_NAMES) # 0 and 7 = SUN + + +def _is_uint(s: str) -> bool: + return s.isascii() and s.isdigit() + + +@dataclass(frozen=True) +class CronField: + """ + Single cron field. + + Accepted values: + *: any value + N: one value + A-B: a range, A <= B + */S: a step within the whole range + A-B/S: a step inside the range + Comma separated lists of the above. + """ + + value: str + spec: FieldSpec + + def __post_init__(self) -> None: + for part in self.value.split(","): + self._check_part(part) + + def _check_part(self, part: str) -> None: + name = self.spec.name + + base, has_step, step = part.partition("/") + if has_step: + if not _is_uint(step) or int(step) == 0: + raise CronError(f"{name}: the step should be > 0, got {step!r}") + if base != "*" and "-" not in base: + raise CronError( + f"{name}: the step is allowed only after '*' or a range, " + f"got {part!r}", + ) + + if base == "*": + return + + start, has_range, end = base.partition("-") + lo = self._to_number(start) + if has_range: + hi = self._to_number(end) + if lo > hi: + raise CronError(f"{name}: start must be lower than the end {part!r}") + + def _to_number(self, token: str) -> int: + spec = self.spec + alias = token.upper() + if alias in spec.aliases: + return spec.aliases[alias] + if not _is_uint(token): + raise CronError(f"{spec.name}: a number expected, got {token!r}") + n = int(token) + if not spec.low <= n <= spec.high: + raise CronError(f"{spec.name}: {n} is out of range {spec.low}-{spec.high}") + return n + + def to_numeric(self) -> str: + """Convert the field to numbers only, in the form pycron understands.""" + is_weekday = self.spec is WEEKDAY + parts: list[str] = [] + for part in self.value.split(","): + base, _, step = part.partition("/") + if base == "*": + parts.append(part) + continue + start, has_range, end = base.partition("-") + lo = self._to_number(start) + if not has_range: + parts.append(str(lo % 7 if is_weekday else lo)) + elif is_weekday: + values = range(lo, self._to_number(end) + 1, int(step or 1)) + parts.extend(str(v % 7) for v in values) + else: + suffix = f"/{step}" if step else "" + parts.append(f"{lo}-{self._to_number(end)}{suffix}") + return ",".join(parts) + + +@dataclass +class Cron: + """Validated 5-field cron expression.""" + + minute: CronField + hour: CronField + day: CronField + month: CronField + weekday: CronField + + def __post_init__(self) -> None: + expected = (MINUTE, HOUR, DAY, MONTH, WEEKDAY) + actual = (self.minute, self.hour, self.day, self.month, self.weekday) + for f, spec in zip(actual, expected, strict=True): + if f.spec is not spec: + raise CronError( + f"field {spec.name} is created with wrong spec: {f.spec.name}", + ) + + @classmethod + def parse(cls, expr: str) -> "Cron": + """ + Parse cron string. + + :param expr: cron expression with 5 fields. + :raises CronError: if the expression is invalid. + :return: cron object. + """ + parts = expr.split() + if len(parts) != 5: + raise CronError(f"Need 5 fields, got: {len(parts)}: {expr!r}") + + specs = (MINUTE, HOUR, DAY, MONTH, WEEKDAY) + return cls(*(CronField(v, s) for v, s in zip(parts, specs, strict=True))) + + def to_pycron(self) -> str: + """Cron string with numeric values only, suitable for pycron.""" + fields = (self.minute, self.hour, self.day, self.month, self.weekday) + return " ".join(f.to_numeric() for f in fields) + + def __str__(self) -> str: + return ( + f"{self.minute.value} {self.hour.value} {self.day.value} " + f"{self.month.value} {self.weekday.value}" + ) + + +def _validate_cron(value: Any) -> Cron: + if isinstance(value, Cron): + return value + if isinstance(value, str): + return Cron.parse(value) + raise CronError(f"cron must be a string or Cron, got {type(value).__name__}") + + +CronExpression = Annotated[ + Cron | str, + PlainValidator(_validate_cron), + PlainSerializer(str, return_type=str), + WithJsonSchema({"type": "string"}), +] diff --git a/taskiq/scheduler/scheduled_task/cron_spec.py b/taskiq/scheduler/scheduled_task/cron_spec.py index 1bce0b1d..0cf5028b 100644 --- a/taskiq/scheduler/scheduled_task/cron_spec.py +++ b/taskiq/scheduler/scheduled_task/cron_spec.py @@ -1,5 +1,6 @@ -from pydantic import BaseModel +from pydantic import BaseModel, model_validator +from taskiq.scheduler.scheduled_task.cron import Cron from taskiq.scheduler.scheduled_task.validators import CronOffset @@ -14,6 +15,13 @@ class CronSpec(BaseModel): offset: CronOffset | None = None - def to_cron(self) -> str: # pragma: no cover - """Converts cron spec to cron string.""" - return f"{self.minutes} {self.hours} {self.days} {self.months} {self.weekdays}" + @model_validator(mode="after") + def _validate_cron(self) -> "CronSpec": + self.to_cron() + return self + + def to_cron(self) -> Cron: + """Converts cron spec to a validated cron object.""" + return Cron.parse( + f"{self.minutes} {self.hours} {self.days} {self.months} {self.weekdays}", + ) diff --git a/taskiq/scheduler/scheduled_task/task.py b/taskiq/scheduler/scheduled_task/task.py index 86ec0586..73d1c15a 100644 --- a/taskiq/scheduler/scheduled_task/task.py +++ b/taskiq/scheduler/scheduled_task/task.py @@ -5,6 +5,7 @@ from pydantic import BaseModel, Field, model_validator +from taskiq.scheduler.scheduled_task.cron import CronExpression from taskiq.scheduler.scheduled_task.validators import ( CronOffset, validate_interval_value, @@ -25,7 +26,7 @@ class ScheduledTask(BaseModel): kwargs: dict[str, Any] task_id: str | None = None schedule_id: str = Field(default_factory=lambda: uuid.uuid4().hex) - cron: str | None = None + cron: CronExpression | None = None cron_offset: CronOffset | None = None time: datetime | None = None interval: int | timedelta | None = None diff --git a/tests/scheduler/test_cron.py b/tests/scheduler/test_cron.py new file mode 100644 index 00000000..9cc9c14e --- /dev/null +++ b/tests/scheduler/test_cron.py @@ -0,0 +1,89 @@ +from datetime import datetime, timezone + +import pycron +import pytest +from pydantic import ValidationError + +from taskiq.scheduler.scheduled_task import Cron, CronError, CronSpec, ScheduledTask + + +@pytest.mark.parametrize( + "expr", + [ + "* * * * *", + "0 0 1 1 *", + "*/5 1-3 1,15 JAN-MAR MON-FRI", + "0-30/10 * * * sun", + "0 0 * * 7", + "0 0 * * 5-7", + ], +) +def test_valid(expr: str) -> None: + assert str(Cron.parse(expr)) == expr + + +@pytest.mark.parametrize( + "expr", + [ + "invalid cron", + "* * * *", + "60 * * * *", + "* 24 * * *", + "* * 0 * *", + "* * * 13 *", + "* * * * 8", + "*/0 * * * *", + "5/10 * * * *", + "5-1 * * * *", + "1,,2 * * * *", + "@daily", + ], +) +def test_invalid(expr: str) -> None: + with pytest.raises(CronError): + Cron.parse(expr) + + +@pytest.mark.parametrize( + ("expr", "expected"), + [ + ("* * * * MON-FRI", "* * * * 1,2,3,4,5"), + ("* * * * 7", "* * * * 0"), + ("* * * * 5-7", "* * * * 5,6,0"), + ("* * * jan-mar *", "* * * 1-3 *"), + ("*/5 0-20/10 * * *", "*/5 0-20/10 * * *"), + ], +) +def test_to_pycron(expr: str, expected: str) -> None: + assert Cron.parse(expr).to_pycron() == expected + + +def test_pycron_matches_sunday() -> None: + sunday = datetime(2023, 1, 1, tzinfo=timezone.utc) + for expr in ("* * * * SUN", "* * * * 7", "* * * * 5-7"): + assert pycron.is_now(Cron.parse(expr).to_pycron(), sunday) + + +def test_scheduled_task_cron() -> None: + task = ScheduledTask( + task_name="a", + labels={}, + args=[], + kwargs={}, + cron="*/5 * * * *", + ) + assert isinstance(task.cron, Cron) + dumped = task.model_dump(mode="json") + assert dumped["cron"] == "*/5 * * * *" + assert ScheduledTask.model_validate(dumped).cron == task.cron + + +def test_scheduled_task_invalid_cron() -> None: + with pytest.raises(ValidationError): + ScheduledTask(task_name="a", labels={}, args=[], kwargs={}, cron="bad") + + +def test_cron_spec() -> None: + assert str(CronSpec(minutes=5).to_cron()) == "5 * * * *" + with pytest.raises(ValidationError): + CronSpec(minutes=99)