Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/schemas/schedules.py: 15%
256 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-12-05 13:38 +0000
1"""
2Schedule schemas
3"""
5from __future__ import annotations 1a
7import datetime 1a
8import sys 1a
9from typing import ( 1a
10 Annotated,
11 Any,
12 ClassVar,
13 Generator,
14 List,
15 Optional,
16 Tuple,
17 Union,
18)
19from zoneinfo import ZoneInfo 1a
21import dateutil 1a
22import dateutil.rrule 1a
23import pytz 1a
24from pydantic import ConfigDict, Field, field_validator, model_validator 1a
25from typing_extensions import TypeAlias 1a
27from prefect._internal.schemas.validators import ( 1a
28 default_timezone,
29 validate_cron_string,
30 validate_rrule_string,
31)
32from prefect._vendor.croniter import croniter 1a
33from prefect.server.utilities.schemas.bases import PrefectBaseModel 1a
34from prefect.types import DateTime, TimeZone 1a
35from prefect.types._datetime import create_datetime_instance, now 1a
37MAX_ITERATIONS = 1000 1a
39if sys.version_info >= (3, 13): 39 ↛ 40line 39 didn't jump to line 40 because the condition on line 39 was never true1a
40 AnchorDate: TypeAlias = datetime.datetime
41else:
42 from pydantic import AfterValidator 1a
44 from prefect._internal.schemas.validators import default_anchor_date 1a
46 AnchorDate: TypeAlias = Annotated[DateTime, AfterValidator(default_anchor_date)] 1a
49def _prepare_scheduling_start_and_end( 1a
50 start: Any, end: Any, timezone: str
51) -> Tuple[DateTime, Optional[DateTime]]:
52 """Uniformly prepares the start and end dates for any Schedule's get_dates call,
53 coercing the arguments into timezone-aware datetimes."""
54 timezone = timezone or "UTC"
56 if start is not None:
57 if sys.version_info >= (3, 13):
58 start = create_datetime_instance(start).astimezone(ZoneInfo(timezone))
59 else:
60 start = create_datetime_instance(start).in_tz(timezone)
62 if end is not None:
63 if sys.version_info >= (3, 13):
64 end = create_datetime_instance(end).astimezone(ZoneInfo(timezone))
65 else:
66 end = create_datetime_instance(end).in_tz(timezone)
68 return start, end
71class IntervalSchedule(PrefectBaseModel): 1a
72 """
73 A schedule formed by adding `interval` increments to an `anchor_date`. If no
74 `anchor_date` is supplied, the current UTC time is used. If a
75 timezone-naive datetime is provided for `anchor_date`, it is assumed to be
76 in the schedule's timezone (or UTC). Even if supplied with an IANA timezone,
77 anchor dates are always stored as UTC offsets, so a `timezone` can be
78 provided to determine localization behaviors like DST boundary handling. If
79 none is provided it will be inferred from the anchor date.
81 NOTE: If the `IntervalSchedule` `anchor_date` or `timezone` is provided in a
82 DST-observing timezone, then the schedule will adjust itself appropriately.
83 Intervals greater than 24 hours will follow DST conventions, while intervals
84 of less than 24 hours will follow UTC intervals. For example, an hourly
85 schedule will fire every UTC hour, even across DST boundaries. When clocks
86 are set back, this will result in two runs that *appear* to both be
87 scheduled for 1am local time, even though they are an hour apart in UTC
88 time. For longer intervals, like a daily schedule, the interval schedule
89 will adjust for DST boundaries so that the clock-hour remains constant. This
90 means that a daily schedule that always fires at 9am will observe DST and
91 continue to fire at 9am in the local time zone.
93 Args:
94 interval (datetime.timedelta): an interval to schedule on.
95 anchor_date (DateTime, optional): an anchor date to schedule increments against;
96 if not provided, the current timestamp will be used.
97 timezone (str, optional): a valid timezone string.
98 """
100 model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") 1a
102 interval: datetime.timedelta = Field(gt=datetime.timedelta(0)) 1a
103 anchor_date: AnchorDate = Field( 1a
104 default_factory=lambda: now("UTC"),
105 examples=["2020-01-01T00:00:00Z"],
106 )
107 timezone: Optional[str] = Field(default=None, examples=["America/New_York"]) 1a
109 @model_validator(mode="after") 1a
110 def validate_timezone(self): 1a
111 self.timezone = default_timezone(self.timezone, self.model_dump())
112 return self
114 async def get_dates( 1a
115 self,
116 n: Optional[int] = None,
117 start: Optional[datetime.datetime] = None,
118 end: Optional[datetime.datetime] = None,
119 ) -> List[DateTime]:
120 """Retrieves dates from the schedule. Up to 1,000 candidate dates are checked
121 following the start date.
123 Args:
124 n (int): The number of dates to generate
125 start (datetime.datetime, optional): The first returned date will be on or
126 after this date. Defaults to None. If a timezone-naive datetime is
127 provided, it is assumed to be in the schedule's timezone.
128 end (datetime.datetime, optional): The maximum scheduled date to return. If
129 a timezone-naive datetime is provided, it is assumed to be in the
130 schedule's timezone.
132 Returns:
133 List[DateTime]: A list of dates
134 """
135 return sorted(self._get_dates_generator(n=n, start=start, end=end))
137 def _get_dates_generator( 1a
138 self,
139 n: Optional[int] = None,
140 start: Optional[datetime.datetime] = None,
141 end: Optional[datetime.datetime] = None,
142 ) -> Generator[DateTime, None, None]:
143 """Retrieves dates from the schedule. Up to 1,000 candidate dates are checked
144 following the start date.
146 Args:
147 n (Optional[int]): The number of dates to generate
148 start (Optional[datetime.datetime]): The first returned date will be on or
149 after this date. Defaults to None. If a timezone-naive datetime is
150 provided, it is assumed to be in the schedule's timezone.
151 end (Optional[datetime.datetime]): The maximum scheduled date to return. If
152 a timezone-naive datetime is provided, it is assumed to be in the
153 schedule's timezone.
155 Returns:
156 List[DateTime]: a list of dates
157 """
158 if n is None:
159 # if an end was supplied, we do our best to supply all matching dates (up to
160 # MAX_ITERATIONS)
161 if end is not None:
162 n = MAX_ITERATIONS
163 else:
164 n = 1
166 if sys.version_info >= (3, 13):
167 # `pendulum` is not supported in Python 3.13, so we use `whenever` instead
168 from whenever import PlainDateTime, ZonedDateTime
170 if start is None:
171 start = ZonedDateTime.now("UTC").py_datetime()
173 target_timezone = self.timezone or "UTC"
175 def to_local_zdt(dt: datetime.datetime | None) -> ZonedDateTime | None:
176 if dt is None:
177 return None
178 if dt.tzinfo is None:
179 return PlainDateTime.from_py_datetime(dt).assume_tz(target_timezone)
180 if isinstance(dt.tzinfo, ZoneInfo):
181 return ZonedDateTime.from_py_datetime(dt).to_tz(target_timezone)
182 # For offset-based tzinfo instances (e.g. datetime.timezone(+09:00)),
183 # use astimezone to preserve the instant, then convert to ZonedDateTime.
184 return ZonedDateTime.from_py_datetime(
185 dt.astimezone(ZoneInfo(target_timezone))
186 )
188 anchor_zdt = to_local_zdt(self.anchor_date)
189 assert anchor_zdt is not None
191 local_start = to_local_zdt(start)
192 assert local_start is not None
194 local_end = to_local_zdt(end)
196 offset = (
197 local_start - anchor_zdt
198 ).in_seconds() / self.interval.total_seconds()
199 next_date = anchor_zdt.add(
200 seconds=self.interval.total_seconds() * int(offset)
201 )
203 # break the interval into `days` and `seconds` because the datetime
204 # library will handle DST boundaries properly if days are provided, but not
205 # if we add `total seconds`. Therefore, `next_date + self.interval`
206 # fails while `next_date.add(days=days, seconds=seconds)` works.
207 interval_days = self.interval.days
208 interval_seconds = self.interval.total_seconds() - (
209 interval_days * 24 * 60 * 60
210 )
212 while next_date < local_start:
213 next_date = next_date.add(days=interval_days, seconds=interval_seconds)
215 counter = 0
216 dates: set[ZonedDateTime] = set()
218 while True:
219 # if the end date was exceeded, exit
220 if local_end and next_date > local_end:
221 break
223 # ensure no duplicates; weird things can happen with DST
224 if next_date not in dates:
225 dates.add(next_date)
226 yield next_date.py_datetime()
228 # if enough dates have been collected or enough attempts were made, exit
229 if len(dates) >= n or counter > MAX_ITERATIONS:
230 break
232 counter += 1
234 next_date = next_date.add(days=interval_days, seconds=interval_seconds)
236 else:
237 if start is None:
238 start = now("UTC")
239 anchor_tz = self.anchor_date.in_tz(self.timezone)
240 start, end = _prepare_scheduling_start_and_end(start, end, self.timezone)
242 # compute the offset between the anchor date and the start date to jump to the
243 # next date
244 offset = (start - anchor_tz).total_seconds() / self.interval.total_seconds()
245 next_date = anchor_tz.add(
246 seconds=self.interval.total_seconds() * int(offset)
247 )
249 # break the interval into `days` and `seconds` because the datetime
250 # library will handle DST boundaries properly if days are provided, but not
251 # if we add `total seconds`. Therefore, `next_date + self.interval`
252 # fails while `next_date.add(days=days, seconds=seconds)` works.
253 interval_days = self.interval.days
254 interval_seconds = self.interval.total_seconds() - (
255 interval_days * 24 * 60 * 60
256 )
258 # daylight saving time boundaries can create a situation where the next date is
259 # before the start date, so we advance it if necessary
260 while next_date < start:
261 next_date = next_date.add(days=interval_days, seconds=interval_seconds)
263 counter = 0
264 dates = set()
266 while True:
267 # if the end date was exceeded, exit
268 if end and next_date > end:
269 break
271 # ensure no duplicates; weird things can happen with DST
272 if next_date not in dates:
273 dates.add(next_date)
274 yield next_date
276 # if enough dates have been collected or enough attempts were made, exit
277 if len(dates) >= n or counter > MAX_ITERATIONS:
278 break
280 counter += 1
282 next_date = next_date.add(days=interval_days, seconds=interval_seconds)
285class CronSchedule(PrefectBaseModel): 1a
286 """
287 Cron schedule
289 NOTE: If the timezone is a DST-observing one, then the schedule will adjust
290 itself appropriately. Cron's rules for DST are based on schedule times, not
291 intervals. This means that an hourly cron schedule will fire on every new
292 schedule hour, not every elapsed hour; for example, when clocks are set back
293 this will result in a two-hour pause as the schedule will fire *the first
294 time* 1am is reached and *the first time* 2am is reached, 120 minutes later.
295 Longer schedules, such as one that fires at 9am every morning, will
296 automatically adjust for DST.
298 Args:
299 cron (str): a valid cron string
300 timezone (str): a valid timezone string in IANA tzdata format (for example,
301 America/New_York).
302 day_or (bool, optional): Control how croniter handles `day` and `day_of_week`
303 entries. Defaults to True, matching cron which connects those values using
304 OR. If the switch is set to False, the values are connected using AND. This
305 behaves like fcron and enables you to e.g. define a job that executes each
306 2nd friday of a month by setting the days of month and the weekday.
307 """
309 model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") 1a
311 cron: str = Field(default=..., examples=["0 0 * * *"]) 1a
312 timezone: Optional[str] = Field(default=None, examples=["America/New_York"]) 1a
313 day_or: bool = Field( 1a
314 default=True,
315 description=(
316 "Control croniter behavior for handling day and day_of_week entries."
317 ),
318 )
320 @model_validator(mode="after") 1a
321 def validate_timezone(self): 1a
322 self.timezone = default_timezone(self.timezone, self.model_dump())
323 return self
325 @field_validator("cron") 1a
326 @classmethod 1a
327 def valid_cron_string(cls, v: str) -> str: 1a
328 return validate_cron_string(v)
330 async def get_dates( 1a
331 self,
332 n: Optional[int] = None,
333 start: Optional[datetime.datetime] = None,
334 end: Optional[datetime.datetime] = None,
335 ) -> List[DateTime]:
336 """Retrieves dates from the schedule. Up to 1,000 candidate dates are checked
337 following the start date.
339 Args:
340 n (int): The number of dates to generate
341 start (datetime.datetime, optional): The first returned date will be on or
342 after this date. Defaults to None. If a timezone-naive datetime is
343 provided, it is assumed to be in the schedule's timezone.
344 end (datetime.datetime, optional): The maximum scheduled date to return. If
345 a timezone-naive datetime is provided, it is assumed to be in the
346 schedule's timezone.
348 Returns:
349 List[DateTime]: A list of dates
350 """
351 return sorted(self._get_dates_generator(n=n, start=start, end=end))
353 def _get_dates_generator( 1a
354 self,
355 n: Optional[int] = None,
356 start: Optional[datetime.datetime] = None,
357 end: Optional[datetime.datetime] = None,
358 ) -> Generator[DateTime, None, None]:
359 """Retrieves dates from the schedule. Up to 1,000 candidate dates are checked
360 following the start date.
362 Args:
363 n (int): The number of dates to generate
364 start (datetime.datetime, optional): The first returned date will be on or
365 after this date. Defaults to the current date. If a timezone-naive
366 datetime is provided, it is assumed to be in the schedule's timezone.
367 end (datetime.datetime, optional): No returned date will exceed this date.
368 If a timezone-naive datetime is provided, it is assumed to be in the
369 schedule's timezone.
371 Returns:
372 List[DateTime]: a list of dates
373 """
374 if start is None:
375 start = now("UTC")
377 start, end = _prepare_scheduling_start_and_end(start, end, self.timezone)
379 if n is None:
380 # if an end was supplied, we do our best to supply all matching dates (up to
381 # MAX_ITERATIONS)
382 if end is not None:
383 n = MAX_ITERATIONS
384 else:
385 n = 1
387 if self.timezone:
388 if sys.version_info >= (3, 13):
389 start = start.astimezone(ZoneInfo(self.timezone or "UTC"))
390 else:
391 start = start.in_tz(self.timezone)
393 # subtract one second from the start date, so that croniter returns it
394 # as an event (if it meets the cron criteria)
395 start = start - datetime.timedelta(seconds=1)
397 # Respect microseconds by rounding up
398 if start.microsecond > 0:
399 start += datetime.timedelta(seconds=1)
401 # croniter's DST logic interferes with all other datetime libraries except pytz
402 if sys.version_info >= (3, 13):
403 start_localized = start.astimezone(ZoneInfo(self.timezone or "UTC"))
404 start_naive_tz = start.replace(tzinfo=None)
405 else:
406 start_localized = pytz.timezone(start.tz.name).localize(
407 datetime.datetime(
408 year=start.year,
409 month=start.month,
410 day=start.day,
411 hour=start.hour,
412 minute=start.minute,
413 second=start.second,
414 microsecond=start.microsecond,
415 )
416 )
417 start_naive_tz = start.naive()
419 cron = croniter(self.cron, start_naive_tz, day_or=self.day_or) # type: ignore
420 dates = set()
421 counter = 0
423 while True:
424 # croniter does not handle DST properly when the start time is
425 # in and around when the actual shift occurs. To work around this,
426 # we use the naive start time to get the next cron date delta, then
427 # add that time to the original scheduling anchor.
428 next_time = cron.get_next(datetime.datetime)
429 delta = next_time - start_naive_tz
430 if sys.version_info >= (3, 13):
431 from whenever import ZonedDateTime
433 # Use `whenever` to handle DST correctly
434 next_date = (
435 ZonedDateTime.from_py_datetime(start_localized + delta)
436 .to_tz(self.timezone or "UTC")
437 .py_datetime()
438 )
439 else:
440 next_date = create_datetime_instance(start_localized + delta)
442 # if the end date was exceeded, exit
443 if end and next_date > end:
444 break
445 # ensure no duplicates; weird things can happen with DST
446 if next_date not in dates:
447 dates.add(next_date)
448 yield next_date
450 # if enough dates have been collected or enough attempts were made, exit
451 if len(dates) >= n or counter > MAX_ITERATIONS:
452 break
454 counter += 1
457DEFAULT_ANCHOR_DATE = datetime.date(2020, 1, 1) 1a
460class RRuleSchedule(PrefectBaseModel): 1a
461 """
462 RRule schedule, based on the iCalendar standard
463 ([RFC 5545](https://datatracker.ietf.org/doc/html/rfc5545)) as
464 implemented in `dateutils.rrule`.
466 RRules are appropriate for any kind of calendar-date manipulation, including
467 irregular intervals, repetition, exclusions, week day or day-of-month
468 adjustments, and more.
470 Note that as a calendar-oriented standard, `RRuleSchedules` are sensitive to
471 to the initial timezone provided. A 9am daily schedule with a daylight saving
472 time-aware start date will maintain a local 9am time through DST boundaries;
473 a 9am daily schedule with a UTC start date will maintain a 9am UTC time.
475 Args:
476 rrule (str): a valid RRule string
477 timezone (str, optional): a valid timezone string
478 """
480 model_config: ClassVar[ConfigDict] = ConfigDict(extra="forbid") 1a
482 rrule: str 1a
483 timezone: Optional[TimeZone] = "UTC" 1a
485 @field_validator("rrule") 1a
486 @classmethod 1a
487 def validate_rrule_str(cls, v: str) -> str: 1a
488 return validate_rrule_string(v)
490 @classmethod 1a
491 def from_rrule( 1a
492 cls, rrule: dateutil.rrule.rrule | dateutil.rrule.rruleset
493 ) -> "RRuleSchedule":
494 if isinstance(rrule, dateutil.rrule.rrule):
495 if rrule._dtstart.tzinfo is not None:
496 timezone = getattr(rrule._dtstart.tzinfo, "name", None) or getattr(
497 rrule._dtstart.tzinfo, "key", "UTC"
498 )
499 else:
500 timezone = "UTC"
501 return RRuleSchedule(rrule=str(rrule), timezone=timezone)
502 elif isinstance(rrule, dateutil.rrule.rruleset):
503 dtstarts = [rr._dtstart for rr in rrule._rrule if rr._dtstart is not None]
504 unique_dstarts = set(
505 create_datetime_instance(d).astimezone(ZoneInfo("UTC"))
506 for d in dtstarts
507 )
508 unique_timezones = set(d.tzinfo for d in dtstarts if d.tzinfo is not None)
510 if len(unique_timezones) > 1:
511 raise ValueError(
512 f"rruleset has too many dtstart timezones: {unique_timezones}"
513 )
515 if len(unique_dstarts) > 1:
516 raise ValueError(f"rruleset has too many dtstarts: {unique_dstarts}")
518 if unique_dstarts and unique_timezones:
519 tzinfo = dtstarts[0].tzinfo
520 timezone = getattr(tzinfo, "name", None) or getattr(
521 tzinfo, "key", "UTC"
522 )
523 else:
524 timezone = "UTC"
526 rruleset_string = ""
527 if rrule._rrule:
528 rruleset_string += "\n".join(str(r) for r in rrule._rrule)
529 if rrule._exrule:
530 rruleset_string += "\n" if rruleset_string else ""
531 rruleset_string += "\n".join(str(r) for r in rrule._exrule).replace(
532 "RRULE", "EXRULE"
533 )
534 if rrule._rdate:
535 rruleset_string += "\n" if rruleset_string else ""
536 rruleset_string += "RDATE:" + ",".join(
537 rd.strftime("%Y%m%dT%H%M%SZ") for rd in rrule._rdate
538 )
539 if rrule._exdate:
540 rruleset_string += "\n" if rruleset_string else ""
541 rruleset_string += "EXDATE:" + ",".join(
542 exd.strftime("%Y%m%dT%H%M%SZ") for exd in rrule._exdate
543 )
544 return RRuleSchedule(rrule=rruleset_string, timezone=timezone)
545 else:
546 raise ValueError(f"Invalid RRule object: {rrule}")
548 def to_rrule(self) -> dateutil.rrule.rrule: 1a
549 """
550 Since rrule doesn't properly serialize/deserialize timezones, we localize dates
551 here
552 """
553 rrule = dateutil.rrule.rrulestr(
554 self.rrule,
555 dtstart=DEFAULT_ANCHOR_DATE,
556 cache=True,
557 )
558 timezone = dateutil.tz.gettz(self.timezone)
559 if isinstance(rrule, dateutil.rrule.rrule):
560 kwargs = dict(dtstart=rrule._dtstart.replace(tzinfo=timezone))
561 if rrule._until:
562 kwargs.update(
563 until=rrule._until.replace(tzinfo=timezone),
564 )
565 return rrule.replace(**kwargs)
566 elif isinstance(rrule, dateutil.rrule.rruleset):
567 # update rrules
568 localized_rrules = []
569 for rr in rrule._rrule:
570 kwargs = dict(dtstart=rr._dtstart.replace(tzinfo=timezone))
571 if rr._until:
572 kwargs.update(
573 until=rr._until.replace(tzinfo=timezone),
574 )
575 localized_rrules.append(rr.replace(**kwargs))
576 rrule._rrule = localized_rrules
578 # update exrules
579 localized_exrules = []
580 for exr in rrule._exrule:
581 kwargs = dict(dtstart=exr._dtstart.replace(tzinfo=timezone))
582 if exr._until:
583 kwargs.update(
584 until=exr._until.replace(tzinfo=timezone),
585 )
586 localized_exrules.append(exr.replace(**kwargs))
587 rrule._exrule = localized_exrules
589 # update rdates
590 localized_rdates = []
591 for rd in rrule._rdate:
592 localized_rdates.append(rd.replace(tzinfo=timezone))
593 rrule._rdate = localized_rdates
595 # update exdates
596 localized_exdates = []
597 for exd in rrule._exdate:
598 localized_exdates.append(exd.replace(tzinfo=timezone))
599 rrule._exdate = localized_exdates
601 return rrule
603 async def get_dates( 1a
604 self,
605 n: Optional[int] = None,
606 start: datetime.datetime = None,
607 end: datetime.datetime = None,
608 ) -> List[DateTime]:
609 """Retrieves dates from the schedule. Up to 1,000 candidate dates are checked
610 following the start date.
612 Args:
613 n (int): The number of dates to generate
614 start (datetime.datetime, optional): The first returned date will be on or
615 after this date. Defaults to None. If a timezone-naive datetime is
616 provided, it is assumed to be in the schedule's timezone.
617 end (datetime.datetime, optional): The maximum scheduled date to return. If
618 a timezone-naive datetime is provided, it is assumed to be in the
619 schedule's timezone.
621 Returns:
622 List[DateTime]: A list of dates
623 """
624 return sorted(self._get_dates_generator(n=n, start=start, end=end))
626 def _get_dates_generator( 1a
627 self,
628 n: Optional[int] = None,
629 start: Optional[datetime.datetime] = None,
630 end: Optional[datetime.datetime] = None,
631 ) -> Generator[DateTime, None, None]:
632 """Retrieves dates from the schedule. Up to 1,000 candidate dates are checked
633 following the start date.
635 Args:
636 n (int): The number of dates to generate
637 start (datetime.datetime, optional): The first returned date will be on or
638 after this date. Defaults to the current date. If a timezone-naive
639 datetime is provided, it is assumed to be in the schedule's timezone.
640 end (datetime.datetime, optional): No returned date will exceed this date.
641 If a timezone-naive datetime is provided, it is assumed to be in the
642 schedule's timezone.
644 Returns:
645 List[DateTime]: a list of dates
646 """
647 if start is None:
648 start = now("UTC")
650 start, end = _prepare_scheduling_start_and_end(start, end, self.timezone)
652 if n is None:
653 # if an end was supplied, we do our best to supply all matching dates (up
654 # to MAX_ITERATIONS)
655 if end is not None:
656 n = MAX_ITERATIONS
657 else:
658 n = 1
660 dates = set()
661 counter = 0
663 # pass count = None to account for discrepancies with duplicates around DST
664 # boundaries
665 for next_date in self.to_rrule().xafter(start, count=None, inc=True):
666 next_date = create_datetime_instance(next_date).astimezone(
667 ZoneInfo(self.timezone)
668 )
670 # if the end date was exceeded, exit
671 if end and next_date > end:
672 break
674 # ensure no duplicates; weird things can happen with DST
675 if next_date not in dates:
676 dates.add(next_date)
677 yield next_date
679 # if enough dates have been collected or enough attempts were made, exit
680 if len(dates) >= n or counter > MAX_ITERATIONS:
681 break
683 counter += 1
686SCHEDULE_TYPES = Union[IntervalSchedule, CronSchedule, RRuleSchedule] 1a