Chore: Refactor context run and _run (#3328)

Author: crericha

sqlmesh/core/context.py

@@ -617,6 +617,8 @@ def run(
             True if the run was successful, False otherwise.
         """
         environment = environment or self.config.default_target_environment
+        if not skip_janitor and environment.lower() == c.PROD:
+            self._run_janitor()
 
         self.notification_target_manager.notify(
             NotificationEvent.RUN_START, environment=environment
@@ -627,30 +629,76 @@ def run(
         )
         self._load_materializations_and_signals()
 
+        env_check_attempts_num = max(
+            1,
+            self.config.run.environment_check_max_wait
+            // self.config.run.environment_check_interval,
+        )
+
+        def _block_until_finalized() -> str:
+            for _ in range(env_check_attempts_num):
+                assert environment is not None  # mypy
+                environment_state = self.state_sync.get_environment(environment)
+                if not environment_state:
+                    raise SQLMeshError(f"Environment '{environment}' was not found.")
+                if environment_state.finalized_ts:
+                    return environment_state.plan_id
+                logger.warning(
+                    "Environment '%s' is being updated by plan '%s'. Retrying in %s seconds...",
+                    environment,
+                    environment_state.plan_id,
+                    self.config.run.environment_check_interval,
+                )
+                time.sleep(self.config.run.environment_check_interval)
+            raise SQLMeshError(
+                f"Exceeded the maximum wait time for environment '{environment}' to be ready. "
+                "This means that the environment either failed to update or the update is taking longer than expected. "
+                "See https://sqlmesh.readthedocs.io/en/stable/reference/configuration/#run to adjust the timeout settings."
+            )
+
         success = False
         interrupted = False
-        try:
-            success = self._run(
-                environment=environment,
-                start=start,
-                end=end,
-                execution_time=execution_time,
-                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, interrupted=False, error=e
-            )
-            raise e
+        done = False
+        while not done:
+            plan_id_at_start = _block_until_finalized()
+
+            def _has_environment_changed() -> bool:
+                assert environment is not None  # mypy
+                current_environment_state = self.state_sync.get_environment(environment)
+                return (
+                    not current_environment_state
+                    or current_environment_state.plan_id != plan_id_at_start
+                    or not current_environment_state.finalized_ts
+                )
+
+            try:
+                success = self._run(
+                    environment,
+                    start=start,
+                    end=end,
+                    execution_time=execution_time,
+                    ignore_cron=ignore_cron,
+                    select_models=select_models,
+                    circuit_breaker=_has_environment_changed,
+                )
+                done = True
+            except CircuitBreakerError:
+                logger.warning(
+                    "Environment '%s' has been modified while running. Restarting the run...",
+                    environment,
+                )
+                if exit_on_env_update:
+                    interrupted = True
+                    done = 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, interrupted=False, error=e
+                )
+                raise e
 
         if success or interrupted:
             self.notification_target_manager.notify(
@@ -1822,41 +1870,10 @@ def _run(
         start: t.Optional[TimeLike],
         end: t.Optional[TimeLike],
         execution_time: t.Optional[TimeLike],
-        skip_janitor: bool,
         ignore_cron: bool,
         select_models: t.Optional[t.Collection[str]],
-        exit_on_env_update: bool,
+        circuit_breaker: t.Optional[t.Callable[[], bool]],
     ) -> bool:
-        if not skip_janitor and environment.lower() == c.PROD:
-            self._run_janitor()
-
-        env_check_attempts_num = max(
-            1,
-            self.config.run.environment_check_max_wait
-            // self.config.run.environment_check_interval,
-        )
-
-        def _block_until_finalized() -> str:
-            for _ in range(env_check_attempts_num):
-                assert environment is not None  # mypy
-                environment_state = self.state_sync.get_environment(environment)
-                if not environment_state:
-                    raise SQLMeshError(f"Environment '{environment}' was not found.")
-                if environment_state.finalized_ts:
-                    return environment_state.plan_id
-                logger.warning(
-                    "Environment '%s' is being updated by plan '%s'. Retrying in %s seconds...",
-                    environment,
-                    environment_state.plan_id,
-                    self.config.run.environment_check_interval,
-                )
-                time.sleep(self.config.run.environment_check_interval)
-            raise SQLMeshError(
-                f"Exceeded the maximum wait time for environment '{environment}' to be ready. "
-                "This means that the environment either failed to update or the update is taking longer than expected. "
-                "See https://sqlmesh.readthedocs.io/en/stable/reference/configuration/#run to adjust the timeout settings."
-            )
-
         scheduler = self.scheduler(environment=environment)
         snapshots = scheduler.snapshots
 
@@ -1870,39 +1887,15 @@ def _block_until_finalized() -> str:
             model_selector = self._new_selector(models=models, dag=dag)
             select_models = set(dag.subdag(*model_selector.expand_model_selections(select_models)))
 
-        done = False
-        while not done:
-            plan_id_at_start = _block_until_finalized()
-
-            def _has_environment_changed() -> bool:
-                assert environment is not None  # mypy
-                current_environment_state = self.state_sync.get_environment(environment)
-                return (
-                    not current_environment_state
-                    or current_environment_state.plan_id != plan_id_at_start
-                    or not current_environment_state.finalized_ts
-                )
-
-            try:
-                success = scheduler.run(
-                    environment,
-                    start=start,
-                    end=end,
-                    execution_time=execution_time,
-                    ignore_cron=ignore_cron,
-                    circuit_breaker=_has_environment_changed,
-                    selected_snapshots=select_models,
-                )
-                done = True
-            except CircuitBreakerError:
-                logger.warning(
-                    "Environment '%s' has been modified while running. Restarting the run...",
-                    environment,
-                )
-                if exit_on_env_update:
-                    raise
-
-        return success
+        return scheduler.run(
+            environment,
+            start=start,
+            end=end,
+            execution_time=execution_time,
+            ignore_cron=ignore_cron,
+            circuit_breaker=circuit_breaker,
+            selected_snapshots=select_models,
+        )
 
     def _apply(self, plan: Plan, circuit_breaker: t.Optional[t.Callable[[], bool]]) -> None:
         self._scheduler.create_plan_evaluator(self).evaluate(