Repository: PrefectHQ/prefect

Issue:
Weird daylight savings time behavior with croniter
As identified in a unit test in #325, cron seems to give an unexpected result during a DST transition. 

Thinking about DST is hard enough, let alone trying to figure out if this is a bug, so I've marked the test as `xfail` for now. It may be related to this issue: https://github.com/taichino/croniter/issues/116.

If this *is* a bug, then we should probably be converting all cron dates to UTC when the schedule is created (in fact, this might not be a bad idea in general).


Discussion:
Link to test: https://github.com/PrefectHQ/prefect/pull/325/files#diff-0be49ac061a72433791cd206d50b9675R101

Resolving patch:
diff --git a/src/prefect/schedules.py b/src/prefect/schedules.py
--- a/src/prefect/schedules.py
+++ b/src/prefect/schedules.py
@@ -1,18 +1,32 @@
 # Licensed under LICENSE.md; also available at https://www.prefect.io/licenses/alpha-eula
-
+import math
 import itertools
 from datetime import datetime, timedelta
-import pendulum
 from typing import Iterable, List
+
+import pendulum
+from crontab import CronTab
+
 from prefect.utilities.datetimes import ensure_tz_aware
-import croniter
 
 
 class Schedule:
     """
     Base class for Schedules
+
+    Args:
+        - start_date (datetime, optional): an optional start date for the schedule
+        - end_date (datetime, optional): an optional end date for the schedule
     """
 
+    def __init__(self, start_date: datetime = None, end_date: datetime = None):
+        if start_date is not None:
+            start_date = ensure_tz_aware(start_date)
+        if end_date is not None:
+            end_date = ensure_tz_aware(end_date)
+        self.start_date = start_date
+        self.end_date = end_date
+
     def next(self, n: int, on_or_after: datetime = None) -> List[datetime]:
         """
         Retrieve next scheduled dates.
@@ -39,16 +53,23 @@ class IntervalSchedule(Schedule):
     Args:
         - start_date (datetime): first date of schedule
         - interval (timedelta): interval on which this schedule occurs
+        - end_date (datetime, optional): an optional end date for the schedule
 
     Raises:
+        - TypeError: if start_date is not a datetime
         - ValueError: if provided interval is negative
     """
 
-    def __init__(self, start_date: datetime, interval: timedelta) -> None:
-        if interval.total_seconds() <= 0:
+    def __init__(
+        self, start_date: datetime, interval: timedelta, end_date: datetime = None
+    ) -> None:
+        if not isinstance(start_date, datetime):
+            raise TypeError("`start_date` must be a datetime.")
+        elif interval.total_seconds() <= 0:
             raise ValueError("Interval must be positive")
-        self.start_date = ensure_tz_aware(start_date)
+
         self.interval = interval
+        super().__init__(start_date=start_date, end_date=end_date)
 
     def next(self, n: int, on_or_after: datetime = None) -> List[datetime]:
         """
@@ -63,18 +84,26 @@ def next(self, n: int, on_or_after: datetime = None) -> List[datetime]:
         """
         if on_or_after is None:
             on_or_after = pendulum.now("utc")
+
         assert isinstance(on_or_after, datetime)  # mypy assertion
+        assert isinstance(self.start_date, datetime)  # mypy assertion
+
         on_or_after = ensure_tz_aware(on_or_after)
 
