ref(crons): Move schedule logic to its own module with typing (#57050)

Moves this logic into it's own file and significantly improves typing of
the schedule config

src/sentry/monitors/models.py

@@ -2,13 +2,11 @@ from __future__ import annotations
 
 import logging
 from datetime import datetime, timedelta
-from typing import TYPE_CHECKING, Any, Dict, Optional
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
 from uuid import uuid4
 
 import jsonschema
 import pytz
-from croniter import croniter
-from dateutil import rrule
 from django.conf import settings
 from django.db import models
 from django.db.models.signals import pre_save
@@ -33,6 +31,7 @@ from sentry.db.models import (
 from sentry.db.models.utils import slugify_instance
 from sentry.locks import locks
 from sentry.models import Environment, Rule, RuleSource
+from sentry.monitors.types import CrontabSchedule, IntervalSchedule
 from sentry.utils.retries import TimedRetryPolicy
 
 logger = logging.getLogger(__name__)
@@ -40,15 +39,6 @@ logger = logging.getLogger(__name__)
 if TYPE_CHECKING:
     from sentry.models import Project
 
-SCHEDULE_INTERVAL_MAP = {
-    "year": rrule.YEARLY,
-    "month": rrule.MONTHLY,
-    "week": rrule.WEEKLY,
-    "day": rrule.DAILY,
-    "hour": rrule.HOURLY,
-    "minute": rrule.MINUTELY,
-}
-
 MONITOR_CONFIG = {
     "type": "object",
     "properties": {
@@ -82,43 +72,6 @@ class MonitorEnvironmentValidationFailed(Exception):
     pass
 
 
-def get_next_schedule(reference_ts: datetime, schedule_type: int, schedule):
-    """
-    Given the schedule type and schedule, determine the next timestamp for a
-    schedule from the reference_ts
-
-    Examples:
-
-    >>> get_next_schedule('05:30', ScheduleType.CRONTAB, '0 * * * *')
-    >>> 06:00
-
-    >>> get_next_schedule('05:30', ScheduleType.CRONTAB, '30 * * * *')
-    >>> 06:30
-
-    >>> get_next_schedule('05:35', ScheduleType.INTERVAL, [2, 'hour'])
-    >>> 07:35
-    """
-    if schedule_type == ScheduleType.CRONTAB:
-        next_schedule = croniter(schedule, reference_ts).get_next(datetime)
-    elif schedule_type == ScheduleType.INTERVAL:
-        interval, unit_name = schedule
-        rule = rrule.rrule(
-            freq=SCHEDULE_INTERVAL_MAP[unit_name],
-            interval=interval,
-            dtstart=reference_ts,
-            count=2,
-        )
-        next_schedule = rule.after(reference_ts)
-    else:
-        raise NotImplementedError("unknown schedule_type")
-
-    # Ensure we clamp the expected time down to the minute, that is the level
-    # of granularity we're able to support
-    next_schedule = next_schedule.replace(second=0, microsecond=0)
-
-    return next_schedule
-
-
 class MonitorStatus:
     """
     The monitor status is an extension of the ObjectStatus constants. In this
@@ -264,6 +217,18 @@ class Monitor(Model):
                 )
         return super().save(*args, **kwargs)
 
+    @property
+    def schedule(self) -> Union[CrontabSchedule, IntervalSchedule]:
+        schedule_type = self.config["schedule_type"]
+        schedule = self.config["schedule"]
+
+        if schedule_type == ScheduleType.CRONTAB:
+            return CrontabSchedule(crontab=schedule)
+        if schedule_type == ScheduleType.INTERVAL:
+            return IntervalSchedule(interval=schedule[0], unit=schedule[1])
+
+        raise NotImplementedError("unknown schedule_type")
+
     def get_schedule_type_display(self):
         return ScheduleType.get_name(self.config["schedule_type"])
 
@@ -274,12 +239,10 @@ class Monitor(Model):
         """
         Computes the next expected checkin time given the most recent checkin time
         """
+        from sentry.monitors.schedule import get_next_schedule
+
         tz = pytz.timezone(self.config.get("timezone") or "UTC")
-        schedule_type = self.config.get("schedule_type", ScheduleType.CRONTAB)
-        next_checkin = get_next_schedule(
-            last_checkin.astimezone(tz), schedule_type, self.config["schedule"]
-        )
-        return next_checkin
+        return get_next_schedule(last_checkin.astimezone(tz), self.schedule)
 
     def get_next_expected_checkin_latest(self, last_checkin: datetime) -> datetime:
         """

src/sentry/monitors/schedule.py

@@ -0,0 +1,54 @@
+from datetime import datetime
+from typing import Dict
+
+from croniter import croniter
+from dateutil import rrule
+
+from sentry.monitors.types import IntervalUnit, ScheduleConfig
+
+SCHEDULE_INTERVAL_MAP: Dict[IntervalUnit, int] = {
+    "year": rrule.YEARLY,
+    "month": rrule.MONTHLY,
+    "week": rrule.WEEKLY,
+    "day": rrule.DAILY,
+    "hour": rrule.HOURLY,
+    "minute": rrule.MINUTELY,
+}
+
+
+def get_next_schedule(
+    reference_ts: datetime,
+    schedule: ScheduleConfig,
+):
+    """
+    Given the schedule type and schedule, determine the next timestamp for a
+    schedule from the reference_ts
+
+    Examples:
+
+    >>> get_next_schedule('05:30', CrontabSchedule('0 * * * *'))
+    >>> 06:00
+
+    >>> get_next_schedule('05:30', CrontabSchedule('30 * * * *'))
+    >>> 06:30
+
+    >>> get_next_schedule('05:35', IntervalSchedule(interval=2, unit='hour'))
+    >>> 07:35
+    """
+    # Ensure we clamp the expected time down to the minute, that is the level
+    # of granularity we're able to support
+
+    if schedule.type == "crontab":
+        iterator = croniter(schedule.crontab, reference_ts)
+        return iterator.get_next(datetime).replace(second=0, microsecond=0)
+
+    if schedule.type == "interval":
+        rule = rrule.rrule(
+            freq=SCHEDULE_INTERVAL_MAP[schedule.unit],
+            interval=schedule.interval,
+            dtstart=reference_ts,
+            count=2,
+        )
+        return rule.after(reference_ts).replace(second=0, microsecond=0)
+
+    raise NotImplementedError("unknown schedule_type")

src/sentry/monitors/types.py

@@ -1,4 +1,5 @@
-from typing import Dict, Literal, TypedDict
+from dataclasses import dataclass
+from typing import Dict, Literal, TypedDict, Union
 
 from typing_extensions import NotRequired
 
@@ -33,3 +34,22 @@ class CheckinPayload(TypedDict):
     duration: NotRequired[int]
     monitor_config: NotRequired[Dict]
     contexts: NotRequired[CheckinContexts]
+
+
+IntervalUnit = Literal["year", "month", "week", "day", "hour", "minute"]
+
+
+@dataclass
+class CrontabSchedule:
+    crontab: str
+    type: Literal["crontab"] = "crontab"
+
+
+@dataclass
+class IntervalSchedule:
+    interval: int
+    unit: IntervalUnit
+    type: Literal["interval"] = "interval"
+
+
+ScheduleConfig = Union[CrontabSchedule, IntervalSchedule]

tests/sentry/monitors/test_schedule.py

@@ -0,0 +1,27 @@
+from datetime import datetime
+
+from django.utils import timezone
+
+from sentry.monitors.schedule import get_next_schedule
+from sentry.monitors.types import CrontabSchedule, IntervalSchedule
+
+
+def test_get_next_schedule():
+    ts = datetime(2019, 1, 1, 5, 30, 0, tzinfo=timezone.utc)
+
+    # 00 * * * *: 5:30 -> 6:00
+    expect = ts.replace(hour=6, minute=0)
+    assert get_next_schedule(ts, CrontabSchedule("0 * * * *")) == expect
+
+    # 30 * * * *: 5:30 -> 6:30
+    expect = ts.replace(hour=6, minute=30)
+    assert get_next_schedule(ts, CrontabSchedule("30 * * * *")) == expect
+
+    # 2 hour interval: 5:30 -> 7:30
+    expect = ts.replace(hour=7, minute=30)
+    assert get_next_schedule(ts, IntervalSchedule(interval=2, unit="hour")) == expect
+
+    # 2 hour interval: 5:42 -> 7:42
+    interval_ts = ts.replace(hour=5, minute=42)
+    expect = ts.replace(hour=7, minute=42)
+    assert get_next_schedule(interval_ts, IntervalSchedule(interval=2, unit="hour")) == expect