Fix: Inference of plan start and end (#1519)

Author: izeigerman

sqlmesh/core/plan/definition.py

@@ -31,9 +31,7 @@
 from sqlmesh.utils.dag import DAG
 from sqlmesh.utils.date import (
     TimeLike,
-    make_inclusive_end,
     now,
-    to_date,
     to_datetime,
     to_timestamp,
     validate_date_range,
@@ -144,7 +142,13 @@ def uncategorized(self) -> t.List[Snapshot]:
     def start(self) -> TimeLike:
         """Returns the start of the plan or the earliest date of all snapshots."""
         if not self.override_start and not self._missing_intervals:
-            return earliest_start_date(self.snapshots)
+            earliest_start = earliest_start_date(self.snapshots)
+            earliest_interval_starts = [s.intervals[0][0] for s in self.snapshots if s.intervals]
+            return (
+                min(earliest_start, to_datetime(min(earliest_interval_starts)))
+                if earliest_interval_starts
+                else earliest_start
+            )
         return self._start or (
             min(
                 start
@@ -166,40 +170,10 @@ def set_start(self, new_start: TimeLike) -> None:
         self.__missing_intervals = None
         self._refresh_dag_and_ignored_snapshots()
 
-    @classmethod
-    def _get_end_date(cls, end_and_units: t.List[t.Tuple[int, IntervalUnit]]) -> TimeLike:
-        if end_and_units:
-            end, unit = max(end_and_units)
-
-            if unit.is_date_granularity:
-                return to_date(make_inclusive_end(end))
-            return end
-        return now()
-
     @property
     def end(self) -> TimeLike:
         """Returns the end of the plan or now."""
-        if not self._end or not self.override_end:
-            if self._missing_intervals:
-                return self._get_end_date(
-                    [
-                        (end, snapshot.node.interval_unit)
-                        for snapshot in self.snapshots
-                        if (snapshot.name, snapshot.version_get_or_generate())
-                        in self._missing_intervals
-                        for _, end in self._missing_intervals[
-                            (snapshot.name, snapshot.version_get_or_generate())
-                        ]
-                    ]
-                )
-            return self._get_end_date(
-                [
-                    (snapshot.intervals[-1][1], snapshot.node.interval_unit)
-                    for snapshot in self.snapshots
-                    if snapshot.intervals
-                ]
-            )
-        return self._end
+        return self._end or now()
 
     @end.setter
     def end(self, new_end: TimeLike) -> None:

sqlmesh/core/snapshot/definition.py

@@ -1208,8 +1208,14 @@ def missing_intervals(
     """Returns all missing intervals given a collection of snapshots."""
     missing = {}
     cache: t.Dict[str, datetime] = {}
-    start_dt = to_datetime(start or earliest_start_date(snapshots, cache))
     end_date = end or now()
+    start_dt = (
+        to_datetime(start)
+        if start
+        else earliest_start_date(
+            snapshots, cache, default_value=to_date(end_date) - timedelta(days=1)
+        )
+    )
     restatements = restatements or {}
 
     for snapshot in snapshots:
@@ -1239,7 +1245,9 @@ def missing_intervals(
 
 
 def earliest_start_date(
-    snapshots: t.Iterable[Snapshot], cache: t.Optional[t.Dict[str, datetime]] = None
+    snapshots: t.Iterable[Snapshot],
+    cache: t.Optional[t.Dict[str, datetime]] = None,
+    default_value: t.Optional[TimeLike] = None,
 ) -> datetime:
     """Get the earliest start date from a collection of snapshots.
 
@@ -1251,7 +1259,7 @@ def earliest_start_date(
     """
     cache = {} if cache is None else cache
     snapshots = list(snapshots)
-    earliest = to_datetime(yesterday().date())
+    earliest = to_datetime(default_value or yesterday().date())
     if snapshots:
         return min(start_date(snapshot, snapshots, cache) or earliest for snapshot in snapshots)
     return earliest

tests/core/test_context.py

@@ -20,7 +20,7 @@
 from sqlmesh.core.environment import Environment
 from sqlmesh.core.model import load_sql_based_model
 from sqlmesh.core.plan import BuiltInPlanEvaluator, Plan
-from sqlmesh.utils.date import now, to_date, yesterday_ds
+from sqlmesh.utils.date import make_inclusive_end, now, to_date, yesterday_ds
 from sqlmesh.utils.errors import ConfigError
 from tests.utils.test_filesystem import create_temp_file
 
@@ -438,12 +438,12 @@ def test_plan_default_end(sushi_context_pre_scheduling: Context):
     dev_plan = sushi_context_pre_scheduling.plan(
         "test_env", no_prompts=True, include_unmodified=True, skip_backfill=True, auto_apply=True
     )
-    assert dev_plan.end == plan_end
+    assert to_date(make_inclusive_end(dev_plan.end)) == plan_end
 
     forward_only_dev_plan = sushi_context_pre_scheduling.plan(
         "test_env_forward_only", no_prompts=True, include_unmodified=True, forward_only=True
     )
-    assert forward_only_dev_plan.end == plan_end
+    assert to_date(make_inclusive_end(forward_only_dev_plan.end)) == plan_end
     assert forward_only_dev_plan._start == plan_end
 
 

tests/core/test_plan.py

@@ -15,7 +15,7 @@
     SnapshotFingerprint,
 )
 from sqlmesh.utils.dag import DAG
-from sqlmesh.utils.date import now, to_date, to_ds, to_timestamp
+from sqlmesh.utils.date import now, to_date, to_datetime, to_timestamp
 from sqlmesh.utils.errors import PlanError
 
 
@@ -63,8 +63,8 @@ def test_forward_only_dev(make_snapshot, mocker: MockerFixture):
     )
 
     expected_start = to_date("2022-01-01")
-    expected_end = to_date("2022-01-02")
-    expected_interval_end = to_timestamp(to_date("2022-01-03"))
+    expected_end = to_date("2022-01-03")
+    expected_interval_end = to_timestamp(to_date("2022-01-04"))
 
     context_diff_mock = mocker.Mock()
     context_diff_mock.snapshots = {"a": snapshot_a}
@@ -342,6 +342,12 @@ def test_start_inference(make_snapshot, mocker: MockerFixture):
     assert plan.missing_intervals[0].snapshot_name == snapshot_a.name
     assert plan.start == to_timestamp("2022-01-01")
 
+    # Test inference from existing intervals
+    context_diff_mock.snapshots = {"b": snapshot_b}
+    plan = Plan(context_diff_mock)
+    assert not plan.missing_intervals
+    assert plan.start == to_datetime("2022-01-01")
+
 
 def test_auto_categorization(make_snapshot, mocker: MockerFixture):
     snapshot = make_snapshot(SqlModel(name="a", query=parse_one("select 1, ds")))
@@ -394,33 +400,6 @@ def test_auto_categorization_missing_schema_downstream(make_snapshot, mocker: Mo
     assert updated_snapshot.change_category == SnapshotChangeCategory.BREAKING
 
 
-def test_end_from_missing_instead_of_now(make_snapshot, mocker: MockerFixture):
-    snapshot_a = make_snapshot(
-        SqlModel(
-            name="a",
-            query=parse_one("select 1, ds"),
-            kind=IncrementalByTimeRangeKind(time_column="ds"),
-        )
-    )
-
-    context_diff_mock = mocker.Mock()
-    context_diff_mock.snapshots = {"a": snapshot_a}
-    context_diff_mock.added = set()
-    context_diff_mock.removed_snapshots = set()
-    context_diff_mock.modified_snapshots = {}
-    context_diff_mock.new_snapshots = {snapshot_a.snapshot_id: snapshot_a}
-
-    start_mock = mocker.patch("sqlmesh.core.snapshot.definition.earliest_start_date")
-    start_mock.return_value = to_ds("2022-01-01")
-    now_mock = mocker.patch("sqlmesh.core.plan.definition.now")
-    now_mock.return_value = to_ds("2022-01-30")
-    snapshot_a.add_interval("2022-01-01", "2022-01-05")
-
-    plan = Plan(context_diff_mock, is_dev=True)
-    assert plan.start == to_timestamp("2022-01-06")
-    assert plan.end == to_date("2022-01-29")
-
-
 def test_broken_references(make_snapshot, mocker: MockerFixture):
     snapshot_b = make_snapshot(SqlModel(name="b", query=parse_one("select 2, ds FROM a")))
     snapshot_b.categorize_as(SnapshotChangeCategory.BREAKING)