Pass audit results directly to evaluation prog bar output

Author: treysp

sqlmesh/core/console.py

@@ -29,7 +29,6 @@
 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,
@@ -93,12 +92,6 @@ 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,
@@ -114,7 +107,12 @@ def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
 
     @abc.abstractmethod
     def update_snapshot_evaluation_progress(
-        self, snapshot: Snapshot, batch_idx: int, duration_ms: t.Optional[int]
+        self,
+        snapshot: Snapshot,
+        batch_idx: int,
+        duration_ms: t.Optional[int],
+        num_audits_passed: int,
+        num_audits_failed: int,
     ) -> None:
         """Updates the snapshot evaluation progress."""
 
@@ -350,11 +348,6 @@ 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,
         batches: t.Dict[Snapshot, Intervals],
@@ -367,7 +360,12 @@ def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
         pass
 
     def update_snapshot_evaluation_progress(
-        self, snapshot: Snapshot, batch_idx: int, duration_ms: t.Optional[int]
+        self,
+        snapshot: Snapshot,
+        batch_idx: int,
+        duration_ms: t.Optional[int],
+        num_audits_passed: int,
+        num_audits_failed: int,
     ) -> None:
         pass
 
@@ -551,7 +549,6 @@ def __init__(
         self.evaluation_model_batch_sizes: t.Dict[Snapshot, int] = {}
         self.evaluation_model_info: t.Dict[Snapshot, t.Dict[str, t.Any]] = {}
         self.evaluation_model_column_widths: t.Dict[str, int] = {}
-        self.evaluation_audit_results: t.Dict[Snapshot, t.List[AuditResult]] = {}
 
         # Put in temporary values that are replaced when evaluating
         self.environment_naming_info = EnvironmentNamingInfo()
@@ -590,11 +587,6 @@ 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,
         batched_intervals: t.Dict[Snapshot, Intervals],
@@ -654,7 +646,12 @@ def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
             )
 
     def update_snapshot_evaluation_progress(
-        self, snapshot: Snapshot, batch_idx: int, duration_ms: t.Optional[int]
+        self,
+        snapshot: Snapshot,
+        batch_idx: int,
+        duration_ms: t.Optional[int],
+        num_audits_passed: int,
+        num_audits_failed: int,
     ) -> None:
         """Update the snapshot evaluation progress."""
         if (
@@ -672,14 +669,6 @@ def update_snapshot_evaluation_progress(
                 )
 
                 annotation = self.evaluation_model_info[snapshot]["annotation"][batch_idx]
-                num_audits = 0
-                num_audits_passed = 0
-                if snapshot in self.evaluation_audit_results:
-                    num_audits = len(self.evaluation_audit_results[snapshot])
-                    num_audits_passed = sum(
-                        result.count == 0 for result in self.evaluation_audit_results[snapshot]
-                    )
-                num_audits_failed = num_audits - num_audits_passed
                 if num_audits_passed:
                     annotation += f", {num_audits_passed} audits pass"
                 if num_audits_failed:
@@ -721,7 +710,6 @@ def stop_evaluation_progress(self, success: bool = True) -> None:
         self.evaluation_model_batch_sizes = {}
         self.evaluation_model_info = {}
         self.evaluation_model_column_widths = {}
-        self.evaluation_audit_results = {}
         self.environment_naming_info = EnvironmentNamingInfo()
         self.default_catalog = None
 
@@ -2348,7 +2336,12 @@ def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
             )
 
     def update_snapshot_evaluation_progress(
-        self, snapshot: Snapshot, batch_idx: int, duration_ms: t.Optional[int]
+        self,
+        snapshot: Snapshot,
+        batch_idx: int,
+        duration_ms: t.Optional[int],
+        num_audits_passed: int,
+        num_audits_failed: int,
     ) -> None:
         view_name, loaded_batches = self.evaluation_batch_progress[snapshot.snapshot_id]
         total_batches = self.evaluation_model_batch_sizes[snapshot]
