Add directly modified and restatement triggers

Author: treysp

sqlmesh/core/console.py

@@ -3830,12 +3830,10 @@ def update_snapshot_evaluation_progress(
                 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:
-                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.directly_modified_triggers:
+                message += f" | directly_modified_triggers={','.join(trigger.name for trigger in snapshot_evaluation_triggers.directly_modified_triggers)}"
+            if snapshot_evaluation_triggers.restatement_triggers:
+                message += f" | restatement_triggers={','.join(trigger.name for trigger in snapshot_evaluation_triggers.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}"

sqlmesh/core/context.py

@@ -2298,7 +2298,7 @@ def check_intervals(
         if select_models:
             selected, _ = self._select_models_for_run(select_models, True, snapshots.values())
         else:
-            selected = t.cast(t.Set[str], snapshots.keys())
+            selected = set(snapshots.keys())
 
         results = {}
         execution_context = self.execution_context(snapshots=snapshots)

sqlmesh/core/plan/builder.py

@@ -293,7 +293,7 @@ def build(self) -> Plan:
             else DeployabilityIndex.all_deployable()
         )
 
-        restatements = self._build_restatements(
+        restatements, restatement_triggers = self._build_restatements(
             dag,
             earliest_interval_start(self._context_diff.snapshots.values(), self.execution_time),
         )
@@ -330,6 +330,7 @@ def build(self) -> Plan:
             indirectly_modified=indirectly_modified,
             deployability_index=deployability_index,
             restatements=restatements,
+            restatement_triggers=restatement_triggers,
             start_override_per_model=self._start_override_per_model,
             end_override_per_model=end_override_per_model,
             selected_models_to_backfill=self._backfill_models,
@@ -351,14 +352,14 @@ def _build_dag(self) -> DAG[SnapshotId]:
 
     def _build_restatements(
         self, dag: DAG[SnapshotId], earliest_interval_start: TimeLike
-    ) -> t.Dict[SnapshotId, Interval]:
+    ) -> t.Tuple[t.Dict[SnapshotId, Interval], t.Dict[SnapshotId, t.List[SnapshotId]]]:
         restate_models = self._restate_models
         if restate_models == set():
             # This is a warning but we print this as error since the Console is lacking API for warnings.
             self._console.log_error(
                 "Provided restated models do not match any models. No models will be included in plan."
             )
-            return {}
+            return {}, {}
 
         restatements: t.Dict[SnapshotId, Interval] = {}
         forward_only_preview_needed = self._forward_only_preview_needed
@@ -380,7 +381,7 @@ def _build_restatements(
             is_preview = True
 
         if not restate_models:
-            return {}
+            return {}, {}
 
         start = self._start or earliest_interval_start
         end = self._end or now()
@@ -390,6 +391,7 @@ def _build_restatements(
             if model_fqn not in self._model_fqn_to_snapshot:
                 raise PlanError(f"Cannot restate model '{model_fqn}'. Model does not exist.")
 
+        restatement_triggers: t.Dict[SnapshotId, t.List[SnapshotId]] = {}
         # Get restatement intervals for all restated snapshots and make sure that if an incremental snapshot expands it's
         # restatement range that it's downstream dependencies all expand their restatement ranges as well.
         for s_id in dag:
@@ -422,6 +424,13 @@ def _build_restatements(
                     logger.info("Skipping restatement for model '%s'", snapshot.name)
                     continue
 
+            if snapshot.name in restate_models:
+                restatement_triggers[s_id] = [s_id]
+            if restating_parents:
+                restatement_triggers[s_id] = restatement_triggers.get(s_id, []) + [
+                    s.snapshot_id for s in restating_parents
+                ]
+
             possible_intervals = {
                 restatements[p.snapshot_id] for p in restating_parents if p.is_incremental
             }
@@ -450,7 +459,7 @@ def _build_restatements(
 
             restatements[s_id] = (snapshot_start, snapshot_end)
 
-        return restatements
+        return restatements, restatement_triggers
 
     def _build_directly_and_indirectly_modified(
         self, dag: DAG[SnapshotId]

sqlmesh/core/plan/definition.py

@@ -57,6 +57,7 @@ class Plan(PydanticModel, frozen=True):
 
     deployability_index: DeployabilityIndex
     restatements: t.Dict[SnapshotId, Interval]
+    restatement_triggers: t.Dict[SnapshotId, t.List[SnapshotId]] = {}
     start_override_per_model: t.Optional[t.Dict[str, datetime]]
     end_override_per_model: t.Optional[t.Dict[str, datetime]]
 
@@ -254,6 +255,7 @@ def to_evaluatable(self) -> EvaluatablePlan:
             skip_backfill=self.skip_backfill,
             empty_backfill=self.empty_backfill,
             restatements={s.name: i for s, i in self.restatements.items()},
+            restatement_triggers=self.restatement_triggers,
             is_dev=self.is_dev,
             allow_destructive_models=self.allow_destructive_models,
             forward_only=self.forward_only,
@@ -295,6 +297,7 @@ class EvaluatablePlan(PydanticModel):
     skip_backfill: bool
     empty_backfill: bool
     restatements: t.Dict[str, Interval]
+    restatement_triggers: t.Dict[SnapshotId, t.List[SnapshotId]] = {}
     is_dev: bool
     allow_destructive_models: t.Set[str]
     forward_only: bool

sqlmesh/core/plan/evaluator.py

@@ -37,6 +37,7 @@
     SnapshotCreationFailedError,
     SnapshotNameVersion,
 )
+from sqlmesh.core.snapshot.definition import SnapshotEvaluationTriggers
 from sqlmesh.utils import to_snake_case
 from sqlmesh.core.state_sync import StateSync
 from sqlmesh.utils import CorrelationId
@@ -234,6 +235,27 @@ def visit_backfill_stage(self, stage: stages.BackfillStage, plan: EvaluatablePla
             self.console.log_success("SKIP: No model batches to execute")
             return
 
+        directly_modified_triggers: t.Dict[SnapshotId, t.List[SnapshotId]] = {}
+        for parent, children in plan.indirectly_modified_snapshots.items():
+            parent_id = stage.all_snapshots[parent].snapshot_id
+            directly_modified_triggers[parent_id] = directly_modified_triggers.get(
+                parent_id, []
+            ) + [parent_id]
+            for child in children:
+                directly_modified_triggers[child] = directly_modified_triggers.get(child, []) + [
+                    parent_id
+                ]
+        directly_modified_triggers = {
+            k: list(dict.fromkeys(v)) for k, v in directly_modified_triggers.items()
+        }
+        snapshot_evaluation_triggers = {
+            s_id: SnapshotEvaluationTriggers(
+                directly_modified_triggers=directly_modified_triggers.get(s_id, []),
+                restatement_triggers=plan.restatement_triggers.get(s_id, []),
+            )
+            for s_id in [s.snapshot_id for s in stage.all_snapshots.values()]
+        }
+
         scheduler = self.create_scheduler(stage.all_snapshots.values(), self.snapshot_evaluator)
         # Convert model name restatements to snapshot ID restatements
         restatements_by_snapshot_id = {
@@ -249,6 +271,7 @@ def visit_backfill_stage(self, stage: stages.BackfillStage, plan: EvaluatablePla
             start=plan.start,
             end=plan.end,
             restatements=restatements_by_snapshot_id,
+            snapshot_evaluation_triggers=snapshot_evaluation_triggers,
         )
         if errors:
             raise PlanError("Plan application failed.")

sqlmesh/core/snapshot/definition.py

@@ -330,6 +330,8 @@ class SnapshotEvaluationTriggers(PydanticModel):
     cron_ready: t.Optional[bool] = None
     auto_restatement_triggers: t.List[SnapshotId] = []
     select_snapshot_triggers: t.List[SnapshotId] = []
+    directly_modified_triggers: t.List[SnapshotId] = []
+    restatement_triggers: t.List[SnapshotId] = []
 
 
 class SnapshotInfoMixin(ModelKindMixin):

tests/core/test_integration.py

@@ -26,6 +26,7 @@
 
 
 from sqlmesh import CustomMaterialization
+import sqlmesh
 from sqlmesh.cli.project_init import init_example_project
 from sqlmesh.core import constants as c
 from sqlmesh.core import dialect as d
@@ -1784,26 +1785,97 @@ def test_snapshot_triggers(init_and_plan_context: t.Callable, mocker: MockerFixt
     context, plan = init_and_plan_context("examples/sushi")
     context.apply(plan)
 
+    # modify 3 models
+    # - 2 breaking changes for testing plan directly modified triggers
+    # - 1 adding an auto-restatement for subsequent `run` test
+    marketing = context.get_model("sushi.marketing")
+    marketing_kwargs = {
+        **marketing.dict(),
+        "query": d.parse_one(
+            f"{marketing.query.sql(dialect='duckdb')} ORDER BY customer_id", dialect="duckdb"
+        ),
+    }
+    context.upsert_model(SqlModel.parse_obj(marketing_kwargs))
+
+    customers = context.get_model("sushi.customers")
+    customers_kwargs = {
+        **customers.dict(),
+        "query": d.parse_one(
+            f"{customers.query.sql(dialect='duckdb')} ORDER BY customer_id", dialect="duckdb"
+        ),
+    }
+    context.upsert_model(SqlModel.parse_obj(customers_kwargs))
+
     # add auto restatement to orders
-    model = context.get_model("sushi.orders")
-    kind = {
-        **model.kind.dict(),
+    orders = context.get_model("sushi.orders")
+    orders_kind = {
+        **orders.kind.dict(),
         "auto_restatement_cron": "@hourly",
     }
-    kwargs = {
-        **model.dict(),
-        "kind": kind,
+    orders_kwargs = {
+        **orders.dict(),
+        "kind": orders_kind,
     }
-    context.upsert_model(PythonModel.parse_obj(kwargs))
-    plan = context.plan_builder(skip_tests=True).build()
-    context.apply(plan)
+    context.upsert_model(PythonModel.parse_obj(orders_kwargs))
 
-    # Mock run_merged_intervals to capture triggers arg
-    scheduler = context.scheduler()
-    run_merged_intervals_mock = mocker.patch.object(
-        scheduler, "run_merged_intervals", return_value=([], [])
+    spy = mocker.spy(sqlmesh.core.scheduler.Scheduler, "run_merged_intervals")
+
+    context.plan(auto_apply=True, no_prompts=True, categorizer_config=CategorizerConfig.all_full())
+
+    # PLAN: directly modified triggers
+    actual_triggers = spy.call_args.kwargs["snapshot_evaluation_triggers"]
+    actual_triggers_name = {
+        k.name: sorted([s.name for s in v.directly_modified_triggers])
+        for k, v in actual_triggers.items()
+        if v.directly_modified_triggers
+    }
+    marketing_name = '"memory"."sushi"."marketing"'
+    customers_name = '"memory"."sushi"."customers"'
+    marketing_customers_names = sorted([marketing_name, customers_name])
+    children_names = [
+        f'"memory"."sushi"."{model}"'
+        for model in {
+            "waiter_as_customer_by_day",
+            "active_customers",
+            "count_customers_active",
+            "count_customers_inactive",
+        }
+    ]
+    assert actual_triggers_name == {
+        marketing_name: [marketing_name],
+        customers_name: [customers_name],
+        **{k: marketing_customers_names for k in children_names},
+    }
+
+    # PLAN: restatement triggers
+    spy.reset_mock()
+    context.plan(
+        restate_models=[
+            '"memory"."sushi"."marketing"',
+            '"memory"."sushi"."order_items"',
+            '"memory"."sushi"."waiter_revenue_by_day"',
+        ],
+        auto_apply=True,
+        no_prompts=True,
     )
 
+    order_items_name = '"memory"."sushi"."order_items"'
+    waiter_revenue_by_day_name = '"memory"."sushi"."waiter_revenue_by_day"'
+    actual_triggers = spy.call_args.kwargs["snapshot_evaluation_triggers"]
+    actual_triggers_name = {
+        k.name: sorted([s.name for s in v.restatement_triggers])
+        for k, v in actual_triggers.items()
+        if v.restatement_triggers
+    }
+    assert actual_triggers_name == {
+        waiter_revenue_by_day_name: [waiter_revenue_by_day_name, order_items_name],
+        order_items_name: [order_items_name],
+        '"memory"."sushi"."top_waiters"': [waiter_revenue_by_day_name],
+        '"memory"."sushi"."customer_revenue_by_day"': [order_items_name],
+        '"memory"."sushi"."customer_revenue_lifetime"': [order_items_name],
+    }
+
+    # RUN: select and auto-restatement triggers
     # 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"}
@@ -1814,6 +1886,11 @@ def test_snapshot_triggers(init_and_plan_context: t.Callable, mocker: MockerFixt
         f'"memory"."sushi"."{model}"' for model in selected_models
     }
 
+    scheduler = context.scheduler()
+    run_merged_intervals_mock = mocker.patch.object(
+        scheduler, "run_merged_intervals", return_value=([], [])
+    )
+
     with time_machine.travel("2023-01-09 00:00:01 UTC"):
         scheduler.run(
             environment=c.PROD,