Print auto-restated trigger of model evaluation in debug console

Author: treysp

sqlmesh/core/console.py

@@ -428,6 +428,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
+        auto_restatement_trigger: t.Optional[SnapshotId] = None,
     ) -> None:
         """Updates the snapshot evaluation progress."""
 
@@ -575,6 +576,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
+        auto_restatement_trigger: t.Optional[SnapshotId] = None,
     ) -> None:
         pass
 
@@ -1056,6 +1058,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
+        auto_restatement_trigger: t.Optional[SnapshotId] = None,
     ) -> None:
         """Update the snapshot evaluation progress."""
         if (
@@ -3635,6 +3638,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
+        auto_restatement_trigger: t.Optional[SnapshotId] = None,
     ) -> None:
         view_name, loaded_batches = self.evaluation_batch_progress[snapshot.snapshot_id]
 
@@ -3804,9 +3808,13 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
+        auto_restatement_trigger: t.Optional[SnapshotId] = None,
     ) -> None:
         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 auto_restatement_trigger:
+            message += f" | evaluation_triggered_by={auto_restatement_trigger.name}"
+
         if audit_only:
             message = f"Auditing {snapshot.name} duration={duration_ms}ms | num_audits_passed={num_audits_passed} | num_audits_failed={num_audits_failed}"
 

sqlmesh/core/scheduler.py

@@ -374,6 +374,7 @@ def run_merged_intervals(
         run_environment_statements: bool = False,
         audit_only: bool = False,
         restatements: t.Optional[t.Dict[SnapshotId, Interval]] = None,
+        auto_restatement_triggers: t.Dict[SnapshotId, SnapshotId] = {},
     ) -> t.Tuple[t.List[NodeExecutionFailedError[SchedulingUnit]], t.List[SchedulingUnit]]:
         """Runs precomputed batches of missing intervals.
 
@@ -476,6 +477,7 @@ def evaluate_node(node: SchedulingUnit) -> None:
                     evaluation_duration_ms,
                     num_audits - num_audits_failed,
                     num_audits_failed,
+                    auto_restatement_trigger=auto_restatement_triggers.get(snapshot.snapshot_id),
                 )
 
         try:
@@ -639,8 +641,11 @@ def _run_or_audit(
         for s_id, interval in (remove_intervals or {}).items():
             self.snapshots[s_id].remove_interval(interval)
 
+        auto_restatement_triggers: t.Dict[SnapshotId, SnapshotId] = {}
         if auto_restatement_enabled:
-            auto_restated_intervals = apply_auto_restatements(self.snapshots, execution_time)
+            auto_restated_intervals, auto_restatement_triggers = apply_auto_restatements(
+                self.snapshots, execution_time
+            )
             self.state_sync.add_snapshots_intervals(auto_restated_intervals)
             self.state_sync.update_auto_restatements(
                 {s.name_version: s.next_auto_restatement_ts for s in self.snapshots.values()}
@@ -672,6 +677,7 @@ def _run_or_audit(
             run_environment_statements=run_environment_statements,
             audit_only=audit_only,
             restatements=remove_intervals,
+            auto_restatement_triggers=auto_restatement_triggers,
         )
 
         return CompletionStatus.FAILURE if errors else CompletionStatus.SUCCESS

sqlmesh/core/snapshot/definition.py

@@ -2082,7 +2082,7 @@ def snapshots_to_dag(snapshots: t.Collection[Snapshot]) -> DAG[SnapshotId]:
 
 def apply_auto_restatements(
     snapshots: t.Dict[SnapshotId, Snapshot], execution_time: TimeLike
-) -> t.List[SnapshotIntervals]:
+) -> t.Tuple[t.List[SnapshotIntervals], t.Dict[SnapshotId, SnapshotId]]:
     """Applies auto restatements to the snapshots.
 
     This operation results in the removal of intervals for snapshots that are ready to be restated based
