PR Feedback 1

Author: VaggelisD

sqlmesh/core/console.py

@@ -350,11 +350,14 @@ def start_evaluation_progress(
         batched_intervals: t.Dict[Snapshot, Intervals],
         environment_naming_info: EnvironmentNamingInfo,
         default_catalog: t.Optional[str],
+        audit_only: bool = False,
     ) -> None:
-        """Indicates that a new snapshot evaluation progress has begun."""
+        """Indicates that a new snapshot evaluation/auditing progress has begun."""
 
     @abc.abstractmethod
-    def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
+    def start_snapshot_evaluation_progress(
+        self, snapshot: Snapshot, audit_only: bool = False
+    ) -> None:
         """Starts the snapshot evaluation progress."""
 
     @abc.abstractmethod
@@ -366,6 +369,7 @@ def update_snapshot_evaluation_progress(
         duration_ms: t.Optional[int],
         num_audits_passed: int,
         num_audits_failed: int,
+        audit_only: bool = False,
     ) -> None:
         """Updates the snapshot evaluation progress."""
 
@@ -507,10 +511,13 @@ def start_evaluation_progress(
         batched_intervals: t.Dict[Snapshot, Intervals],
         environment_naming_info: EnvironmentNamingInfo,
         default_catalog: t.Optional[str],
+        audit_only: bool = False,
     ) -> None:
         pass
 
