Fix: User-provided plan end must take precedence over the max interval end per model (#3042)

Author: izeigerman

sqlmesh/core/context.py

@@ -1134,7 +1134,7 @@ def plan_builder(
         # to prevent unintended evaluation of the entire DAG.
         default_end: t.Optional[int] = None
         max_interval_end_per_model: t.Optional[t.Dict[str, int]] = None
-        if not run:
+        if not run and not end:
             models_for_interval_end: t.Optional[t.Set[str]] = None
             if backfill_models is not None:
                 models_for_interval_end = set()

sqlmesh/core/plan/builder.py

@@ -257,6 +257,12 @@ def build(self) -> Plan:
             dag, earliest_interval_start(filtered_snapshots.values())
         )
 
+        interval_end_per_model = self._interval_end_per_model
+        if interval_end_per_model and self.override_end:
+            # If the end date was provided explicitly by a user, then interval end for each individual
+            # model should be ignored.
+            interval_end_per_model = None
+
         plan = Plan(
             context_diff=self._context_diff,
             plan_id=self._plan_id,
@@ -275,7 +281,7 @@ def build(self) -> Plan:
             ignored=ignored,
             deployability_index=deployability_index,
             restatements=restatements,
-            interval_end_per_model=self._interval_end_per_model,
+            interval_end_per_model=interval_end_per_model,
             selected_models_to_backfill=self._backfill_models,
             models_to_backfill=models_to_backfill,
             effective_from=self._effective_from,

sqlmesh/core/snapshot/definition.py

@@ -1628,9 +1628,14 @@ def missing_intervals(
     for snapshot in snapshots:
         if not snapshot.evaluatable:
             continue
-        interval = restatements.get(snapshot.snapshot_id)
         snapshot_start_date = start_dt
-        snapshot_end_date = interval_end_per_model.get(snapshot.name, end_date)
+        snapshot_end_date: TimeLike = end_date
+
+        existing_interval_end = interval_end_per_model.get(snapshot.name)
+        if existing_interval_end and existing_interval_end > to_timestamp(snapshot_start_date):
+            snapshot_end_date = existing_interval_end
+
+        interval = restatements.get(snapshot.snapshot_id)
         if interval:
             snapshot_start_date, snapshot_end_date = (to_datetime(i) for i in interval)
             snapshot = snapshot.copy()

tests/core/test_integration.py

@@ -1258,6 +1258,22 @@ def test_select_unchanged_model_for_backfill(init_and_plan_context: t.Callable):
     assert {o.name for o in schema_objects} == {"waiter_revenue_by_day", "top_waiters"}
 
 
+@freeze_time("2023-01-08 15:00:00")
+def test_max_interval_end_per_model_not_applied_when_end_is_provided(
+    init_and_plan_context: t.Callable,
+):
+    context, plan = init_and_plan_context("examples/sushi")
+    context.apply(plan)
+
+    with freeze_time("2023-01-09 00:00:00"):
+        context.run()
+
+        plan = context.plan(
+            no_prompts=True, restate_models=["*"], start="2023-01-09", end="2023-01-09"
+        )
+        context.apply(plan)
+
+
 @freeze_time("2023-01-08 15:00:00")
 def test_select_models_for_backfill(init_and_plan_context: t.Callable):
     context, _ = init_and_plan_context("examples/sushi")

tests/core/test_plan.py

@@ -2423,3 +2423,45 @@ def test_plan_start_when_preview_enabled(make_snapshot, mocker: MockerFixture):
         enable_preview=True,
     )
     assert plan_builder.build().start == default_start_for_preview
+
+
+def test_interval_end_per_model(make_snapshot):
+    snapshot = make_snapshot(SqlModel(name="a", query=parse_one("select 1, ds")))
+    snapshot.categorize_as(SnapshotChangeCategory.BREAKING)
+
+    new_snapshot = make_snapshot(SqlModel(name="a", query=parse_one("select 2, ds")))
+
+    context_diff = ContextDiff(
+        environment="test_environment",
+        is_new_environment=True,
+        is_unfinalized_environment=False,
+        normalize_environment_name=True,
+        create_from="prod",
+        added=set(),
+        removed_snapshots={},
+        modified_snapshots={new_snapshot.name: (new_snapshot, snapshot)},
+        snapshots={new_snapshot.snapshot_id: new_snapshot},
+        new_snapshots={new_snapshot.snapshot_id: new_snapshot},
+        previous_plan_id=None,
+        previously_promoted_snapshot_ids=set(),
+        previous_finalized_snapshots=None,
+    )
+
+    plan_builder = PlanBuilder(
+        context_diff,
+        DuckDBEngineAdapter.SCHEMA_DIFFER,
+        interval_end_per_model={snapshot.name: to_timestamp("2023-01-09")},
+    )
+    assert plan_builder.build().interval_end_per_model == {
+        snapshot.name: to_timestamp("2023-01-09")
+    }
+
+    # User-provided end should take precedence.
+    plan_builder = PlanBuilder(
+        context_diff,
+        DuckDBEngineAdapter.SCHEMA_DIFFER,
+        interval_end_per_model={snapshot.name: to_timestamp("2023-01-09")},
+        end="2023-01-10",
+        is_dev=True,
+    )
+    assert plan_builder.build().interval_end_per_model is None

tests/core/test_snapshot.py

@@ -1989,3 +1989,60 @@ def test_snapshot_pickle_intervals(make_snapshot):
     assert not loaded_snapshot.dev_intervals
     assert len(snapshot.intervals) > 0
     assert len(snapshot.dev_intervals) > 0
+
+
+def test_missing_intervals_interval_end_per_model(make_snapshot):
+    snapshot_a = make_snapshot(
+        SqlModel(
+            name="a",
+            start="2023-01-04",
+            query=parse_one("SELECT 1"),
+        ),
+        version="a",
+    )
+
+    snapshot_b = make_snapshot(
+        SqlModel(
+            name="b",
+            start="2023-01-04",
+            query=parse_one("SELECT 2"),
+        ),
+        version="b",
+    )
+
+    assert missing_intervals(
+        [snapshot_a, snapshot_b],
+        start="2023-01-04",
+        end="2023-01-10",
+        interval_end_per_model={
+            snapshot_a.name: to_timestamp("2023-01-09"),
+            snapshot_b.name: to_timestamp("2023-01-06"),
+        },
+    ) == {
+        snapshot_a: [
+            (to_timestamp("2023-01-04"), to_timestamp("2023-01-05")),
+            (to_timestamp("2023-01-05"), to_timestamp("2023-01-06")),
+            (to_timestamp("2023-01-06"), to_timestamp("2023-01-07")),
+            (to_timestamp("2023-01-07"), to_timestamp("2023-01-08")),
+            (to_timestamp("2023-01-08"), to_timestamp("2023-01-09")),
+        ],
+        snapshot_b: [
+            (to_timestamp("2023-01-04"), to_timestamp("2023-01-05")),
+            (to_timestamp("2023-01-05"), to_timestamp("2023-01-06")),
+        ],
+    }
+
+    assert missing_intervals(
+        [snapshot_a, snapshot_b],
+        start="2023-01-08",
+        end="2023-01-08",
+        interval_end_per_model={
+            snapshot_a.name: to_timestamp("2023-01-09"),
+            snapshot_b.name: to_timestamp(
+                "2023-01-06"
+            ),  # The interval end is before the start. This should be ignored.
+        },
+    ) == {
+        snapshot_a: [(to_timestamp("2023-01-08"), to_timestamp("2023-01-09"))],
+        snapshot_b: [(to_timestamp("2023-01-08"), to_timestamp("2023-01-09"))],
+    }