Repository: TobikoData/sqlmesh
Pull request #1980 (2024-01-17 21:35:06)

Fix: Normalize max end interval by snapshot when models are selected explicitly
This update ensures that the latest intervals are normalized by snapshot and the minimum value is picked to guarantee that no unexpected backfills get triggered during plan application with explicit model selection.

This also solves the issue of having unexpected backfills when a newly added model is explicitly selected with `--select-model`, in which case there would be no existing intervals to get the max value from and the `now` is used for the plan end as a result, causing unrelated models to be backfilled.


Patch:
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 23632031..82b3cde7 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -956,23 +956,45 @@ class Context(BaseContext):
                     f"Provided restated models do not match any models. No models will be included in plan. Provided: {', '.join(restate_models)}"
                 )
 
+        snapshots = self._snapshots(models_override)
+        context_diff = self._context_diff(
+            environment or c.PROD,
+            snapshots=snapshots,
+            create_from=create_from,
+            force_no_diff=(restate_models is not None and not expanded_restate_models)
+            or (backfill_models is not None and not backfill_models),
+        )
+
         # If no end date is specified, use the max interval end from prod
         # to prevent unintended evaluation of the entire DAG.
-        default_end = (
-            self.state_sync.max_interval_end_for_environment(c.PROD, models=backfill_models)
-            if not run
-            else None
-        )
+        if not run:
+            if backfill_models is not None:
+                # Only consider selected models for the default end value.
+                models_for_default_end = backfill_models.copy()
+                for name in backfill_models:
+                    if name not in snapshots:
+                        continue
+                    snapshot = snapshots[name]
+                    snapshot_id = snapshot.snapshot_id
+                    if (
+                        snapshot_id in context_diff.added
+                        and snapshot_id in context_diff.new_snapshots
+                    ):
+                        # If the selected model is a newly added model, then we should narrow down the intervals
+                        # that should be considered for the default plan end value by including its parents.
+                        models_for_default_end |= {s.name for s in snapshot.parents}
+                default_end = self.state_sync.greatest_common_interval_end(
+                    c.PROD, models_for_default_end
+                )
+            else:
+                default_end = self.state_sync.max_interval_end_for_environment(c.PROD)
+        else:
+            default_end = None
+
         default_start = to_date(default_end) - timedelta(days=1) if default_end and is_dev else None
 
         return PlanBuilder(
-            context_diff=self._context_diff(
-                environment or c.PROD,
-                snapshots=self._snapshots(models_override),
-                create_from=create_from,
-                force_no_diff=(restate_models is not None and not expanded_restate_models)
-                or (backfill_models is not None and not backfill_models),
-            ),
+            context_diff=context_diff,
             start=start,
             end=end,
             execution_time=execution_time,
diff --git a/sqlmesh/core/state_sync/base.py b/sqlmesh/core/state_sync/base.py
index 37cc3087..c68c1025 100644
--- a/sqlmesh/core/state_sync/base.py
+++ b/sqlmesh/core/state_sync/base.py
@@ -138,15 +138,23 @@ class StateReader(abc.ABC):
         """
 
     @abc.abstractmethod
-    def max_interval_end_for_environment(
-        self, environment: str, models: t.Optional[t.Set[str]] = None
-    ) -> t.Optional[int]:
+    def max_interval_end_for_environment(self, environment: str) -> t.Optional[int]:
         """Returns the max interval end for the given environment.
 
         Args:
             environment: The environment.
-            models: The optional model FQNs to pick intervals from. If not provided all
-                models in the environment will be used.
+
+        Returns:
+            A timestamp or None if no interval or environment exists.
+        """
+
+    @abc.abstractmethod
+    def greatest_common_interval_end(self, environment: str, models: t.Set[str]) -> t.Optional[int]:
+        """Returns the greatest common interval end for given models in the target environment.
+
+        Args:
+            environment: The environment.
+            models: The model FQNs to select intervals from.
 
         Returns:
             A timestamp or None if no interval or environment exists.
diff --git a/sqlmesh/core/state_sync/engine_adapter.py b/sqlmesh/core/state_sync/engine_adapter.py
index c176d117..2e2f6209 100644
--- a/sqlmesh/core/state_sync/engine_adapter.py
+++ b/sqlmesh/core/state_sync/engine_adapter.py
@@ -628,18 +628,13 @@ class EngineAdapterStateSync(CommonStateSyncMixin, StateSync):
             s.dev_intervals = []
         return Snapshot.hydrate_with_intervals_by_version(snapshots, intervals)
 
-    def max_interval_end_for_environment(
-        self, environment: str, models: t.Optional[t.Set[str]] = None
-    ) -> t.Optional[int]:
+    def max_interval_end_for_environment(self, environment: str) -> t.Optional[int]:
         env = self._get_environment(environment)
         if not env:
             return None
 
-        snapshots = (
-            [s for s in env.snapshots if s.name in models] if models is not None else env.snapshots
-        )
         max_end = None
-        for where in self._snapshot_name_version_filter(snapshots, "intervals"):
+        for where in self._snapshot_name_version_filter(env.snapshots, "intervals"):
             end = self.engine_adapter.fetchone(
                 exp.select(exp.func("MAX", exp.to_column("end_ts")))
                 .from_(exp.to_table(self.intervals_table).as_("intervals"))
@@ -655,6 +650,49 @@ class EngineAdapterStateSync(CommonStateSyncMixin, StateSync):
 
         return max_end
 
+    def greatest_common_interval_end(self, environment: str, models: t.Set[str]) -> t.Optional[int]:
+        if not models:
+            return None
+
+        env = self._get_environment(environment)
+        if not env:
+            return None
+
+        snapshots = [s for s in env.snapshots if s.name in models]
+        if not snapshots:
+            snapshots = env.snapshots
+
+        greatest_common_end = None
+
+        table_alias = "intervals"
+        name_col = exp.column("name", table=table_alias)
+        version_col = exp.column("version", table=table_alias)
+
+        for where in self._snapshot_name_version_filter(snapshots, table_alias):
+            max_end_subquery = (
+                exp.select(
+                    name_col,
+                    version_col,
+                    exp.func("MAX", exp.column("end_ts", table=table_alias)).as_("max_end_ts"),
+                )
+                .from_(exp.to_table(self.intervals_table).as_(table_alias))
+                .where(where, copy=False)
+                .where(exp.to_column("is_dev").not_(), copy=False)
+                .group_by(name_col, version_col, copy=False)
+            )
+            query = exp.select(exp.func("MIN", exp.column("max_end_ts"))).from_(
+                max_end_subquery.subquery(alias="max_ends")
+            )
+
+            end = self.engine_adapter.fetchone(query, quote_identifiers=True)[0]
+
+            if greatest_common_end is None:
+                greatest_common_end = end
+            elif end is not None:
+                greatest_common_end = min(greatest_common_end, end)
+
+        return greatest_common_end
+
     def recycle(self) -> None:
         self.engine_adapter.recycle()
 
diff --git a/sqlmesh/schedulers/airflow/api.py b/sqlmesh/schedulers/airflow/api.py
index 84d63607..9f282b08 100644
--- a/sqlmesh/schedulers/airflow/api.py
+++ b/sqlmesh/schedulers/airflow/api.py
@@ -75,14 +75,19 @@ def get_environments() -> Response:
 @check_authentication
 def get_max_interval_end(name: str) -> Response:
     with util.scoped_state_sync() as state_sync:
-        models = None
-        if "models" in request.args:
-            models = json.loads(request.args["models"])
+        max_interval_end = state_sync.max_interval_end_for_environment(name)
+        response = common.IntervalEndResponse(environment=name, max_interval_end=max_interval_end)
+        return _success(response)
 
-        max_interval_end = state_sync.max_interval_end_for_environment(name, models=models)
-        response = common.MaxIntervalEndResponse(
-            environment=name, max_interval_end=max_interval_end
-        )
+
+@sqlmesh_api_v1.route("/environments/<name>/greatest_common_interval_end")
+@csrf.exempt
+@check_authentication
+def get_greatest_common_interval_end(name: str) -> Response:
+    with util.scoped_state_sync() as state_sync:
+        models = json.loads(request.args["models"]) if "models" in request.args else []
+        max_interval_end = state_sync.greatest_common_interval_end(name, set(models))
+        response = common.IntervalEndResponse(environment=name, max_interval_end=max_interval_end)
         return _success(response)
 
 
diff --git a/sqlmesh/schedulers/airflow/client.py b/sqlmesh/schedulers/airflow/client.py
index b8769ed7..2fdef98c 100644
--- a/sqlmesh/schedulers/airflow/client.py
+++ b/sqlmesh/schedulers/airflow/client.py
@@ -271,16 +271,18 @@ class AirflowClient(BaseAirflowClient):
         response = self._get(ENVIRONMENTS_PATH)
         return common.EnvironmentsResponse.parse_obj(response).environments
 
-    def max_interval_end_for_environment(
-        self, environment: str, models: t.Optional[t.Collection[str]] = None
+    def max_interval_end_for_environment(self, environment: str) -> t.Optional[int]:
+        response = self._get(f"{ENVIRONMENTS_PATH}/{environment}/max_interval_end")
+        return common.IntervalEndResponse.parse_obj(response).max_interval_end
+
+    def greatest_common_interval_end(
+        self, environment: str, models: t.Collection[str]
     ) -> t.Optional[int]:
-        url = f"{ENVIRONMENTS_PATH}/{environment}/max_interval_end"
-        if models is not None:
-            models_param = _json_query_param(list(models))
-            response = self._get(url, models=models_param)
-        else:
-            response = self._get(url)
-        return common.MaxIntervalEndResponse.parse_obj(response).max_interval_end
+        response = self._get(
+            f"{ENVIRONMENTS_PATH}/{environment}/greatest_common_interval_end",
+            models=_json_query_param(list(models)),
+        )
+        return common.IntervalEndResponse.parse_obj(response).max_interval_end
 
     def invalidate_environment(self, environment: str) -> None:
         response = self._session.delete(
diff --git a/sqlmesh/schedulers/airflow/common.py b/sqlmesh/schedulers/airflow/common.py
index 181ae06b..b11b575d 100644
--- a/sqlmesh/schedulers/airflow/common.py
+++ b/sqlmesh/schedulers/airflow/common.py
@@ -110,7 +110,7 @@ class InvalidateEnvironmentResponse(PydanticModel):
     name: str
 
 
-class MaxIntervalEndResponse(PydanticModel):
+class IntervalEndResponse(PydanticModel):
     environment: str
     max_interval_end: t.Optional[int] = None
 
diff --git a/sqlmesh/schedulers/airflow/state_sync.py b/sqlmesh/schedulers/airflow/state_sync.py
index 9bcf6c7d..06fe191c 100644
--- a/sqlmesh/schedulers/airflow/state_sync.py
+++ b/sqlmesh/schedulers/airflow/state_sync.py
@@ -67,20 +67,28 @@ class HttpStateSync(StateSync):
         """
         return self._client.get_environments()
 
-    def max_interval_end_for_environment(
-        self, environment: str, models: t.Optional[t.Set[str]] = None
-    ) -> t.Optional[int]:
+    def max_interval_end_for_environment(self, environment: str) -> t.Optional[int]:
         """Returns the max interval end for the given environment.
 
         Args:
             environment: The environment.
-            models: The optional model FQNs to pick intervals from. If not provided all
-                models in the environment will be used.
 
         Returns:
             A timestamp or None if no interval or environment exists.
         """
-        return self._client.max_interval_end_for_environment(environment, models=models)
+        return self._client.max_interval_end_for_environment(environment)
+
+    def greatest_common_interval_end(self, environment: str, models: t.Set[str]) -> t.Optional[int]:
+        """Returns the greatest common interval end for given models in the target environment.
+
+        Args:
+            environment: The environment.
+            models: The model FQNs to select intervals from.
+
+        Returns:
+            A timestamp or None if no interval or environment exists.
+        """
+        return self._client.greatest_common_interval_end(environment, models)
 
     def get_snapshots(
         self,


Test patch:
diff --git a/tests/core/test_state_sync.py b/tests/core/test_state_sync.py
index 17c25146..b7c1ba95 100644
--- a/tests/core/test_state_sync.py
+++ b/tests/core/test_state_sync.py
@@ -1608,7 +1608,7 @@ def test_max_interval_end_for_environment(
             name=environment_name,
             snapshots=[snapshot_a.table_info, snapshot_b.table_info],
             start_at="2023-01-01",
-            end_at="2023-01-02",
+            end_at="2023-01-03",
             plan_id="test_plan_id",
         )
     )
@@ -1617,16 +1617,66 @@ def test_max_interval_end_for_environment(
         "2023-01-03"
     )
 
-    assert state_sync.max_interval_end_for_environment(
-        environment_name, models={snapshot_a.name}
+
+def test_greatest_common_interval_end(
+    state_sync: EngineAdapterStateSync, make_snapshot: t.Callable
+) -> None:
+    snapshot_a = make_snapshot(
+        SqlModel(
+            name="a",
+            cron="@daily",
+            query=parse_one("select 1, ds"),
+        ),
+    )
+    snapshot_a.categorize_as(SnapshotChangeCategory.BREAKING)
+
+    snapshot_b = make_snapshot(
+        SqlModel(
+            name="b",
+            cron="@daily",
+            query=parse_one("select 2, ds"),
+        ),
+    )
+    snapshot_b.categorize_as(SnapshotChangeCategory.BREAKING)
+
+    state_sync.push_snapshots([snapshot_a, snapshot_b])
+
+    state_sync.add_interval(snapshot_a, "2023-01-01", "2023-01-01")
+    state_sync.add_interval(snapshot_a, "2023-01-02", "2023-01-02")
+    state_sync.add_interval(snapshot_a, "2023-01-03", "2023-01-03")
+    state_sync.add_interval(snapshot_b, "2023-01-01", "2023-01-01")
+    state_sync.add_interval(snapshot_b, "2023-01-02", "2023-01-02")
+
+    environment_name = "test_max_interval_end_for_environment"
+
+    assert state_sync.greatest_common_interval_end(environment_name, {snapshot_a.name}) is None
+
+    state_sync.promote(
+        Environment(
+            name=environment_name,
+            snapshots=[snapshot_a.table_info, snapshot_b.table_info],
+            start_at="2023-01-01",
+            end_at="2023-01-03",
+            plan_id="test_plan_id",
+        )
+    )
+
+    assert state_sync.greatest_common_interval_end(
+        environment_name, {snapshot_a.name}
+    ) == to_timestamp("2023-01-04")
+
+    assert state_sync.greatest_common_interval_end(
+        environment_name, {snapshot_b.name}
     ) == to_timestamp("2023-01-03")
 
-    assert state_sync.max_interval_end_for_environment(
-        environment_name, models={snapshot_b.name}
-    ) == to_timestamp("2023-01-02")
+    assert state_sync.greatest_common_interval_end(
+        environment_name, {snapshot_a.name, snapshot_b.name}
+    ) == to_timestamp("2023-01-03")
 
-    assert state_sync.max_interval_end_for_environment(environment_name, models={"missing"}) is None
-    assert state_sync.max_interval_end_for_environment(environment_name, models=set()) is None
+    assert state_sync.greatest_common_interval_end(environment_name, {"missing"}) == to_timestamp(
+        "2023-01-03"
+    )
+    assert state_sync.greatest_common_interval_end(environment_name, set()) is None
 
 
 def test_get_snapshots(mocker):
diff --git a/tests/schedulers/airflow/test_client.py b/tests/schedulers/airflow/test_client.py
index eda1eeff..d1ab954a 100644
--- a/tests/schedulers/airflow/test_client.py
+++ b/tests/schedulers/airflow/test_client.py
@@ -350,7 +350,7 @@ def test_get_environments(mocker: MockerFixture, snapshot: Snapshot):
 
 
 def test_max_interval_end_for_environment(mocker: MockerFixture, snapshot: Snapshot):
-    response = common.MaxIntervalEndResponse(
+    response = common.IntervalEndResponse(
         environment="test_environment", max_interval_end=to_timestamp("2023-01-01")
     )
 
@@ -370,8 +370,8 @@ def test_max_interval_end_for_environment(mocker: MockerFixture, snapshot: Snaps
     )
 
 
-def test_max_interval_end_for_environment_with_models(mocker: MockerFixture, snapshot: Snapshot):
-    response = common.MaxIntervalEndResponse(
+def test_greatest_common_interval_end(mocker: MockerFixture, snapshot: Snapshot):
+    response = common.IntervalEndResponse(
         environment="test_environment", max_interval_end=to_timestamp("2023-01-01")
     )
 
@@ -382,12 +382,12 @@ def test_max_interval_end_for_environment_with_models(mocker: MockerFixture, sna
     max_interval_end_mock.return_value = max_interval_end_response_mock
 
     client = AirflowClient(airflow_url=common.AIRFLOW_LOCAL_URL, session=requests.Session())
-    result = client.max_interval_end_for_environment("test_environment", models={"a.b.c"})
+    result = client.greatest_common_interval_end("test_environment", {"a.b.c"})
 
     assert result == response.max_interval_end
 
     max_interval_end_mock.assert_called_once_with(
-        "http://localhost:8080/sqlmesh/api/v1/environments/test_environment/max_interval_end?models=%5B%22a.b.c%22%5D"
+        "http://localhost:8080/sqlmesh/api/v1/environments/test_environment/greatest_common_interval_end?models=%5B%22a.b.c%22%5D"
     )