Collect selected snapshot triggers

Author: treysp

sqlmesh/core/console.py

@@ -3833,7 +3833,17 @@ def update_snapshot_evaluation_progress(
         audit_only: bool = False,
         snapshot_evaluation_triggers: t.Optional[SnapshotEvaluationTriggers] = 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}"
+        message = f"Evaluated {snapshot.name} | batch={batch_idx} | duration={duration_ms}ms | num_audits_passed={num_audits_passed} | num_audits_failed={num_audits_failed}"
+
+        if snapshot_evaluation_triggers:
+            if snapshot_evaluation_triggers.ignore_cron_flag is not None:
+                message += f" | ignore_cron_flag={snapshot_evaluation_triggers.ignore_cron_flag}"
+            if snapshot_evaluation_triggers.cron_ready is not None:
+                message += f" | cron_ready={snapshot_evaluation_triggers.cron_ready}"
+            if snapshot_evaluation_triggers.auto_restatement_triggers:
+                message += f" | auto_restatement_triggers={','.join(trigger.name for trigger in snapshot_evaluation_triggers.auto_restatement_triggers)}"
+            if snapshot_evaluation_triggers.select_snapshot_triggers:
+                message += f" | select_snapshot_triggers={','.join(trigger.name for trigger in snapshot_evaluation_triggers.select_snapshot_triggers)}"
 
         if snapshot_evaluation_triggers:
             if snapshot_evaluation_triggers.auto_restatement_triggers:
