Feat!: Ensure metadata snapshots with modified audits are still audited (#4341)

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,8 +3254,14 @@ def update_snapshot_evaluation_progress(
         duration_ms: t.Optional[int],
         num_audits_passed: int,
         num_audits_failed: int,
+        audit_only: bool = False,
     ) -> None:
         view_name, loaded_batches = self.evaluation_batch_progress[snapshot.snapshot_id]
+
+        if audit_only:
+            print(f"Completed Auditing {view_name}")
+            return
+
         total_batches = self.evaluation_model_batch_sizes[snapshot]
 
         loaded_batches += 1
@@ -3378,13 +3403,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 +3423,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

@@ -1069,6 +1069,37 @@ def _data_hash_values(self) -> t.List[str]:
 
         return data  # type: ignore
 
+    def _audit_metadata_hash_values(self) -> t.List[str]:
+        from sqlmesh.core.audit.builtin import BUILT_IN_AUDITS
+
+        metadata = []
+
+        for audit_name, audit_args in sorted(self.audits, key=lambda a: a[0]):
+            metadata.append(audit_name)
+            if audit_name in BUILT_IN_AUDITS:
+                for arg_name, arg_value in audit_args.items():
+                    metadata.append(arg_name)
+                    metadata.append(gen(arg_value))
+            else:
+                audit = self.audit_definitions[audit_name]
+                query = (
+                    self.render_audit_query(audit, **t.cast(t.Dict[str, t.Any], audit_args))
+                    or audit.query
+                )
+                metadata.extend(
+                    [
+                        gen(query),
+                        audit.dialect,
+                        str(audit.skip),
+                        str(audit.blocking),
+                    ]
+                )
+
+        return metadata
+
+    def audit_metadata_hash(self) -> str:
+        return hash_data(self._audit_metadata_hash_values())
+
     @property
     def metadata_hash(self) -> str:
         """
@@ -1078,8 +1109,6 @@ def metadata_hash(self) -> str:
             The metadata hash for the node.
         """
         if self._metadata_hash is None:
-            from sqlmesh.core.audit.builtin import BUILT_IN_AUDITS
-
             metadata = [
                 self.dialect,
                 self.owner,
@@ -1100,29 +1129,9 @@ def metadata_hash(self) -> str:
                 str(self.allow_partials),
                 gen(self.session_properties_) if self.session_properties_ else None,
                 *[gen(g) for g in self.grains],
+                *self._audit_metadata_hash_values(),
             ]
 
-            for audit_name, audit_args in sorted(self.audits, key=lambda a: a[0]):
-                metadata.append(audit_name)
-                if audit_name in BUILT_IN_AUDITS:
-                    for arg_name, arg_value in audit_args.items():
-                        metadata.append(arg_name)
-                        metadata.append(gen(arg_value))
-                else:
-                    audit = self.audit_definitions[audit_name]
-                    query = (
-                        self.render_audit_query(audit, **t.cast(t.Dict[str, t.Any], audit_args))
-                        or audit.query
-                    )
-                    metadata.extend(
-                        [
-                            gen(query),
-                            audit.dialect,
-                            str(audit.skip),
-                            str(audit.blocking),
-                        ]
-                    )
-
             for key, value in (self.virtual_properties or {}).items():
                 metadata.append(key)
                 metadata.append(gen(value))

sqlmesh/core/plan/evaluator.py

@@ -128,6 +128,8 @@ def evaluate(
                 execution_time=plan.execution_time,
             )
 
+            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:
                 self.console.log_status_update(
@@ -545,6 +547,54 @@ def _restatement_intervals_across_all_environments(
 
         return set(snapshots_to_restate.values())
 
+    def _run_audits_for_metadata_snapshots(
+        self,
+        plan: EvaluatablePlan,
+        new_snapshots: t.Dict[SnapshotId, Snapshot],
+    ) -> None:
+        # 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)
+
+        # Bulk load all the previous snapshots
+        previous_snapshots = self.state_sync.get_snapshots(
+            [
+                s.previous_version.snapshot_id(s.name)
+                for s in metadata_snapshots
+                if s.previous_version
+            ]
+        ).values()
+
+        # 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_hash = snapshot.model.audit_metadata_hash()
+            previous_audit_hash = previous_snapshot.model.audit_metadata_hash()
+
+            if snapshot.model.audits and previous_audit_hash != new_audits_hash:
+                audit_snapshots[snapshot.snapshot_id] = snapshot
+
+        if not audit_snapshots:
+            return
+
+        # 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,
+            end_bounded=plan.end_bounded,
+            interval_end_per_model=plan.interval_end_per_model,
+        )
+
+        if completion_status.is_failure:
+            raise PlanError("Plan application failed.")
+
 
 def update_intervals_for_new_snapshots(
     snapshots: t.Collection[Snapshot], state_sync: StateSync