Overhaul CLI evaluation progress bar

Author: treysp

sqlmesh/core/console.py

@@ -30,6 +30,7 @@
 from sqlmesh.core.environment import EnvironmentNamingInfo
 from sqlmesh.core.linter.rule import RuleViolation
 from sqlmesh.core.model import Model
+from sqlmesh.core.model.definition import AuditResult
 from sqlmesh.core.snapshot import (
     Snapshot,
     SnapshotChangeCategory,
@@ -96,6 +97,12 @@ def start_plan_evaluation(self, plan: EvaluatablePlan) -> None:
     def stop_plan_evaluation(self) -> None:
         """Indicates that the evaluation has ended."""
 
+    @abc.abstractmethod
+    def store_evaluation_audit_results(
+        self, snapshot: Snapshot, audit_results: t.List[AuditResult]
+    ) -> None:
+        """Stores the audit results for the snapshot evaluation."""
+
     @abc.abstractmethod
     def start_evaluation_progress(
         self,
@@ -413,6 +420,11 @@ def start_plan_evaluation(self, plan: EvaluatablePlan) -> None:
     def stop_plan_evaluation(self) -> None:
         pass
 
+    def store_evaluation_audit_results(
+        self, snapshot: Snapshot, audit_results: t.List[AuditResult]
+    ) -> None:
+        pass
+
     def start_evaluation_progress(
         self,
         batch_sizes: t.Dict[Snapshot, int],
@@ -678,6 +690,7 @@ def __init__(
         self.evaluation_total_task: t.Optional[TaskID] = None
         self.evaluation_model_progress: t.Optional[Progress] = None
         self.evaluation_model_tasks: t.Dict[str, TaskID] = {}
+        self.evaluation_model_batch_sizes: t.Dict[Snapshot, int] = {}
 
         # Put in temporary values that are replaced when evaluating
         self.environment_naming_info = EnvironmentNamingInfo()
@@ -726,6 +739,11 @@ def start_plan_evaluation(self, plan: EvaluatablePlan) -> None:
     def stop_plan_evaluation(self) -> None:
         pass
 
+    def store_evaluation_audit_results(
+        self, snapshot: Snapshot, audit_results: t.List[AuditResult]
+    ) -> None:
+        self.evaluation_audit_results[snapshot] = audit_results
+
     def start_evaluation_progress(
         self,
         batch_sizes: t.Dict[Snapshot, int],

sqlmesh/core/scheduler.py

@@ -166,6 +166,8 @@ def evaluate(
         execution_time: TimeLike,
         deployability_index: DeployabilityIndex,
         batch_index: int,
+        environment_naming_info: EnvironmentNamingInfo,
+        default_catalog: t.Optional[str],
         **kwargs: t.Any,
     ) -> t.Tuple[t.List[AuditResult], t.List[AuditError]]:
         """Evaluate a snapshot and add the processed interval to the state sync.
@@ -209,6 +211,7 @@ def evaluate(
             wap_id=wap_id,
             **kwargs,
         )
+        get_console().store_evaluation_audit_results(snapshot, audit_results)
 
         audit_errors_to_raise: t.List[AuditError] = []
         audit_errors_to_warn: t.List[AuditError] = []
@@ -427,7 +430,7 @@ def run_merged_intervals(
         batched_intervals = self.batch_intervals(merged_intervals)
 
         self.console.start_evaluation_progress(
-            {snapshot: len(intervals) for snapshot, intervals in batched_intervals.items()},
+            batched_intervals,
             environment_naming_info,
             self.default_catalog,
         )

tests/cli/test_cli.py

@@ -719,7 +719,7 @@ def test_run_dev(runner, tmp_path, flag):
     # Confirm backfill occurs when we run non-backfilled dev env
     result = runner.invoke(cli, ["--log-file-dir", tmp_path, "--paths", tmp_path, "run", "dev"])
     assert result.exit_code == 0
-    assert_model_batches_executed(result)
+    assert_model_batches_evaluated(result)
 
 
 @time_machine.travel(FREEZE_TIME)
@@ -751,7 +751,7 @@ def test_run_cron_elapsed(runner, tmp_path):
         result = runner.invoke(cli, ["--log-file-dir", tmp_path, "--paths", tmp_path, "run"])
 
     assert result.exit_code == 0
-    assert_model_batches_executed(result)
+    assert_model_batches_evaluated(result)
 
 
 def test_clean(runner, tmp_path):

tests/core/test_scheduler.py

@@ -528,6 +528,8 @@ def _evaluate():
             to_datetime("2022-01-03"),
             DeployabilityIndex.all_deployable(),
             0,
+            EnvironmentNamingInfo(),
+            None,
         )
 
     evaluator_audit_mock.return_value = [

web/server/console.py

@@ -12,6 +12,7 @@
 from sqlmesh.core.environment import EnvironmentNamingInfo
 from sqlmesh.core.plan.definition import EvaluatablePlan
 from sqlmesh.core.snapshot import Snapshot, SnapshotInfoLike
+from sqlmesh.core.snapshot.definition import Intervals
 from sqlmesh.core.test import ModelTest
 from sqlmesh.utils.date import now_timestamp
 from web.server import models