Infer cron ready from auto-restatement intervals

Author: treysp

sqlmesh/core/plan/stages.py

@@ -504,7 +504,7 @@ def _missing_intervals(
         snapshots_by_name: t.Dict[str, Snapshot],
         deployability_index: DeployabilityIndex,
     ) -> SnapshotToIntervals:
-        missing_intervals, _ = merged_missing_intervals(
+        missing_intervals = merged_missing_intervals(
             snapshots=snapshots_by_name.values(),
             start=plan.start,
             end=plan.end,

sqlmesh/core/scheduler.py

@@ -26,10 +26,11 @@
     snapshots_to_dag,
     Intervals,
 )
-from sqlmesh.core.snapshot.definition import check_ready_intervals
 from sqlmesh.core.snapshot.definition import (
     Interval,
     SnapshotEvaluationTriggers,
+    SnapshotIntervals,
+    check_ready_intervals,
     expand_range,
     parent_snapshots_by_name,
 )
@@ -112,7 +113,7 @@ def merged_missing_intervals(
         ignore_cron: bool = False,
         end_bounded: bool = False,
         selected_snapshots: t.Optional[t.Set[str]] = None,
-    ) -> t.Tuple[SnapshotToIntervals, t.List[SnapshotId]]:
+    ) -> SnapshotToIntervals:
         """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,
@@ -134,9 +135,9 @@ def merged_missing_intervals(
             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).
+            A dict containing all snapshots needing to be run with their associated interval params.
         """
-        snapshots_to_intervals, snapshots_naive_cron_ready = merged_missing_intervals(
+        snapshots_to_intervals = merged_missing_intervals(
             snapshots=self.snapshot_per_version.values(),
             start=start,
             end=end,
@@ -154,7 +155,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, snapshots_naive_cron_ready
+        return snapshots_to_intervals
 
     def evaluate(
         self,
@@ -651,6 +652,7 @@ def _run_or_audit(
         for s_id, interval in (remove_intervals or {}).items():
             self.snapshots[s_id].remove_interval(interval)
 
+        auto_restated_intervals: t.List[SnapshotIntervals] = []
         auto_restatement_triggers: t.Dict[SnapshotId, t.List[SnapshotId]] = {}
         if auto_restatement_enabled:
             auto_restated_intervals, auto_restatement_triggers = apply_auto_restatements(
@@ -660,8 +662,9 @@ def _run_or_audit(
             self.state_sync.update_auto_restatements(
                 {s.name_version: s.next_auto_restatement_ts for s in self.snapshots.values()}
             )
+        auto_restated_snapshots = {snapshot.snapshot_id for snapshot in auto_restated_intervals}
 
-        merged_intervals, snapshots_naive_cron_ready = self.merged_missing_intervals(
+        merged_intervals = self.merged_missing_intervals(
             start,
             end,
             execution_time,
@@ -704,7 +707,7 @@ def _run_or_audit(
         all_snapshot_triggers: t.Dict[SnapshotId, SnapshotEvaluationTriggers] = {
             s_id: SnapshotEvaluationTriggers(
                 ignore_cron_flag=ignore_cron,
-                cron_ready=s_id in snapshots_naive_cron_ready,
+                cron_ready=s_id not in auto_restated_snapshots,
                 auto_restatement_triggers=auto_restatement_triggers.get(s_id, []),
                 select_snapshot_triggers=select_snapshot_triggers.get(s_id, []),
             )
@@ -872,7 +875,7 @@ def merged_missing_intervals(
     end_override_per_model: t.Optional[t.Dict[str, datetime]] = None,
     ignore_cron: bool = False,
     end_bounded: bool = False,
-) -> t.Tuple[SnapshotToIntervals, t.List[SnapshotId]]:
+) -> SnapshotToIntervals:
     """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,
@@ -922,7 +925,7 @@ def compute_interval_params(
     end_override_per_model: t.Optional[t.Dict[str, datetime]] = None,
     ignore_cron: bool = False,
     end_bounded: bool = False,
-) -> t.Tuple[SnapshotToIntervals, t.List[SnapshotId]]:
+) -> SnapshotToIntervals:
     """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,
@@ -944,7 +947,7 @@ def compute_interval_params(
             allow_partials, and other attributes that could cause the intervals to exceed the target end date.
 
     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).
+        A dict containing all snapshots needing to be run with their associated interval params.
     """
     snapshot_merged_intervals = {}
 
@@ -972,11 +975,7 @@ def compute_interval_params(
             contiguous_batch.append((next_batch[0][0], next_batch[-1][-1]))
         snapshot_merged_intervals[snapshot] = contiguous_batch
 
-    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
+    return snapshot_merged_intervals
 
 
 def interval_diff(

tests/core/test_scheduler.py

@@ -56,9 +56,7 @@ def test_interval_params(scheduler: Scheduler, sushi_context_fixed_date: Context
     start_ds = "2022-01-01"
     end_ds = "2022-02-05"
 
-    interval_params, _ = compute_interval_params(
-        [orders, waiter_revenue], start=start_ds, end=end_ds
-    )
+    interval_params = compute_interval_params([orders, waiter_revenue], start=start_ds, end=end_ds)
     assert interval_params == {
         orders: [
             (to_timestamp(start_ds), to_timestamp("2022-02-06")),
@@ -79,7 +77,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
@@ -91,7 +89,7 @@ def test_interval_params_nonconsecutive(scheduler: Scheduler, orders: Snapshot):
 
     orders.add_interval("2022-01-10", "2022-01-15")
 
-    interval_params, _ = compute_interval_params([orders], start=start_ds, end=end_ds)
+    interval_params = compute_interval_params([orders], start=start_ds, end=end_ds)
     assert interval_params == {
         orders: [
             (to_timestamp(start_ds), to_timestamp("2022-01-10")),
@@ -108,7 +106,7 @@ def test_interval_params_missing(scheduler: Scheduler, sushi_context_fixed_date:
 
     start_ds = "2022-01-01"
     end_ds = "2022-03-01"
-    interval_params, _ = compute_interval_params(
+    interval_params = compute_interval_params(
         sushi_context_fixed_date.snapshots.values(), start=start_ds, end=end_ds
     )
     assert interval_params[waiters] == [

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