Report auto-restatement triggers only

Author: treysp

sqlmesh/core/console.py

@@ -37,12 +37,7 @@
     SnapshotId,
     SnapshotInfoLike,
 )
-from sqlmesh.core.snapshot.definition import (
-    Interval,
-    Intervals,
-    SnapshotTableInfo,
-    SnapshotEvaluationTriggers,
-)
+from sqlmesh.core.snapshot.definition import Interval, Intervals, SnapshotTableInfo
 from sqlmesh.core.test import ModelTest
 from sqlmesh.utils import rich as srich
 from sqlmesh.utils import Verbosity
@@ -433,7 +428,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
-        snapshot_evaluation_triggers: t.Optional[SnapshotEvaluationTriggers] = None,
+        auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
     ) -> None:
         """Updates the snapshot evaluation progress."""
 
@@ -581,7 +576,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
-        snapshot_evaluation_triggers: t.Optional[SnapshotEvaluationTriggers] = None,
+        auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
     ) -> None:
         pass
 
@@ -1063,7 +1058,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
-        snapshot_evaluation_triggers: t.Optional[SnapshotEvaluationTriggers] = None,
+        auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
     ) -> None:
         """Update the snapshot evaluation progress."""
         if (
@@ -3647,7 +3642,7 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
-        snapshot_evaluation_triggers: t.Optional[SnapshotEvaluationTriggers] = None,
+        auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
     ) -> None:
         view_name, loaded_batches = self.evaluation_batch_progress[snapshot.snapshot_id]
 
@@ -3817,22 +3812,15 @@ def update_snapshot_evaluation_progress(
         num_audits_passed: int,
         num_audits_failed: int,
         audit_only: bool = False,
-        snapshot_evaluation_triggers: t.Optional[SnapshotEvaluationTriggers] = None,
+        auto_restatement_triggers: t.Optional[t.List[SnapshotId]] = None,
     ) -> None:
         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 auto_restatement_triggers:
+            message += f" | Auto-restatement triggers {', '.join(trigger.name for trigger in auto_restatement_triggers)}"
 
         if audit_only:
-            message = f"Audited {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

@@ -2299,9 +2299,11 @@ def check_intervals(
         }
 
         if select_models:
-            selected, _ = self._select_models_for_run(select_models, True, snapshots.values())
+            selected: t.Collection[str] = self._select_models_for_run(
+                select_models, True, snapshots.values()
+            )
         else:
-            selected = set(snapshots.keys())
+            selected = snapshots.keys()
 
         results = {}
         execution_context = self.execution_context(snapshots=snapshots)