@@ -2499,9 +2492,16 @@ def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
         self._write(f"Evaluating {snapshot.name}")
 
     def update_snapshot_evaluation_progress(
-        self, snapshot: Snapshot, batch_idx: int, duration_ms: t.Optional[int]
+        self,
+        snapshot: Snapshot,
+        batch_idx: int,
+        duration_ms: t.Optional[int],
+        num_audits_passed: int,
+        num_audits_failed: int,
     ) -> None:
-        self._write(f"Evaluating {snapshot.name} | batch={batch_idx} | duration={duration_ms}ms")
+        self._write(
+            f"Evaluating {snapshot.name} | batch={batch_idx} | duration={duration_ms}ms | num_audits_passed={num_audits_passed} | num_audits_failed={num_audits_failed}"
+        )
 
     def stop_evaluation_progress(self, success: bool = True) -> None:
         self._write(f"Stopping evaluation with success={success}")

sqlmesh/core/scheduler.py

@@ -7,6 +7,7 @@
 from sqlmesh.core.console import Console, get_console
 from sqlmesh.core.environment import EnvironmentNamingInfo, execute_environment_statements
 from sqlmesh.core.macros import RuntimeStage
+from sqlmesh.core.model.definition import AuditResult
 from sqlmesh.core.node import IntervalUnit
 from sqlmesh.core.notification_target import (
     NotificationEvent,
@@ -168,7 +169,7 @@ def evaluate(
         environment_naming_info: EnvironmentNamingInfo,
         default_catalog: t.Optional[str],
         **kwargs: t.Any,
-    ) -> None:
+    ) -> t.List[AuditResult]:
         """Evaluate a snapshot and add the processed interval to the state sync.
 
         Args:
@@ -180,6 +181,9 @@ def evaluate(
             batch_index: If the snapshot is part of a batch of related snapshots; which index in the batch is it
             auto_restatement_enabled: Whether to enable auto restatements.
             kwargs: Additional kwargs to pass to the renderer.
+
+        Returns:
+            List of audit results from the evaluation.
         """
         validate_date_range(start, end)
 
@@ -207,7 +211,6 @@ def evaluate(
             wap_id=wap_id,
             **kwargs,
         )
-        self.console.store_evaluation_audit_results(snapshot, audit_results)
 
         audit_errors_to_raise: t.List[AuditError] = []
         for audit_result in (result for result in audit_results if result.count):
@@ -241,6 +244,7 @@ def evaluate(
             raise NodeAuditsErrors(audit_errors_to_raise)
 
         self.state_sync.add_interval(snapshot, start, end, is_dev=not is_deployable)
+        return audit_results
 
     def run(
         self,
@@ -474,7 +478,8 @@ def evaluate_node(node: SchedulingUnit) -> None:
             try:
                 assert execution_time  # mypy
                 assert deployability_index  # mypy
-                self.evaluate(
+                audit_results = []
+                audit_results = self.evaluate(
                     snapshot=snapshot,
                     start=start,
                     end=end,
@@ -486,8 +491,14 @@ def evaluate_node(node: SchedulingUnit) -> None:
                 )
                 evaluation_duration_ms = now_timestamp() - execution_start_ts
             finally:
+                num_audits = len(audit_results)
+                num_audits_failed = sum(1 for result in audit_results if result.count)
                 self.console.update_snapshot_evaluation_progress(
-                    snapshot, batch_idx, evaluation_duration_ms
+                    snapshot,
+                    batch_idx,
+                    evaluation_duration_ms,
+                    num_audits - num_audits_failed,
+                    num_audits_failed,
                 )
 
         try:

web/server/console.py

@@ -124,7 +124,12 @@ def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
         self.log_event_plan_apply()
 
     def update_snapshot_evaluation_progress(
-        self, snapshot: Snapshot, batch_idx: int, duration_ms: t.Optional[int]
+        self,
+        snapshot: Snapshot,
+        batch_idx: int,
+        duration_ms: t.Optional[int],
+        audits_passed: int,
+        audits_failed: int,
     ) -> None:
         if self.plan_apply_stage_tracker and self.plan_apply_stage_tracker.backfill:
             task = self.plan_apply_stage_tracker.backfill.tasks[snapshot.name]