@@ -2097,6 +2097,8 @@ def apply_auto_restatements(
         A list of SnapshotIntervals with **new** intervals that need to be restated.
     """
     dag = snapshots_to_dag(snapshots.values())
+    snapshots_with_auto_restatements: t.List[SnapshotId] = []
+    auto_restatement_triggers: t.Dict[SnapshotId, SnapshotId] = {}
     auto_restated_intervals_per_snapshot: t.Dict[SnapshotId, Interval] = {}
     for s_id in dag:
         if s_id not in snapshots:
@@ -2120,6 +2122,23 @@ def apply_auto_restatements(
             )
             auto_restated_intervals.append(next_auto_restated_interval)
 
+            # auto-restated snapshot is its own trigger
+            snapshots_with_auto_restatements.append(s_id)
+            auto_restatement_triggers[s_id] = s_id
+        else:
+            for parent_s_id in snapshot.parents:
+                # first auto-restated parent is the trigger
+                if parent_s_id in snapshots_with_auto_restatements:
+                    auto_restatement_triggers[s_id] = parent_s_id
+                    break
+                # if no trigger yet and parent has trigger, inherit their trigger
+                # - will be overwritten if a different parent is auto-restated
+                if (
+                    parent_s_id in auto_restatement_triggers
+                    and s_id not in auto_restatement_triggers
+                ):
+                    auto_restatement_triggers[s_id] = auto_restatement_triggers[parent_s_id]
+
         if auto_restated_intervals:
             auto_restated_interval_start = sys.maxsize
             auto_restated_interval_end = -sys.maxsize
@@ -2149,20 +2168,22 @@ def apply_auto_restatements(
 
         snapshot.apply_pending_restatement_intervals()
         snapshot.update_next_auto_restatement_ts(execution_time)
-
-    return [
-        SnapshotIntervals(
-            name=snapshots[s_id].name,
-            identifier=None,
-            version=snapshots[s_id].version,
-            dev_version=None,
-            intervals=[],
-            dev_intervals=[],
-            pending_restatement_intervals=[interval],
-        )
-        for s_id, interval in auto_restated_intervals_per_snapshot.items()
-        if s_id in snapshots
-    ]
+    return (
+        [
+            SnapshotIntervals(
+                name=snapshots[s_id].name,
+                identifier=None,
+                version=snapshots[s_id].version,
+                dev_version=None,
+                intervals=[],
+                dev_intervals=[],
+                pending_restatement_intervals=[interval],
+            )
+            for s_id, interval in auto_restated_intervals_per_snapshot.items()
+            if s_id in snapshots
+        ],
+        auto_restatement_triggers,
+    )
 
 
 def parent_snapshots_by_name(

tests/core/test_snapshot.py

@@ -2812,7 +2812,7 @@ def test_apply_auto_restatements(make_snapshot):
         (to_timestamp("2020-01-01"), to_timestamp("2020-01-06")),
     ]
 
-    restated_intervals = apply_auto_restatements(
+    restated_intervals, _ = apply_auto_restatements(
         {
             snapshot_a.snapshot_id: snapshot_a,
             snapshot_b.snapshot_id: snapshot_b,
@@ -2949,7 +2949,7 @@ def test_apply_auto_restatements_disable_restatement_downstream(make_snapshot):
     snapshot_b.add_interval("2020-01-01", "2020-01-05")
     assert snapshot_a.snapshot_id in snapshot_b.parents
 
-    restated_intervals = apply_auto_restatements(
+    restated_intervals, _ = apply_auto_restatements(
         {
             snapshot_a.snapshot_id: snapshot_a,
             snapshot_b.snapshot_id: snapshot_b,

web/server/console.py

@@ -9,7 +9,7 @@
 from sqlmesh.core.console import TerminalConsole
 from sqlmesh.core.environment import EnvironmentNamingInfo
 from sqlmesh.core.plan.definition import EvaluatablePlan
-from sqlmesh.core.snapshot import Snapshot, SnapshotInfoLike, SnapshotTableInfo
+from sqlmesh.core.snapshot import Snapshot, SnapshotInfoLike, SnapshotTableInfo, SnapshotId
 from sqlmesh.core.test import ModelTest
 from sqlmesh.core.test.result import ModelTextTestResult
 from sqlmesh.utils.date import now_timestamp
@@ -142,6 +142,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
+        auto_restatement_trigger: t.Optional[SnapshotId] = None,
     ) -> None:
         if audit_only:
             return