@@ -2451,9 +2453,8 @@ def _run(
         scheduler = self.scheduler(environment=environment)
         snapshots = scheduler.snapshots
 
-        select_models_auto_upstream = None
         if select_models is not None:
-            select_models, select_models_auto_upstream = self._select_models_for_run(
+            select_models = self._select_models_for_run(
                 select_models, no_auto_upstream, snapshots.values()
             )
 
@@ -2465,7 +2466,6 @@ def _run(
             ignore_cron=ignore_cron,
             circuit_breaker=circuit_breaker,
             selected_snapshots=select_models,
-            selected_snapshots_auto_upstream=select_models_auto_upstream,
             auto_restatement_enabled=environment.lower() == c.PROD,
             run_environment_statements=True,
         )
@@ -2881,7 +2881,7 @@ def _select_models_for_run(
         select_models: t.Collection[str],
         no_auto_upstream: bool,
         snapshots: t.Collection[Snapshot],
-    ) -> t.Tuple[t.Set[str], t.Set[str]]:
+    ) -> t.Set[str]:
         models: UniqueKeyDict[str, Model] = UniqueKeyDict(
             "models", **{s.name: s.model for s in snapshots if s.is_model}
         )
@@ -2890,10 +2890,9 @@ 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 no_auto_upstream:
-            return result, set()
-        result_with_upstream = set(dag.subdag(*result))
-        return result_with_upstream, result_with_upstream - result
+        if not no_auto_upstream:
+            result = set(dag.subdag(*result))
+        return result
 
     @cached_property
     def _project_type(self) -> str:

sqlmesh/core/plan/stages.py

@@ -513,7 +513,7 @@ def _missing_intervals(
         snapshots_by_name: t.Dict[str, Snapshot],
         deployability_index: DeployabilityIndex,
     ) -> SnapshotToIntervals:
-        missing_intervals = merged_missing_intervals(
+        return merged_missing_intervals(
             snapshots=snapshots_by_name.values(),
             start=plan.start,
             end=plan.end,
@@ -527,7 +527,6 @@ 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

@@ -26,11 +26,10 @@
     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,
 )
@@ -133,9 +132,6 @@ 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 dict containing all snapshots needing to be run with their associated interval params.
         """
         snapshots_to_intervals = merged_missing_intervals(
             snapshots=self.snapshot_per_version.values(),
@@ -226,7 +222,6 @@ def run(
         ignore_cron: bool = False,
         end_bounded: bool = False,
         selected_snapshots: t.Optional[t.Set[str]] = None,
-        selected_snapshots_auto_upstream: t.Optional[t.Set[str]] = None,
         circuit_breaker: t.Optional[t.Callable[[], bool]] = None,
         deployability_index: t.Optional[DeployabilityIndex] = None,
         auto_restatement_enabled: bool = False,
@@ -243,7 +238,6 @@ def run(
             ignore_cron=ignore_cron,
             end_bounded=end_bounded,
             selected_snapshots=selected_snapshots,
-            selected_snapshots_auto_upstream=selected_snapshots_auto_upstream,
             circuit_breaker=circuit_breaker,
             deployability_index=deployability_index,
             auto_restatement_enabled=auto_restatement_enabled,
@@ -378,6 +372,7 @@ def run_merged_intervals(
         end: t.Optional[TimeLike] = None,
         run_environment_statements: bool = False,
         audit_only: bool = False,
+        restatements: t.Optional[t.Dict[SnapshotId, Interval]] = None,
         auto_restatement_triggers: t.Dict[SnapshotId, t.List[SnapshotId]] = {},
     ) -> t.Tuple[t.List[NodeExecutionFailedError[SchedulingUnit]], t.List[SchedulingUnit]]:
         """Runs precomputed batches of missing intervals.
@@ -476,9 +471,7 @@ def evaluate_node(node: SchedulingUnit) -> None:
                     evaluation_duration_ms,
                     num_audits - num_audits_failed,
                     num_audits_failed,
-                    snapshot_evaluation_triggers=snapshot_evaluation_triggers.get(
-                        snapshot.snapshot_id
-                    ),
+                    auto_restatement_triggers=auto_restatement_triggers.get(snapshot.snapshot_id),
                 )
 
         try:
@@ -589,7 +582,6 @@ def _run_or_audit(
         ignore_cron: bool = False,
         end_bounded: bool = False,
         selected_snapshots: t.Optional[t.Set[str]] = None,
-        selected_snapshots_auto_upstream: t.Optional[t.Set[str]] = None,
         circuit_breaker: t.Optional[t.Callable[[], bool]] = None,
         deployability_index: t.Optional[DeployabilityIndex] = None,
         auto_restatement_enabled: bool = False,
@@ -613,7 +605,6 @@ def _run_or_audit(
             end_bounded: If set to true, the evaluated 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.
-            selected_snapshots_auto_upstream: The set of selected_snapshots that were automatically added because they're upstream of a selected snapshot.
             circuit_breaker: An optional handler which checks if the run should be aborted.
             deployability_index: Determines snapshots that are deployable in the context of this render.
             auto_restatement_enabled: Whether to enable auto restatements.
@@ -672,38 +663,9 @@ def _run_or_audit(
             return CompletionStatus.NOTHING_TO_DO
 
         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
-            selected_snapshots_no_auto_upstream = (
-                selected_snapshots - selected_snapshots_auto_upstream
-            )
-            select_snapshot_triggers = {
-                s_id: [s_id]
-                for s_id in [
-                    snapshot_id
-                    for snapshot_id in merged_intervals_snapshots
-                    if snapshot_id.name in selected_snapshots_no_auto_upstream
-                ]
-            }
 
-            # 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_flag=ignore_cron,
-                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, []),
-            )
-            for s_id in merged_intervals_snapshots
+        auto_restatement_triggers_dict: t.Dict[SnapshotId, t.List[SnapshotId]] = {
+            s_id: auto_restatement_triggers.get(s_id, []) for s_id in merged_intervals_snapshots
         }
 
         errors, _ = self.run_merged_intervals(

sqlmesh/core/snapshot/definition.py

@@ -21,7 +21,7 @@
 from sqlmesh.core.model import Model, ModelKindMixin, ModelKindName, ViewKind, CustomKind
 from sqlmesh.core.model.definition import _Model
 from sqlmesh.core.node import IntervalUnit, NodeType
-from sqlmesh.utils import sanitize_name
+from sqlmesh.utils import sanitize_name, unique
 from sqlmesh.utils.dag import DAG
 from sqlmesh.utils.date import (
     TimeLike,
@@ -327,13 +327,6 @@ def table_name_for_environment(
         return table
 
 
-class SnapshotEvaluationTriggers(PydanticModel):
-    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] = []
-
-
 class SnapshotInfoMixin(ModelKindMixin):
     name: str
     dev_version_: t.Optional[str]
@@ -2226,14 +2219,15 @@ def apply_auto_restatements(
 
             # auto-restated snapshot is its own trigger
             upstream_triggers = [s_id]
+        else:
+            # inherit each parent's auto-restatement triggers (if any)
+            for parent_s_id in snapshot.parents:
+                if parent_s_id in auto_restatement_triggers:
+                    upstream_triggers.extend(auto_restatement_triggers[parent_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
+        # remove duplicate triggers, retaining order and keeping first seen of duplicates
         if upstream_triggers:
-            auto_restatement_triggers[s_id] = list(dict.fromkeys(upstream_triggers))
+            auto_restatement_triggers[s_id] = unique(upstream_triggers)
 
         if auto_restated_intervals:
             auto_restated_interval_start = sys.maxsize