-    def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
+    def start_snapshot_evaluation_progress(
+        self, snapshot: Snapshot, audit_only: bool = False
+    ) -> None:
         pass
 
     def update_snapshot_evaluation_progress(
@@ -521,6 +528,7 @@ def update_snapshot_evaluation_progress(
         duration_ms: t.Optional[int],
         num_audits_passed: int,
         num_audits_failed: int,
+        audit_only: bool = False,
     ) -> None:
         pass
 
@@ -891,11 +899,12 @@ def start_evaluation_progress(
         batched_intervals: t.Dict[Snapshot, Intervals],
         environment_naming_info: EnvironmentNamingInfo,
         default_catalog: t.Optional[str],
+        audit_only: bool = False,
     ) -> None:
-        """Indicates that a new snapshot evaluation progress has begun."""
+        """Indicates that a new snapshot evaluation/auditing progress has begun."""
         if not self.evaluation_progress_live:
             self.evaluation_total_progress = make_progress_bar(
-                "Executing model batches", self.console
+                "Executing model batches" if not audit_only else "Auditing models", self.console
             )
 
             self.evaluation_model_progress = Progress(
@@ -916,8 +925,9 @@ def start_evaluation_progress(
             batch_sizes = {
                 snapshot: len(intervals) for snapshot, intervals in batched_intervals.items()
             }
+            message = "Executing" if not audit_only else "Auditing"
             self.evaluation_total_task = self.evaluation_total_progress.add_task(
-                "Executing models...", total=sum(batch_sizes.values())
+                f"{message} models...", total=sum(batch_sizes.values())
             )
 
             # determine column widths
@@ -943,15 +953,17 @@ def start_evaluation_progress(
             self.environment_naming_info = environment_naming_info
             self.default_catalog = default_catalog
 
-    def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
+    def start_snapshot_evaluation_progress(
+        self, snapshot: Snapshot, audit_only: bool = False
+    ) -> None:
         if self.evaluation_model_progress and snapshot.name not in self.evaluation_model_tasks:
             display_name = snapshot.display_name(
                 self.environment_naming_info,
                 self.default_catalog if self.verbosity < Verbosity.VERY_VERBOSE else None,
                 dialect=self.dialect,
             )
             self.evaluation_model_tasks[snapshot.name] = self.evaluation_model_progress.add_task(
-                f"Evaluating {display_name}...",
+                f"{'Evaluating' if not audit_only else 'Auditing'} {display_name}...",
                 view_name=display_name,
                 total=self.evaluation_model_batch_sizes[snapshot],
             )
@@ -964,6 +976,7 @@ def update_snapshot_evaluation_progress(
         duration_ms: t.Optional[int],
         num_audits_passed: int,
         num_audits_failed: int,
+        audit_only: bool = False,
     ) -> None:
         """Update the snapshot evaluation progress."""
         if (
@@ -1003,7 +1016,7 @@ def update_snapshot_evaluation_progress(
                     self.evaluation_column_widths["duration"]
                 )
 
-                msg = f"{batch} {display_name}   {annotation}   {duration}".replace(
+                msg = f"{f'{batch} ' if not audit_only else ''}{display_name}   {annotation}   {duration}".replace(
                     self.AUDIT_PASS_MARK, self.GREEN_AUDIT_PASS_MARK
                 )
 
@@ -1015,7 +1028,10 @@ def update_snapshot_evaluation_progress(
 
             model_task_id = self.evaluation_model_tasks[snapshot.name]
             self.evaluation_model_progress.update(model_task_id, refresh=True, advance=1)
-            if self.evaluation_model_progress._tasks[model_task_id].completed >= total_batches:
+            if (
+                self.evaluation_model_progress._tasks[model_task_id].completed >= total_batches
+                or audit_only
+            ):
                 self.evaluation_model_progress.remove_task(model_task_id)
 
     def stop_evaluation_progress(self, success: bool = True) -> None:
@@ -3208,14 +3224,17 @@ def start_evaluation_progress(
         batched_intervals: t.Dict[Snapshot, Intervals],
         environment_naming_info: EnvironmentNamingInfo,
         default_catalog: t.Optional[str],
+        audit_only: bool = False,
     ) -> None:
         self.evaluation_model_batch_sizes = {
             snapshot: len(intervals) for snapshot, intervals in batched_intervals.items()
         }
         self.evaluation_environment_naming_info = environment_naming_info
         self.default_catalog = default_catalog
 
-    def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
+    def start_snapshot_evaluation_progress(
+        self, snapshot: Snapshot, audit_only: bool = False
+    ) -> None:
         if not self.evaluation_batch_progress.get(snapshot.snapshot_id):
             display_name = snapshot.display_name(
                 self.evaluation_environment_naming_info,
@@ -3235,7 +3254,11 @@ def update_snapshot_evaluation_progress(
         duration_ms: t.Optional[int],
         num_audits_passed: int,
         num_audits_failed: int,
+        audit_only: bool = False,
     ) -> None:
+        if audit_only:
+            return
+
         view_name, loaded_batches = self.evaluation_batch_progress[snapshot.snapshot_id]
         total_batches = self.evaluation_model_batch_sizes[snapshot]
 
@@ -3378,13 +3401,17 @@ def start_evaluation_progress(
         batched_intervals: t.Dict[Snapshot, Intervals],
         environment_naming_info: EnvironmentNamingInfo,
         default_catalog: t.Optional[str],
+        audit_only: bool = False,
     ) -> None:
+        message = "evaluation" if not audit_only else "auditing"
         self._write(
-            f"Starting evaluation for {sum(len(intervals) for intervals in batched_intervals.values())} snapshots"
+            f"Starting {message} for {sum(len(intervals) for intervals in batched_intervals.values())} snapshots"
         )
 
-    def start_snapshot_evaluation_progress(self, snapshot: Snapshot) -> None:
-        self._write(f"Evaluating {snapshot.name}")
+    def start_snapshot_evaluation_progress(
+        self, snapshot: Snapshot, audit_only: bool = False
+    ) -> None:
+        self._write(f"{'Evaluating' if not audit_only else 'Auditing'} {snapshot.name}")
 
     def update_snapshot_evaluation_progress(
         self,
@@ -3394,10 +3421,14 @@ def update_snapshot_evaluation_progress(
         duration_ms: t.Optional[int],
         num_audits_passed: int,
         num_audits_failed: int,
+        audit_only: bool = False,
     ) -> None:
-        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}"
-        )
+        message = f"Evaluating {snapshot.name} | batch={batch_idx} | duration={duration_ms}ms | num_audits_passed={num_audits_passed} | num_audits_failed={num_audits_failed}"
+
+        if audit_only:
+            message = f"Auditing {snapshot.name} duration={duration_ms}ms | num_audits_passed={num_audits_passed} | num_audits_failed={num_audits_failed}"
+
+        self._write(message)
 
     def stop_evaluation_progress(self, success: bool = True) -> None:
         self._write(f"Stopping evaluation with success={success}")

sqlmesh/core/model/definition.py

@@ -1097,9 +1097,8 @@ def _audit_metadata_hash_values(self) -> t.List[str]:
 
         return metadata
 
-    def audit_metadata_hash(self) -> t.Tuple[t.List[str], str]:
-        hash_values = self._audit_metadata_hash_values()
-        return hash_values, hash_data(hash_values)
+    def audit_metadata_hash(self) -> str:
+        return hash_data(self._audit_metadata_hash_values())
 
     @property
     def metadata_hash(self) -> str:

sqlmesh/core/plan/evaluator.py

@@ -128,7 +128,7 @@ def evaluate(
                 execution_time=plan.execution_time,
             )
 
-            self._run_audits_for_metadata_snapshots(plan, snapshots, new_snapshots)
+            self._run_audits_for_metadata_snapshots(plan, new_snapshots)
 
             push_completion_status = self._push(plan, snapshots, deployability_index_for_creation)
             if push_completion_status.is_nothing_to_do:
@@ -550,18 +550,17 @@ def _restatement_intervals_across_all_environments(
     def _run_audits_for_metadata_snapshots(
         self,
         plan: EvaluatablePlan,
-        snapshots: t.Dict[SnapshotId, Snapshot],
         new_snapshots: t.Dict[SnapshotId, Snapshot],
     ) -> None:
-        # Step 1: Filter out snapshots that are not categorized as metadata changes on models
+        # Filter out snapshots that are not categorized as metadata changes on models
         metadata_snapshots = []
         for snapshot in new_snapshots.values():
             if not snapshot.is_metadata or not snapshot.is_model or not snapshot.evaluatable:
                 continue
 
             metadata_snapshots.append(snapshot)
 
-        # Step 2: Bulk load their previous snapshots from state
+        # Bulk load all the previous snapshots
         previous_snapshots = self.state_sync.get_snapshots(
             [
                 s.previous_version.snapshot_id(s.name)
@@ -570,29 +569,25 @@ def _run_audits_for_metadata_snapshots(
             ]
         ).values()
 
-        # Step 3: Compare the audit metadata hashes to determine if there was a change in the audits field
-        to_be_audited_snapshots = {}
+        # Check if any of the snapshots have modifications to the audits field by comparing the hashes
+        audit_snapshots = {}
         for snapshot, previous_snapshot in zip(metadata_snapshots, previous_snapshots):
-            new_audits, new_audits_hash = snapshot.model.audit_metadata_hash()
-            _, previous_audit_hash = previous_snapshot.model.audit_metadata_hash()
+            new_audits_hash = snapshot.model.audit_metadata_hash()
+            previous_audit_hash = previous_snapshot.model.audit_metadata_hash()
 
-            if previous_audit_hash != new_audits_hash and new_audits:
-                snapshot_start = min(i[0] for i in snapshot.intervals)
-                snapshot_end = max(i[1] for i in snapshot.intervals)
-                to_be_audited_snapshots[snapshot.snapshot_id] = (snapshot_start, snapshot_end)
+            if snapshot.model.audits and previous_audit_hash != new_audits_hash:
+                audit_snapshots[snapshot.snapshot_id] = snapshot
 
-        if not to_be_audited_snapshots:
+        if not audit_snapshots:
             return
 
-        # Step 4: If there are any snapshots to be audited, we'll reuse the scheduler's
-        # internals to audit them by utilizing the restatement logic
-        scheduler = self.create_scheduler(snapshots.values())
+        # If there are any snapshots to be audited, we'll reuse the scheduler's internals to audit them
+        scheduler = self.create_scheduler(audit_snapshots.values())
         completion_status = scheduler.audit(
             plan.environment,
             plan.start,
             plan.end,
             execution_time=plan.execution_time,
-            restatements=to_be_audited_snapshots,
             end_bounded=plan.end_bounded,
             interval_end_per_model=plan.interval_end_per_model,
         )