Repository: TobikoData/sqlmesh
Pull request #3308 (2024-10-29 19:36:05)

Fix: Only backfill models selected for restatement when restatements are provided
This fix ensures that unrelated models will not be evaluated as part of the restatement plan even if they are behind.


Patch:
diff --git a/docs/reference/cli.md b/docs/reference/cli.md
index 15c6aab0..334c44c8 100644
--- a/docs/reference/cli.md
+++ b/docs/reference/cli.md
@@ -398,16 +398,19 @@ Usage: sqlmesh run [OPTIONS] [ENVIRONMENT]
   Evaluate missing intervals for the target environment.
 
 Options:
-  -s, --start TEXT     The start datetime of the interval for which this
-                       command will be applied.
-  -e, --end TEXT       The end datetime of the interval for which this command
-                       will be applied.
-  --skip-janitor       Skip the janitor task.
-  --ignore-cron        Run for all missing intervals, ignoring individual cron
-                       schedules.
-  --select-model TEXT  Select specific models to run. Note: this always
-                       includes upstream dependencies.
-  --help               Show this message and exit.
+  -s, --start TEXT              The start datetime of the interval for which
+                                this command will be applied.
+  -e, --end TEXT                The end datetime of the interval for which
+                                this command will be applied.
+  --skip-janitor                Skip the janitor task.
+  --ignore-cron                 Run for all missing intervals, ignoring
+                                individual cron schedules.
+  --select-model TEXT           Select specific models to run. Note: this
+                                always includes upstream dependencies.
+  --exit-on-env-update INTEGER  If set, the command will exit with the
+                                specified code if the run is interrupted by an
+                                update to the target environment.
+  --help                        Show this message and exit.
 ```
 
 ## table_diff
diff --git a/sqlmesh/cli/main.py b/sqlmesh/cli/main.py
index 284b504f..5ea0b641 100644
--- a/sqlmesh/cli/main.py
+++ b/sqlmesh/cli/main.py
@@ -444,6 +444,11 @@ def plan(
     multiple=True,
     help="Select specific models to run. Note: this always includes upstream dependencies.",
 )
+@click.option(
+    "--exit-on-env-update",
+    type=int,
+    help="If set, the command will exit with the specified code if the run is interrupted by an update to the target environment.",
+)
 @click.pass_context
 @error_handler
 @cli_analytics
diff --git a/sqlmesh/core/analytics/collector.py b/sqlmesh/core/analytics/collector.py
index 8661eeef..e2779d02 100644
--- a/sqlmesh/core/analytics/collector.py
+++ b/sqlmesh/core/analytics/collector.py
@@ -252,11 +252,15 @@ class AnalyticsCollector:
         )
         return run_id
 
-    def on_run_end(self, *, run_id: str, succeeded: bool, error: t.Optional[t.Any] = None) -> None:
+    def on_run_end(
+        self, *, run_id: str, succeeded: bool, interrupted: bool, error: t.Optional[t.Any] = None
+    ) -> None:
         """Called after a run ends.
 
         Args:
             run_id: The ID of the run.
+            succeeded: Whether the run succeeded.
+            interrupted: Whether the run was interrupted.
             error: The error that occurred during the run, if any.
         """
         self._add_event(
@@ -264,6 +268,7 @@ class AnalyticsCollector:
             {
                 "run_id": run_id,
                 "succeeded": succeeded,
+                "interrupted": interrupted,
                 "error": type(error).__name__ if error else None,
             },
         )
diff --git a/sqlmesh/core/context.py b/sqlmesh/core/context.py
index 9acd20fd..9c32b7a0 100644
--- a/sqlmesh/core/context.py
+++ b/sqlmesh/core/context.py
@@ -36,6 +36,7 @@ from __future__ import annotations
 import abc
 import collections
 import logging
+import sys
 import time
 import traceback
 import typing as t
@@ -592,6 +593,7 @@ class GenericContext(BaseContext, t.Generic[C]):
         skip_janitor: bool = False,
         ignore_cron: bool = False,
         select_models: t.Optional[t.Collection[str]] = None,
+        exit_on_env_update: t.Optional[int] = None,
     ) -> bool:
         """Run the entire dag through the scheduler.
 
@@ -604,6 +606,8 @@ class GenericContext(BaseContext, t.Generic[C]):
             ignore_cron: Whether to ignore the model's cron schedule and run all available missing intervals.
             select_models: A list of model selection expressions to filter models that should run. Note that
                 upstream dependencies of selected models will also be evaluated.
+            exit_on_env_update: If set, exits with the provided code if the run is interrupted by an update
+                to the target environment.
 
         Returns:
             True if the run was successful, False otherwise.
@@ -620,6 +624,7 @@ class GenericContext(BaseContext, t.Generic[C]):
         self._load_materializations_and_signals()
 
         success = False
+        interrupted = False
         try:
             success = self._run(
                 environment=environment,
@@ -629,16 +634,21 @@ class GenericContext(BaseContext, t.Generic[C]):
                 skip_janitor=skip_janitor,
                 ignore_cron=ignore_cron,
                 select_models=select_models,
+                exit_on_env_update=exit_on_env_update is not None,
             )
+        except CircuitBreakerError:
+            interrupted = True
         except Exception as e:
             self.notification_target_manager.notify(
                 NotificationEvent.RUN_FAILURE, traceback.format_exc()
             )
             logger.error(f"Run Failure: {traceback.format_exc()}")
-            analytics.collector.on_run_end(run_id=analytics_run_id, succeeded=False, error=e)
+            analytics.collector.on_run_end(
+                run_id=analytics_run_id, succeeded=False, interrupted=False, error=e
+            )
             raise e
 
-        if success:
+        if success or interrupted:
             self.notification_target_manager.notify(
                 NotificationEvent.RUN_END, environment=environment
             )
@@ -648,7 +658,12 @@ class GenericContext(BaseContext, t.Generic[C]):
                 NotificationEvent.RUN_FAILURE, "See console logs for details."
             )
 
-        analytics.collector.on_run_end(run_id=analytics_run_id, succeeded=success)
+        analytics.collector.on_run_end(
+            run_id=analytics_run_id, succeeded=success, interrupted=interrupted
+        )
+
+        if interrupted and exit_on_env_update is not None:
+            sys.exit(exit_on_env_update)
 
         return success
 
@@ -1800,6 +1815,7 @@ class GenericContext(BaseContext, t.Generic[C]):
         skip_janitor: bool,
         ignore_cron: bool,
         select_models: t.Optional[t.Collection[str]],
+        exit_on_env_update: bool,
     ) -> bool:
         if not skip_janitor and environment.lower() == c.PROD:
             self._run_janitor()
@@ -1873,6 +1889,8 @@ class GenericContext(BaseContext, t.Generic[C]):
                     "Environment '%s' has been modified while running. Restarting the run...",
                     environment,
                 )
+                if exit_on_env_update:
+                    raise
 
         return success
 
diff --git a/sqlmesh/core/plan/builder.py b/sqlmesh/core/plan/builder.py
index af0a3489..13eae9e5 100644
--- a/sqlmesh/core/plan/builder.py
+++ b/sqlmesh/core/plan/builder.py
@@ -233,10 +233,10 @@ class PlanBuilder:
             else DeployabilityIndex.all_deployable()
         )
 
-        models_to_backfill = self._build_models_to_backfill(dag)
         restatements = self._build_restatements(
             dag, earliest_interval_start(self._context_diff.snapshots.values())
         )
+        models_to_backfill = self._build_models_to_backfill(dag, restatements)
 
         interval_end_per_model = self._interval_end_per_model
         if interval_end_per_model and self.override_end:
@@ -393,15 +393,25 @@ class PlanBuilder:
             indirectly_modified,
         )
 
-    def _build_models_to_backfill(self, dag: DAG[SnapshotId]) -> t.Optional[t.Set[str]]:
-        if self._backfill_models is None:
+    def _build_models_to_backfill(
+        self, dag: DAG[SnapshotId], restatements: t.Collection[SnapshotId]
+    ) -> t.Optional[t.Set[str]]:
+        backfill_models = (
+            self._backfill_models
+            if self._backfill_models is not None
+            else [r.name for r in restatements]
+            # Only backfill models explicitly marked for restatement.
+            if self._restate_models
+            else None
+        )
+        if backfill_models is None:
             return None
         return {
             self._context_diff.snapshots[s_id].name
             for s_id in dag.subdag(
                 *[
                     self._model_fqn_to_snapshot[m].snapshot_id
-                    for m in self._backfill_models
+                    for m in backfill_models
                     if m in self._model_fqn_to_snapshot
                 ]
             ).sorted
diff --git a/sqlmesh/magics.py b/sqlmesh/magics.py
index df8b5f26..e91d72e6 100644
--- a/sqlmesh/magics.py
+++ b/sqlmesh/magics.py
@@ -451,6 +451,11 @@ class SQLMeshMagics(Magics):
         nargs="*",
         help="Select specific models to run. Note: this always includes upstream dependencies.",
     )
+    @argument(
+        "--exit-on-env-update",
+        type=int,
+        help="If set, the command will exit with the specified code if the run is interrupted by an update to the target environment.",
+    )
     @line_magic
     @pass_sqlmesh_context
     def run_dag(self, context: Context, line: str) -> None:
@@ -464,6 +469,7 @@ class SQLMeshMagics(Magics):
             skip_janitor=args.skip_janitor,
             ignore_cron=args.ignore_cron,
             select_models=args.select_model,
+            exit_on_env_update=args.exit_on_env_update,
         )
         if not success:
             raise SQLMeshError("Error Running DAG. Check logs for details.")


Test patch:
diff --git a/tests/core/analytics/test_collector.py b/tests/core/analytics/test_collector.py
index 3de73983..b3baa7ed 100644
--- a/tests/core/analytics/test_collector.py
+++ b/tests/core/analytics/test_collector.py
@@ -284,8 +284,11 @@ def test_on_snapshots_created(
 
 def test_on_run(collector: AnalyticsCollector, mocker: MockerFixture):
     run_id = collector.on_run_start(engine_type="bigquery", state_sync_type="mysql")
-    collector.on_run_end(run_id=run_id, succeeded=True)
-    collector.on_run_end(run_id=run_id, succeeded=False, error=SQLMeshError("test_error"))
+    collector.on_run_end(run_id=run_id, succeeded=True, interrupted=False)
+    collector.on_run_end(
+        run_id=run_id, succeeded=False, interrupted=False, error=SQLMeshError("test_error")
+    )
+    collector.on_run_end(run_id=run_id, succeeded=False, interrupted=True)
 
     collector.flush()
 
@@ -309,7 +312,7 @@ def test_on_run(collector: AnalyticsCollector, mocker: MockerFixture):
                 {
                     "seq_num": 1,
                     "event_type": "RUN_END",
-                    "event": f'{{"run_id": "{run_id}", "succeeded": true, "error": null}}',
+                    "event": f'{{"run_id": "{run_id}", "succeeded": true, "interrupted": false, "error": null}}',
                     **common_fields,
                 }
             ),
@@ -317,7 +320,15 @@ def test_on_run(collector: AnalyticsCollector, mocker: MockerFixture):
                 {
                     "seq_num": 2,
                     "event_type": "RUN_END",
-                    "event": f'{{"run_id": "{run_id}", "succeeded": false, "error": "SQLMeshError"}}',
+                    "event": f'{{"run_id": "{run_id}", "succeeded": false, "interrupted": false, "error": "SQLMeshError"}}',
+                    **common_fields,
+                }
+            ),
+            call(
+                {
+                    "seq_num": 3,
+                    "event_type": "RUN_END",
+                    "event": f'{{"run_id": "{run_id}", "succeeded": false, "interrupted": true, "error": null}}',
                     **common_fields,
                 }
             ),
diff --git a/tests/core/test_plan.py b/tests/core/test_plan.py
index 09869e16..8013c9d0 100644
--- a/tests/core/test_plan.py
+++ b/tests/core/test_plan.py
@@ -745,14 +745,25 @@ def test_restate_models(sushi_context_pre_scheduling: Context):
         ),
     }
     assert plan.requires_backfill
+    assert plan.models_to_backfill == {
+        '"memory"."sushi"."customer_revenue_by_day"',
+        '"memory"."sushi"."customer_revenue_lifetime"',
+        '"memory"."sushi"."items"',
+        '"memory"."sushi"."order_items"',
+        '"memory"."sushi"."orders"',
+        '"memory"."sushi"."top_waiters"',
+        '"memory"."sushi"."waiter_revenue_by_day"',
+    }
 
     plan = sushi_context_pre_scheduling.plan(restate_models=["unknown_model"], no_prompts=True)
     assert not plan.has_changes
     assert not plan.restatements
+    assert plan.models_to_backfill is None
 
     plan = sushi_context_pre_scheduling.plan(restate_models=["tag:unknown_tag"], no_prompts=True)
     assert not plan.has_changes
     assert not plan.restatements
+    assert plan.models_to_backfill is None
 
 
 @pytest.mark.slow
@@ -808,6 +819,13 @@ def test_restate_models_with_existing_missing_intervals(sushi_context: Context):
         ),
     ]
     assert plan.requires_backfill
+    assert plan.models_to_backfill == {
+        top_waiters_snapshot_id.name,
+        waiter_revenue_by_day_snapshot_id.name,
+        '"memory"."sushi"."items"',
+        '"memory"."sushi"."order_items"',
+        '"memory"."sushi"."orders"',
+    }
 
 
 def test_restate_symbolic_model(make_snapshot, mocker: MockerFixture):