-        # infinite generator of all dates in the series
-        all_dates = (
-            self.start_date + i * self.interval
-            for i in itertools.count(start=0, step=1)
+        first_interval = math.ceil(
+            (on_or_after - self.start_date).total_seconds()
+            / self.interval.total_seconds()
         )
-        # filter generator for only dates on or after the requested date
-        upcoming_dates = filter(lambda d: d >= on_or_after, all_dates)  # type: ignore
-        # get the next n items from the generator
-        return list(itertools.islice(upcoming_dates, n))
+
+        dates = []
+
+        for i in range(n):
+            next_date = self.start_date + self.interval * (first_interval + i)
+            if self.end_date and next_date > self.end_date:
+                break
+            dates.append(next_date)
+
+        return dates
 
 
 class CronSchedule(Schedule):
@@ -83,12 +112,20 @@ class CronSchedule(Schedule):
 
     Args:
         - cron (str): a valid cron string
+        - start_date (datetime, optional): an optional start date for the schedule
+        - end_date (datetime, optional): an optional end date for the schedule
+
+    Raises:
+        - ValueError: if the cron string is invalid
     """
 
-    def __init__(self, cron: str) -> None:
+    def __init__(
+        self, cron: str, start_date: datetime = None, end_date: datetime = None
+    ) -> None:
         # build cron object to check the cron string - will raise an error if it's invalid
-        croniter.croniter(cron)
+        CronTab(cron)
         self.cron = cron
+        super().__init__(start_date=start_date, end_date=end_date)
 
     def next(self, n: int, on_or_after: datetime = None) -> List[datetime]:
         """
@@ -103,11 +140,25 @@ def next(self, n: int, on_or_after: datetime = None) -> List[datetime]:
         """
         if on_or_after is None:
             on_or_after = pendulum.now("utc")
+
+        if self.start_date is not None:
+            on_or_after = max(on_or_after, self.start_date)
+
         assert isinstance(on_or_after, datetime)  # mypy assertion
         on_or_after = ensure_tz_aware(on_or_after)
 
-        # croniter only supports >, not >=, so we subtract a second
-        on_or_after -= timedelta(seconds=1)
+        cron = CronTab(self.cron)
+
+        # subtract one second because we want to include on_or_after as a possible date
+        next_date = on_or_after - timedelta(seconds=1)  # type: pendulum.DateTime
+        dates = []
+
+        for i in range(n):
+            next_date += timedelta(
+                seconds=cron.next(next_date.in_tz("utc"), default_utc=True)
+            )
+            if self.end_date and next_date > self.end_date:
+                break
+            dates.append(next_date)
 
-        cron = croniter.croniter(self.cron, on_or_after)
-        return list(itertools.islice(cron.all_next(datetime), n))
+        return dates
diff --git a/src/prefect/serialization/schedule.py b/src/prefect/serialization/schedule.py
--- a/src/prefect/serialization/schedule.py
+++ b/src/prefect/serialization/schedule.py
@@ -18,6 +18,7 @@ class Meta:
         object_class = prefect.schedules.IntervalSchedule
 
     start_date = fields.DateTime(required=True)
+    end_date = fields.DateTime(allow_none=True)
     interval = fields.TimeDelta(required=True)
 
 
@@ -26,6 +27,8 @@ class CronScheduleSchema(VersionedSchema):
     class Meta:
         object_class = prefect.schedules.CronSchedule
 
+    start_date = fields.DateTime(allow_none=True)
+    end_date = fields.DateTime(allow_none=True)
     cron = fields.String(required=True)
 
 


Test patch:
diff --git a/tests/core/test_flow.py b/tests/core/test_flow.py
--- a/tests/core/test_flow.py
+++ b/tests/core/test_flow.py
@@ -1245,7 +1245,7 @@ def test_deserialization(self):
             name="hi",
             version="2",
             tasks=[p1, t2, t3],
-            schedule=prefect.schedules.CronSchedule("* * 0 0 0"),
+            schedule=prefect.schedules.CronSchedule("0 0 * * *"),
         )
         f.add_edge(p1, t2)
         f.add_edge(p1, t3)
diff --git a/tests/serialization/test_flows.py b/tests/serialization/test_flows.py
--- a/tests/serialization/test_flows.py
+++ b/tests/serialization/test_flows.py
@@ -46,7 +46,7 @@ class NewFlow(Flow):
 
 
 def test_deserialize_schedule():
-    schedule = prefect.schedules.CronSchedule("* * 0 0 0")
+    schedule = prefect.schedules.CronSchedule("0 0 * * *")
     f = Flow(schedule=schedule)
     serialized = FlowSchema().dump(f)
     deserialized = FlowSchema().load(serialized)
diff --git a/tests/serialization/test_schedules.py b/tests/serialization/test_schedules.py
--- a/tests/serialization/test_schedules.py
+++ b/tests/serialization/test_schedules.py
@@ -19,13 +19,19 @@
 @pytest.fixture()
 def interval_schedule():
     return schedules.IntervalSchedule(
-        interval=datetime.timedelta(hours=1), start_date=datetime.datetime(2020, 1, 1)
+        interval=datetime.timedelta(hours=1),
+        start_date=datetime.datetime(2020, 1, 1),
+        end_date=datetime.datetime(2020, 5, 1),
     )
 
 
 @pytest.fixture()
 def cron_schedule():
-    return schedules.CronSchedule(cron="0 0 * * *")
+    return schedules.CronSchedule(
+        cron="0 0 * * *",
+        start_date=datetime.datetime(2020, 1, 1),
+        end_date=datetime.datetime(2020, 5, 1),
+    )
 
 
 def test_all_schedules_have_serialization_schemas():
@@ -65,6 +71,8 @@ def test_serialize_cron_schedule(cron_schedule):
     assert schema.dump(cron_schedule) == {
         "cron": cron_schedule.cron,
         "__version__": __version__,
+        "start_date": cron_schedule.start_date.isoformat(),
+        "end_date": cron_schedule.end_date.isoformat(),
     }
 
 
@@ -72,6 +80,7 @@ def test_serialize_interval_schedule(interval_schedule):
     schema = schemas.IntervalScheduleSchema()
     assert schema.dump(interval_schedule) == {
         "start_date": interval_schedule.start_date.isoformat(),
+        "end_date": interval_schedule.end_date.isoformat(),
         "interval": interval_schedule.interval.total_seconds(),
         "__version__": __version__,
     }
diff --git a/tests/test_schedules.py b/tests/test_schedules.py
--- a/tests/test_schedules.py
+++ b/tests/test_schedules.py
@@ -26,108 +26,197 @@ def test_base_schedule_next_no_implemented():
         s.next(1)
 
 
-def test_create_interval_schedule():
-    assert schedules.IntervalSchedule(
-        start_date=pendulum.now("utc"), interval=timedelta(days=1)
-    )
-
-
-def test_interval_schedule_interval_must_be_positive():
-    with pytest.raises(ValueError):
-        schedules.IntervalSchedule(pendulum.now("utc"), interval=timedelta(hours=-1))
-
-
-def test_interval_schedule_next_n():
-
-    start_date = pendulum.datetime(2018, 1, 1)
-    today = pendulum.today("utc")
-    s = schedules.IntervalSchedule(start_date, timedelta(days=1))
-    assert s.next(3) == [today.add(days=1), today.add(days=2), today.add(days=3)]
-
-
-def test_interval_schedule_next_n_with_on_or_after_argument():
-    start_date = pendulum.datetime(2018, 1, 1)
-    today = pendulum.today("utc")
-    s = schedules.IntervalSchedule(start_date, timedelta(days=1))
-    assert s.next(3, on_or_after=start_date) == [
-        start_date,
-        start_date.add(days=1),
-        start_date.add(days=2),
-    ]
-
-
-def test_interval_schedule_daylight_savings_time():
-    dt = pendulum.datetime(2018, 11, 4, tz="America/New_York")
-    s = schedules.IntervalSchedule(dt, timedelta(hours=1))
-    next_4 = s.next(4, on_or_after=dt)
-    assert [t.hour for t in next_4] == [0, 1, 1, 2]
-
-
-def test_create_cron_schedule():
-    assert schedules.CronSchedule("* * * * *")
-
-
-def test_create_cron_schedule_with_invalid_cron_string_raises_error():
-    with pytest.raises(Exception):
-        schedules.CronSchedule("*")
-    with pytest.raises(Exception):
-        schedules.CronSchedule("hello world")
-    with pytest.raises(Exception):
-        schedules.CronSchedule(1)
-
-
-def test_cron_schedule_next_n():
-    every_day = "0 0 * * *"
-    s = schedules.CronSchedule(every_day)
-    assert s.next(3) == [
-        pendulum.today("utc").add(days=1),
-        pendulum.today("utc").add(days=2),
-        pendulum.today("utc").add(days=3),
-    ]
-
-
-def test_cron_schedule_next_n_with_on_or_after_argument():
-    every_day = "0 0 * * *"
-    s = schedules.CronSchedule(every_day)
-    start_date = pendulum.datetime(2018, 1, 1)
-
-    assert s.next(3, on_or_after=start_date) == [
-        start_date,
-        start_date.add(days=1),
-        start_date.add(days=2),
-    ]
-
-
-@pytest.mark.xfail(reason="Cron seems to have issues with DST")
-def test_cron_schedule_daylight_savings_time():
-    """
-    Cron's behavior is to skip the 2am hour altogether, which seems wrong??
+class TestIntervalSchedule:
+    def test_create_interval_schedule(self):
+        assert schedules.IntervalSchedule(
+            start_date=pendulum.now("utc"), interval=timedelta(days=1)
+        )
 
-    See also https://github.com/taichino/croniter/issues/116
-    """
-    dt = pendulum.datetime(2018, 11, 4, tz="America/New_York")
-    every_day = "0 * * * *"
-    s = schedules.CronSchedule(every_day)
-    next_4 = s.next(4, on_or_after=dt)
-    assert [t.hour for t in next_4] == [0, 1, 1, 2]
+    def test_start_date_must_be_datetime(self):
+        with pytest.raises(TypeError):
+            schedules.IntervalSchedule(start_date=None, interval=timedelta(hours=-1))
+
+    def test_interval_schedule_interval_must_be_positive(self):
+        with pytest.raises(ValueError):
+            schedules.IntervalSchedule(
+                pendulum.now("utc"), interval=timedelta(hours=-1)
+            )
+
+    def test_interval_schedule_next_n(self):
+        """Test that default on_or_after is *now*"""
+        start_date = pendulum.datetime(2018, 1, 1)
+        today = pendulum.today("utc")
+        s = schedules.IntervalSchedule(start_date, timedelta(days=1))
+        assert s.next(3) == [today.add(days=1), today.add(days=2), today.add(days=3)]
+
+    def test_interval_schedule_next_n_with_on_or_after_argument(self):
+        start_date = pendulum.datetime(2018, 1, 1)
+        today = pendulum.today("utc")
+        s = schedules.IntervalSchedule(start_date, timedelta(days=1))
+        assert s.next(3, on_or_after=today) == [
+            today,
+            today.add(days=1),
+            today.add(days=2),
+        ]
+
+    def test_interval_schedule_start_daylight_savings_time(self):
+        """
+        On 3/11/2018, at 2am, EST switched clocks forward an hour.
+        """
+        dt = pendulum.datetime(2018, 3, 11, tz="America/New_York")
+        s = schedules.IntervalSchedule(dt, timedelta(hours=1))
+        next_4 = s.next(4, on_or_after=dt)
+        assert [t.in_tz("utc").hour for t in next_4] == [5, 6, 7, 8]
+
+    def test_interval_schedule_daylight_savings_time(self):
+        """
+        11/4/2018, at 2am, EST switched clocks back an hour.
+        """
+        dt = pendulum.datetime(2018, 11, 4, tz="America/New_York")
+        s = schedules.IntervalSchedule(dt, timedelta(hours=1))
+        next_4 = s.next(4, on_or_after=dt)
+        assert [t.in_tz("utc").hour for t in next_4] == [4, 5, 6, 7]
+
+    def test_interval_schedule_end_date(self):
+        start_date = pendulum.datetime(2018, 1, 1)
+        end_date = pendulum.datetime(2018, 1, 2)
+        s = schedules.IntervalSchedule(
+            start_date=start_date, interval=timedelta(days=1), end_date=end_date
+        )
+        assert s.next(3, on_or_after=start_date) == [start_date, start_date.add(days=1)]
+        assert s.next(3, on_or_after=pendulum.datetime(2018, 2, 1)) == []
+
+    def test_interval_schedule_respects_on_or_after_in_middle_of_interval(self):
+        """
+        If the "on_or_after" date is in the middle of an interval, then the IntervalSchedule
+        should advance to the next interval.
+        """
+        start_date = pendulum.datetime(2018, 1, 1)
+        s = schedules.IntervalSchedule(
+            start_date=start_date, interval=timedelta(hours=1)
+        )
+        assert s.next(2, on_or_after=start_date + timedelta(microseconds=1)) == [
+            start_date.add(hours=1),
+            start_date.add(hours=2),
+        ]
+
+
+class TestCronSchedule:
+    def test_create_cron_schedule(self):
+        assert schedules.CronSchedule("* * * * *")
+
+    def test_create_cron_schedule_with_invalid_cron_string_raises_error(self):
+        with pytest.raises(Exception):
+            schedules.CronSchedule("*")
+        with pytest.raises(Exception):
+            schedules.CronSchedule("hello world")
+        with pytest.raises(Exception):
+            schedules.CronSchedule(1)
+        with pytest.raises(Exception):
+            schedules.CronSchedule("* * 0 0 0")
+
+    def test_cron_schedule_next_n(self):
+        every_day = "0 0 * * *"
+        s = schedules.CronSchedule(every_day)
+        assert s.next(3) == [
+            pendulum.today("utc").add(days=1),
+            pendulum.today("utc").add(days=2),
+            pendulum.today("utc").add(days=3),
+        ]
+
+    def test_cron_schedule_next_n_with_on_or_after_argument(self):
+        every_day = "0 0 * * *"
+        s = schedules.CronSchedule(every_day)
+        start_date = pendulum.datetime(2018, 1, 1)
+
+        assert s.next(3, on_or_after=start_date) == [
+            start_date,
+            start_date.add(days=1),
+            start_date.add(days=2),
+        ]
+
+    def test_cron_schedule_start_daylight_savings_time(self):
+        """
+        On 3/11/2018, at 2am, EST switched clocks forward an hour.
+        """
+        dt = pendulum.datetime(2018, 3, 11, tz="America/New_York")
+        every_hour = "0 * * * *"
+        s = schedules.CronSchedule(every_hour)
+        next_4 = s.next(4, on_or_after=dt)
+
+        assert [t.in_tz("utc").hour for t in next_4] == [5, 6, 7, 8]
+
+    def test_cron_schedule_end_daylight_savings_time(self):
+        """
+        11/4/2018, at 2am, EST switched clocks back an hour.
+        """
+        dt = pendulum.datetime(2018, 11, 4, tz="America/New_York")
+        every_hour = "0 * * * *"
+        s = schedules.CronSchedule(every_hour)
+        next_4 = s.next(4, on_or_after=dt)
+        assert [t.in_tz("utc").hour for t in next_4] == [4, 5, 6, 7]
+
+    def test_cron_schedule_start_date(self):
+        every_day = "0 0 * * *"
+        start_date = pendulum.datetime(2050, 1, 1)
+        s = schedules.CronSchedule(every_day, start_date=start_date)
+        assert s.next(3) == [start_date, start_date.add(days=1), start_date.add(days=2)]
+
+    def test_cron_schedule_end_date(self):
+        every_day = "0 0 * * *"
+        start_date = pendulum.datetime(2050, 1, 1)
+        end_date = pendulum.datetime(2050, 1, 2)
+        s = schedules.CronSchedule(every_day, start_date=start_date, end_date=end_date)
+        assert s.next(3) == [start_date, start_date.add(days=1)]
 
 
 class TestSerialization:
     def test_serialize_cron_schedule(self):
         schedule = schedules.CronSchedule("0 0 * * *")
         assert schedule.serialize() == {
+            "end_date": None,
+            "start_date": None,
+            "type": "CronSchedule",
+            "cron": "0 0 * * *",
+            "__version__": __version__,
+        }
+
+    def test_serialize_cron_schedule_with_start_and_end_dates(self):
+        sd = pendulum.datetime(2000, 1, 1)
+        ed = pendulum.datetime(2010, 1, 1)
+        schedule = schedules.CronSchedule("0 0 * * *", start_date=sd, end_date=ed)
+        assert schedule.serialize() == {
+            "end_date": str(ed),
+            "start_date": str(sd),
             "type": "CronSchedule",
             "cron": "0 0 * * *",
             "__version__": __version__,
         }
 
     def test_serialize_interval_schedule(self):
+        sd = pendulum.datetime(2000, 1, 1)
+        schedule = schedules.IntervalSchedule(
+            interval=timedelta(hours=1), start_date=sd
+        )
+        assert schedule.serialize() == {
+            "type": "IntervalSchedule",
+            "start_date": str(sd),
+            "end_date": None,
+            "interval": 3600,
+            "__version__": __version__,
+        }
+
+    def test_serialize_interval_schedule_with_end_date(self):
+        sd = pendulum.datetime(2000, 1, 1)
+        ed = pendulum.datetime(2010, 1, 1)
+
         schedule = schedules.IntervalSchedule(
-            interval=timedelta(hours=1), start_date=START_DATE
+            interval=timedelta(hours=1), start_date=sd, end_date=ed
         )
         assert schedule.serialize() == {
             "type": "IntervalSchedule",
-            "start_date": "2018-01-01T00:00:00+00:00",
+            "start_date": str(sd),
+            "end_date": str(ed),
             "interval": 3600,
             "__version__": __version__,
         }