Repository: TobikoData/sqlmesh
Pull request #3307 (2024-10-29 17:43:54)

Feat: Support exit from the run command with a provided code when interrupted by an update to the target env
Restarting the run may sometimes lead to the run job taking longer than expected. This gives users ability to terminate the run job instead of restarting it in place if an update to the target environment took place.


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/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,
                 }
             ),