@@ -3842,7 +3852,7 @@ def update_snapshot_evaluation_progress(
                 message += f" | select_snapshot_triggers={','.join(trigger.name for trigger in snapshot_evaluation_triggers.select_snapshot_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}"
+            message = f"Audited {snapshot.name} duration={duration_ms}ms | num_audits_passed={num_audits_passed} | num_audits_failed={num_audits_failed}"
 
         self._write(message)
 

sqlmesh/core/context.py

@@ -2887,9 +2887,10 @@ def _select_models_for_run(
             dag.add(fqn, model.depends_on)
         model_selector = self._new_selector(models=models, dag=dag)
         result = set(model_selector.expand_model_selections(select_models))
-        if not no_auto_upstream:
-            result_with_upstream = set(dag.subdag(*result))
-        return result, result_with_upstream - result
+        if no_auto_upstream:
+            return result, set()
+        result_with_upstream = set(dag.subdag(*result))
+        return result_with_upstream, result_with_upstream - result
 
     @cached_property
     def _project_type(self) -> str:

sqlmesh/core/plan/stages.py

@@ -516,7 +516,7 @@ def _missing_intervals(
         snapshots_by_name: t.Dict[str, Snapshot],
         deployability_index: DeployabilityIndex,
     ) -> SnapshotToIntervals:
-        return merged_missing_intervals(
+        missing_intervals, _ = merged_missing_intervals(
             snapshots=snapshots_by_name.values(),
             start=plan.start,
             end=plan.end,
@@ -530,6 +530,7 @@ def _missing_intervals(
             start_override_per_model=plan.start_override_per_model,
             end_override_per_model=plan.end_override_per_model,
         )
+        return missing_intervals
 
     def _get_audit_only_snapshots(
         self, new_snapshots: t.Dict[SnapshotId, Snapshot]

sqlmesh/core/scheduler.py

@@ -112,7 +112,7 @@ def merged_missing_intervals(
         ignore_cron: bool = False,
         end_bounded: bool = False,
         selected_snapshots: t.Optional[t.Set[str]] = None,
-    ) -> SnapshotToIntervals:
+    ) -> t.Tuple[SnapshotToIntervals, t.List[SnapshotId]]:
         """Find the largest contiguous date interval parameters based only on what is missing.
 
         For each node name, find all dependencies and look for a stored snapshot from the metastore. If a snapshot is found,
@@ -132,8 +132,11 @@ def merged_missing_intervals(
             end_bounded: If set to true, the returned intervals will be bounded by the target end date, disregarding lookback,
                 allow_partials, and other attributes that could cause the intervals to exceed the target end date.
             selected_snapshots: A set of snapshot names to run. If not provided, all snapshots will be run.
+
+        Returns:
+            A tuple containing a dict containing all snapshots needing to be run with their associated interval params and a list of snapshots that are ready to run based on their naive cron schedule (ignoring plan/run context and other attributes).
         """
-        snapshots_to_intervals = merged_missing_intervals(
+        snapshots_to_intervals, snapshots_naive_cron_ready = merged_missing_intervals(
             snapshots=self.snapshot_per_version.values(),
             start=start,
             end=end,
@@ -151,7 +154,7 @@ def merged_missing_intervals(
             snapshots_to_intervals = {
                 s: i for s, i in snapshots_to_intervals.items() if s.name in selected_snapshots
             }
-        return snapshots_to_intervals
+        return snapshots_to_intervals, snapshots_naive_cron_ready
 
     def evaluate(
         self,
@@ -658,7 +661,7 @@ def _run_or_audit(
                 {s.name_version: s.next_auto_restatement_ts for s in self.snapshots.values()}
             )
 
-        merged_intervals = self.merged_missing_intervals(
+        merged_intervals, snapshots_naive_cron_ready = self.merged_missing_intervals(
             start,
             end,
             execution_time,
@@ -673,9 +676,7 @@ def _run_or_audit(
         if not merged_intervals:
             return CompletionStatus.NOTHING_TO_DO
 
-        merged_intervals_snapshots = {
-            snapshot.snapshot_id: snapshot for snapshot in merged_intervals.keys()
-        }
+        merged_intervals_snapshots = {snapshot.snapshot_id for snapshot in merged_intervals}
         select_snapshot_triggers: t.Dict[SnapshotId, t.List[SnapshotId]] = {}
         if selected_snapshots and selected_snapshots_auto_upstream:
             # actually selected snapshots are their own triggers
@@ -691,24 +692,25 @@ def _run_or_audit(
                 ]
             }
 
-            # trace upstream by reversing dag of all snapshots to evaluate
-            reversed_intervals_dag = snapshots_to_dag(merged_intervals_snapshots.values()).reversed
-            for s_id in reversed_intervals_dag:
-                if s_id not in select_snapshot_triggers:
-                    triggers = []
-                    for parent_s_id in merged_intervals_snapshots[s_id].parents:
-                        triggers.extend(select_snapshot_triggers[parent_s_id])
+            # trace upstream by walking downstream on reversed dag
+            reversed_dag = snapshots_to_dag(self.snapshots.values()).reversed
+            for s_id in reversed_dag:
+                if s_id in merged_intervals_snapshots:
+                    triggers = select_snapshot_triggers.get(s_id, [])
+                    for parent_s_id in reversed_dag.graph.get(s_id, set()):
+                        triggers.extend(select_snapshot_triggers.get(parent_s_id, []))
                     select_snapshot_triggers[s_id] = list(dict.fromkeys(triggers))
 
         all_snapshot_triggers: t.Dict[SnapshotId, SnapshotEvaluationTriggers] = {
             s_id: SnapshotEvaluationTriggers(
-                ignore_cron=ignore_cron,
+                ignore_cron_flag=ignore_cron,
+                cron_ready=s_id in snapshots_naive_cron_ready,
                 auto_restatement_triggers=auto_restatement_triggers.get(s_id, []),
                 select_snapshot_triggers=select_snapshot_triggers.get(s_id, []),
             )
             for s_id in merged_intervals_snapshots
-            if ignore_cron or s_id in auto_restatement_triggers or s_id in select_snapshot_triggers
         }
+
         errors, _ = self.run_merged_intervals(
             merged_intervals=merged_intervals,
             deployability_index=deployability_index,
@@ -870,7 +872,7 @@ def merged_missing_intervals(
     end_override_per_model: t.Optional[t.Dict[str, datetime]] = None,
     ignore_cron: bool = False,
     end_bounded: bool = False,
-) -> SnapshotToIntervals:
+) -> t.Tuple[SnapshotToIntervals, t.List[SnapshotId]]:
     """Find the largest contiguous date interval parameters based only on what is missing.
 
     For each node name, find all dependencies and look for a stored snapshot from the metastore. If a snapshot is found,
@@ -920,7 +922,7 @@ def compute_interval_params(
     end_override_per_model: t.Optional[t.Dict[str, datetime]] = None,
     ignore_cron: bool = False,
     end_bounded: bool = False,
-) -> SnapshotToIntervals:
+) -> t.Tuple[SnapshotToIntervals, t.List[SnapshotId]]:
     """Find the largest contiguous date interval parameters based only on what is missing.
 
     For each node name, find all dependencies and look for a stored snapshot from the metastore. If a snapshot is found,
@@ -942,7 +944,7 @@ def compute_interval_params(
             allow_partials, and other attributes that could cause the intervals to exceed the target end date.
 
     Returns:
-        A dict containing all snapshots needing to be run with their associated interval params.
+        A tuple containing a dict containing all snapshots needing to be run with their associated interval params and a list of snapshots that are ready to run based on their naive cron schedule (ignoring plan/run context and other attributes).
     """
     snapshot_merged_intervals = {}
 
@@ -970,7 +972,11 @@ def compute_interval_params(
             contiguous_batch.append((next_batch[0][0], next_batch[-1][-1]))
         snapshot_merged_intervals[snapshot] = contiguous_batch
 
-    return snapshot_merged_intervals
+    snapshots_naive_cron_ready = [
+        snap.snapshot_id for snap in missing_intervals(snapshots, execution_time=execution_time)
+    ]
+
+    return snapshot_merged_intervals, snapshots_naive_cron_ready
 
 
 def interval_diff(

sqlmesh/core/snapshot/definition.py

@@ -326,11 +326,10 @@ def table_name_for_environment(
 
 
 class SnapshotEvaluationTriggers(PydanticModel):
-    ignore_cron: bool
+    ignore_cron_flag: t.Optional[bool] = None
+    cron_ready: t.Optional[bool] = None
     auto_restatement_triggers: t.List[SnapshotId] = []
     select_snapshot_triggers: t.List[SnapshotId] = []
-    directly_modified_triggers: t.List[SnapshotId] = []
-    manual_restatement_triggers: t.List[SnapshotId] = []
 
 
 class SnapshotInfoMixin(ModelKindMixin):

tests/core/test_integration.py

@@ -1779,6 +1779,77 @@ def test_select_unchanged_model_for_backfill(init_and_plan_context: t.Callable):
     assert {o.name for o in schema_objects} == {"waiter_revenue_by_day", "top_waiters"}
 
 
+@time_machine.travel("2023-01-08 00:00:00 UTC")
+def test_snapshot_triggers(init_and_plan_context: t.Callable, mocker: MockerFixture):
+    context, plan = init_and_plan_context("examples/sushi")
+    context.apply(plan)
+
+    # add auto restatement to orders
+    model = context.get_model("sushi.orders")
+    kind = {
+        **model.kind.dict(),
+        "auto_restatement_cron": "@hourly",
+    }
+    kwargs = {
+        **model.dict(),
+        "kind": kind,
+    }
+    context.upsert_model(PythonModel.parse_obj(kwargs))
+    plan = context.plan_builder(skip_tests=True).build()
+    context.apply(plan)
+
+    # Mock run_merged_intervals to capture triggers arg
+    scheduler = context.scheduler()
+    run_merged_intervals_mock = mocker.patch.object(
+        scheduler, "run_merged_intervals", return_value=([], [])
+    )
+
+    # User selects top_waiters and waiter_revenue_by_day, others added as auto-upstream
+    selected_models = {"top_waiters", "waiter_revenue_by_day"}
+    selected_models_auto_upstream = {"order_items", "orders", "items"}
+    selected_snapshots = {
+        f'"memory"."sushi"."{model}"' for model in selected_models | selected_models_auto_upstream
+    }
+    selected_snapshots_auto_upstream = selected_snapshots - {
+        f'"memory"."sushi"."{model}"' for model in selected_models
+    }
+
+    with time_machine.travel("2023-01-09 00:00:01 UTC"):
+        scheduler.run(
+            environment=c.PROD,
+            selected_snapshots=selected_snapshots,
+            selected_snapshots_auto_upstream=selected_snapshots_auto_upstream,
+            start="2023-01-01",
+            auto_restatement_enabled=True,
+        )
+
+    assert run_merged_intervals_mock.called
+
+    actual_triggers = run_merged_intervals_mock.call_args.kwargs["snapshot_evaluation_triggers"]
+
+    # validate ignore_cron not passed and all model crons ready
+    assert all(
+        not trigger.ignore_cron_flag and trigger.cron_ready for trigger in actual_triggers.values()
+    )
+
+    for id, trigger in actual_triggers.items():
+        # top_waiters is its own trigger, waiter_revenue_by_day is upstream of it, everyone else is upstream of both
+        select_triggers = [t.name for t in trigger.select_snapshot_triggers]
+        assert (
+            select_triggers == ['"memory"."sushi"."top_waiters"']
+            if id.name == '"memory"."sushi"."top_waiters"'
+            else ['"memory"."sushi"."waiter_revenue_by_day"', '"memory"."sushi"."top_waiters"']
+        )
+
+        # everyone other than items is downstream of orders
+        auto_restatement_triggers = [t.name for t in trigger.auto_restatement_triggers]
+        assert (
+            auto_restatement_triggers == []
+            if id.name == '"memory"."sushi"."items"'
+            else ['"memory"."sushi"."orders"']
+        )
+
+
 @time_machine.travel("2023-01-08 15:00:00 UTC")
 def test_max_interval_end_per_model_not_applied_when_end_is_provided(
     init_and_plan_context: t.Callable,
@@ -6597,13 +6668,13 @@ def plan_with_output(ctx: Context, environment: str):
     assert "Differences from the `prod` environment" in output.stdout
 
     assert (
-        """MODEL (                                                                        
-   name test.a,                                                                 
-+  owner test,                                                                  
-   kind FULL                                                                    
- )                                                                              
- SELECT                                                                         
--  5 AS col                                                                     
+        """MODEL (
+   name test.a,
++  owner test,
+   kind FULL
+ )
+ SELECT
+-  5 AS col
 +  10 AS col"""
         in output.stdout
     )

tests/core/test_scheduler.py

@@ -76,7 +76,7 @@ def _get_batched_missing_intervals(
         end: TimeLike,
         execution_time: t.Optional[TimeLike] = None,
     ) -> SnapshotToIntervals:
-        merged_intervals = scheduler.merged_missing_intervals(start, end, execution_time)
+        merged_intervals, _ = scheduler.merged_missing_intervals(start, end, execution_time)
         return scheduler.batch_intervals(merged_intervals, mocker.Mock(), mocker.Mock())
 
     return _get_batched_missing_intervals
@@ -104,9 +104,10 @@ def test_interval_params_missing(scheduler: Scheduler, sushi_context_fixed_date:
 
     start_ds = "2022-01-01"
     end_ds = "2022-03-01"
-    assert compute_interval_params(
+    interval_params, _ = compute_interval_params(
         sushi_context_fixed_date.snapshots.values(), start=start_ds, end=end_ds
-    )[waiters] == [
+    )
+    assert interval_params[waiters] == [
         (to_timestamp(start_ds), to_timestamp("2022-03-02")),
     ]
 

web/server/api/endpoints/plan.py

@@ -132,7 +132,7 @@ def _get_plan_changes(context: Context, plan: Plan) -> models.PlanChanges:
 
 def _get_plan_backfills(context: Context, plan: Plan) -> t.Dict[str, t.Any]:
     """Get plan backfills"""
-    merged_intervals = context.scheduler().merged_missing_intervals()
+    merged_intervals, _ = context.scheduler().merged_missing_intervals()
     batches = context.scheduler().batch_intervals(merged_intervals, None, EnvironmentNamingInfo())
     tasks = {snapshot.name: len(intervals) for snapshot, intervals in batches.items()}
     snapshots = plan.context_diff.snapshots