List all upstream auto-restated models

Author: treysp

sqlmesh/core/console.py

@@ -428,7 +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,
+        auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
     ) -> None:
         """Updates the snapshot evaluation progress."""
 
@@ -576,7 +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,
+        auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
     ) -> None:
         pass
 
@@ -1058,7 +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,
+        auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
     ) -> None:
         """Update the snapshot evaluation progress."""
         if (
@@ -3638,7 +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,
+        auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
     ) -> None:
         view_name, loaded_batches = self.evaluation_batch_progress[snapshot.snapshot_id]
 
@@ -3808,12 +3808,12 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
-        auto_restatement_trigger: t.Optional[SnapshotId] = None,
+        auto_restatement_triggers: t.Optional[t.List[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 auto_restatement_triggers:
+            message += f" | auto_restatement_triggers={','.join(trigger.name for trigger in auto_restatement_triggers)}"
 
         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,7 +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] = {},
+        auto_restatement_triggers: t.Dict[SnapshotId, t.List[SnapshotId]] = {},
     ) -> t.Tuple[t.List[NodeExecutionFailedError[SchedulingUnit]], t.List[SchedulingUnit]]:
         """Runs precomputed batches of missing intervals.
 
@@ -477,7 +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),
+                    auto_restatement_triggers=auto_restatement_triggers.get(snapshot.snapshot_id),
                 )
 
         try:
@@ -641,7 +641,7 @@ 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] = {}
+        auto_restatement_triggers: t.Dict[SnapshotId, t.List[SnapshotId]] = {}
         if auto_restatement_enabled:
             auto_restated_intervals, auto_restatement_triggers = apply_auto_restatements(
                 self.snapshots, execution_time

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.Tuple[t.List[SnapshotIntervals], t.Dict[SnapshotId, SnapshotId]]:
+) -> t.Tuple[t.List[SnapshotIntervals], t.Dict[SnapshotId, t.List[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,8 +2097,7 @@ 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_restatement_triggers: t.Dict[SnapshotId, t.List[SnapshotId]] = {}
     auto_restated_intervals_per_snapshot: t.Dict[SnapshotId, Interval] = {}
     for s_id in dag:
         if s_id not in snapshots:
@@ -2113,6 +2112,7 @@ def apply_auto_restatements(
             for parent_s_id in snapshot.parents
             if parent_s_id in auto_restated_intervals_per_snapshot
         ]
+        upstream_triggers = []
         if next_auto_restated_interval:
             logger.info(
                 "Calculated the next auto restated interval (%s, %s) for snapshot %s",
@@ -2123,21 +2123,15 @@ 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]
+            upstream_triggers = [s_id]
+
+        for parent_s_id in snapshot.parents:
+            if parent_s_id in auto_restatement_triggers:
+                upstream_triggers.extend(auto_restatement_triggers[parent_s_id])
+
+        # remove duplicate triggers
+        if upstream_triggers:
+            auto_restatement_triggers[s_id] = list(dict.fromkeys(upstream_triggers))
 
         if auto_restated_intervals:
             auto_restated_interval_start = sys.maxsize

tests/core/test_snapshot.py

@@ -2989,6 +2989,111 @@ def test_apply_auto_restatements_disable_restatement_downstream(make_snapshot):
     ]
 
 
+def test_auto_restatement_triggers(make_snapshot):
+    model_a = SqlModel(
+        name="test_model_a",
+        kind=IncrementalByTimeRangeKind(
+            time_column=TimeColumn(column="ds"),
+            auto_restatement_cron="0 10 * * *",
+            auto_restatement_intervals=24,
+        ),
+        start="2020-01-01",
+        cron="@daily",
+        query=parse_one("SELECT 1 as ds"),
+    )
+    snapshot_a = make_snapshot(model_a, version="1")
+    snapshot_a.add_interval("2020-01-01", "2020-01-05")
+    snapshot_a.next_auto_restatement_ts = to_timestamp("2020-01-06 10:00:00")
+
+    model_b = SqlModel(
+        name="test_model_b",
+        kind=IncrementalByTimeRangeKind(
+            time_column=TimeColumn(column="ds"),
+        ),
+        start="2020-01-01",
+        cron="@daily",
+        query=parse_one("SELECT ds FROM test_model_a"),
+    )
+    snapshot_b = make_snapshot(model_b, nodes={model_a.fqn: model_a}, version="1")
+    snapshot_b.add_interval("2020-01-01", "2020-01-05")
+
+    model_c = SqlModel(
+        name="test_model_c",
+        kind=IncrementalByTimeRangeKind(
+            time_column=TimeColumn(column="ds"),
+            auto_restatement_cron="0 10 * * *",
+            auto_restatement_intervals=24,
+        ),
+        start="2020-01-01",
+        cron="@daily",
+        query=parse_one("SELECT ds FROM test_model_a"),
+    )
+    snapshot_c = make_snapshot(model_c, nodes={model_a.fqn: model_a}, version="1")
+    snapshot_c.add_interval("2020-01-01", "2020-01-05")
+    snapshot_c.next_auto_restatement_ts = to_timestamp("2020-01-06 10:00:00")
+
+    model_d = SqlModel(
+        name="test_model_d",
+        kind=IncrementalByTimeRangeKind(
+            time_column=TimeColumn(column="ds"),
+            auto_restatement_cron="0 10 * * *",
+            auto_restatement_intervals=24,
+        ),
+        start="2020-01-01",
+        cron="@daily",
+        query=parse_one("SELECT 1 as ds"),
+    )
+    snapshot_d = make_snapshot(model_d, version="1")
+    snapshot_d.add_interval("2020-01-01", "2020-01-05")
+    snapshot_d.next_auto_restatement_ts = to_timestamp("2020-01-06 10:00:00")
+
+    model_e = SqlModel(
+        name="test_model_e",
+        kind=IncrementalByTimeRangeKind(
+            time_column=TimeColumn(column="ds"),
+        ),
+        start="2020-01-01",
+        cron="@daily",
+        query=parse_one(
+            "SELECT ds from test_model_b UNION ALL SELECT ds from test_model_c UNION ALL SELECT ds from test_model_d"
+        ),
+    )
+    snapshot_e = make_snapshot(
+        model_e,
+        nodes={
+            model_a.fqn: model_a,
+            model_b.fqn: model_b,
+            model_c.fqn: model_c,
+            model_d.fqn: model_d,
+        },
+        version="1",
+    )
+    snapshot_e.add_interval("2020-01-01", "2020-01-05")
+
+    _, auto_restatement_triggers = apply_auto_restatements(
+        {
+            snapshot_a.snapshot_id: snapshot_a,
+            snapshot_b.snapshot_id: snapshot_b,
+            snapshot_c.snapshot_id: snapshot_c,
+            snapshot_d.snapshot_id: snapshot_d,
+            snapshot_e.snapshot_id: snapshot_e,
+        },
+        "2020-01-06 10:01:00",
+    )
+
+    assert auto_restatement_triggers == {
+        snapshot_a.snapshot_id: [snapshot_a.snapshot_id],
+        snapshot_d.snapshot_id: [snapshot_d.snapshot_id],
+        snapshot_b.snapshot_id: [snapshot_a.snapshot_id],
+        snapshot_c.snapshot_id: [snapshot_c.snapshot_id, snapshot_a.snapshot_id],
+        snapshot_e.snapshot_id: [
+            snapshot_d.snapshot_id,
+            snapshot_c.snapshot_id,
+            snapshot_a.snapshot_id,
+        ],
+    }
+
+
 def test_render_signal(make_snapshot, mocker):
     @signal()
     def check_types(batch, env: str, sql: list[SQL], table: exp.Table, default: int = 0):

web/server/console.py

@@ -142,7 +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,
+        auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
     ) -> None:
         if audit_only:
             return