Refactor logic into a built-in signal

Author: VaggelisD

.circleci/continue_config.yml

@@ -310,10 +310,10 @@ workflows:
                 - athena
                 - fabric
                 - gcp-postgres
-          filters:
-            branches:
-              only:
-                - main
+          # filters:
+          #   branches:
+          #     only:
+          #       - main
       - ui_style
       - ui_test
       - vscode_test

sqlmesh/core/scheduler.py

@@ -267,14 +267,6 @@ def evaluate(
 
         snapshots = parent_snapshots_by_name(snapshot, self.snapshots)
 
-        if not is_restatement_plan and self.can_skip_evaluation(snapshot, snapshots):
-            logger.info(f"""
-            Skipping evaluation for snapshot {snapshot.name} as it depends on external models 
-            that have not been updated since the last run.
-            """)
-
-            return []
-
         is_deployable = deployability_index.is_deployable(snapshot)
 
         wap_id = self.snapshot_evaluator.evaluate(
@@ -388,6 +380,7 @@ def batch_intervals(
         deployability_index: t.Optional[DeployabilityIndex],
         environment_naming_info: EnvironmentNamingInfo,
         dag: t.Optional[DAG[SnapshotId]] = None,
+        is_restatement_plan: bool = False,
     ) -> t.Dict[Snapshot, Intervals]:
         dag = dag or snapshots_to_dag(merged_intervals)
 
@@ -427,6 +420,7 @@ def batch_intervals(
                 intervals,
                 context,
                 environment_naming_info,
+                is_restatement_plan=is_restatement_plan,
             )
             unready -= set(intervals)
 
@@ -509,9 +503,12 @@ def run_merged_intervals(
         snapshot_dag = full_dag.subdag(*selected_snapshot_ids_set)
 
         batched_intervals = self.batch_intervals(
-            merged_intervals, deployability_index, environment_naming_info, dag=snapshot_dag
+            merged_intervals,
+            deployability_index,
+            environment_naming_info,
+            dag=snapshot_dag,
+            is_restatement_plan=is_restatement_plan,
         )
-
         self.console.start_evaluation_progress(
             batched_intervals,
             environment_naming_info,
@@ -968,6 +965,7 @@ def _check_ready_intervals(
         intervals: Intervals,
         context: ExecutionContext,
         environment_naming_info: EnvironmentNamingInfo,
+        is_restatement_plan: bool = False,
     ) -> Intervals:
         """Checks if the intervals are ready for evaluation for the given snapshot.
 
@@ -989,13 +987,27 @@ def _check_ready_intervals(
         if not (signals and signals.signals_to_kwargs):
             return intervals
 
+        signal_names = signals.signals_to_kwargs.keys()
+
+        if (
+            is_restatement_plan
+            and len(signal_names) == 1
+            and next(iter(signal_names)) == "freshness"
+        ):
+            # Freshness signal is not checked for restatement plans to allow users
+            # for an escape hatch in reevaluating models
+            return intervals
+
         self.console.start_signal_progress(
             snapshot,
             self.default_catalog,
             environment_naming_info or EnvironmentNamingInfo(),
         )
 
         for signal_idx, (signal_name, kwargs) in enumerate(signals.signals_to_kwargs.items()):
+            if is_restatement_plan and signal_name == "freshness":
+                continue
+
             # Capture intervals before signal check for display
             intervals_to_check = merge_intervals(intervals)
 
@@ -1009,6 +1021,7 @@ def _check_ready_intervals(
                     python_env=signals.python_env,
                     dialect=snapshot.model.dialect,
                     path=snapshot.model._path,
+                    snapshot=snapshot,
                     kwargs=kwargs,
                 )
             except SQLMeshError as e:

sqlmesh/core/signal.py

@@ -1,8 +1,13 @@
 from __future__ import annotations
 
-
+import typing as t
 from sqlmesh.utils import UniqueKeyDict, registry_decorator
 
+if t.TYPE_CHECKING:
+    from sqlmesh.core.context import ExecutionContext
+    from sqlmesh.core.snapshot.definition import Snapshot
+    from sqlmesh.utils.date import DatetimeRanges
+
 
 class signal(registry_decorator):
     """Specifies a function which intervals are ready from a list of scheduled intervals.
@@ -33,3 +38,30 @@ class signal(registry_decorator):
 
 
 SignalRegistry = UniqueKeyDict[str, signal]
+
+
+@signal()
+def freshness(batch: DatetimeRanges, snapshot: Snapshot, context: ExecutionContext) -> bool:
+    adapter = context.engine_adapter
+    if not snapshot.last_altered_ts or not adapter.SUPPORTS_EXTERNAL_MODEL_FRESHNESS:
+        return True
+
+    adapter = context.engine_adapter
+    parent_snapshots = {context.snapshots[p.name] for p in snapshot.parents}
+    if len(parent_snapshots) != len(snapshot.node.depends_on) or not all(
+        p.is_external for p in parent_snapshots
+    ):
+        # The mismatch can happen if e.g an external model is not registered in the project
+        return True
+
+    # Finding new data means that the upstream depedencies have been altered
+    # since the last time the model was evaluated
+    upstream_dep_has_new_data = any(
+        upstream_last_altered_ts > snapshot.last_altered_ts
+        for upstream_last_altered_ts in adapter.get_external_model_freshness(
+            [p.name for p in parent_snapshots]
+        )
+    )
+
+    # Returning true is a no-op, returning False nullifies the batch so the model will not be evaluated.
+    return upstream_dep_has_new_data

sqlmesh/core/snapshot/definition.py

@@ -756,7 +756,9 @@ def hydrate_with_intervals_by_version(
             for interval in snapshot_intervals:
                 snapshot.merge_intervals(interval)
 
-                if interval.last_altered_ts:
+                # Differentiate last_altered_ts between snapshots with shared version but
+                # different dev versions e.g prod vs FORWARD_ONLY dev
+                if snapshot.dev_version == interval.dev_version and interval.last_altered_ts:
                     snapshot.last_altered_ts = max(
                         snapshot.last_altered_ts or -1, interval.last_altered_ts
                     )
@@ -1091,6 +1093,7 @@ def check_ready_intervals(
                     python_env=signals.python_env,
                     dialect=self.model.dialect,
                     path=self.model._path,
+                    snapshot=self,
                     kwargs=kwargs,
                 )
             except SQLMeshError as e:
@@ -2431,6 +2434,7 @@ def check_ready_intervals(
     python_env: t.Dict[str, Executable],
     dialect: DialectType = None,
     path: t.Optional[Path] = None,
+    snapshot: t.Optional[Snapshot] = None,
     kwargs: t.Optional[t.Dict] = None,
 ) -> Intervals:
     checked_intervals: Intervals = []
@@ -2446,6 +2450,7 @@ def check_ready_intervals(
                 provided_args=(batch,),
                 provided_kwargs=(kwargs or {}),
                 context=context,
+                snapshot=snapshot,
             )
         except Exception as ex:
             raise SignalEvalError(format_evaluated_code_exception(ex, python_env))

…s/v0093_add_last_altered_to_intervals.py …s/v0098_add_last_altered_to_intervals.pysqlmesh/migrations/v0093_add_last_altered_to_intervals.py renamed to sqlmesh/migrations/v0098_add_last_altered_to_intervals.py sqlmesh/migrations/v0093_add_last_altered_to_intervals.py renamed to sqlmesh/migrations/v0098_add_last_altered_to_intervals.py

@@ -3,7 +3,7 @@
 from sqlglot import exp
 
 
-def migrate(state_sync, **kwargs):  # type: ignore
+def migrate_schemas(state_sync, **kwargs):  # type: ignore
     engine_adapter = state_sync.engine_adapter
     schema = state_sync.schema
     intervals_table = "_intervals"

tests/core/engine_adapter/integration/test_integration.py

@@ -10,6 +10,8 @@
 from unittest import mock
 from unittest.mock import patch
 import logging
+from IPython.utils.capture import capture_output
+
 
 import time_machine
 from pytest_mock.plugin import MockerFixture
@@ -3846,36 +3848,45 @@ def test_external_model_freshness(ctx: TestContext, mocker: MockerFixture, tmp_p
     if not adapter.SUPPORTS_EXTERNAL_MODEL_FRESHNESS:
         pytest.skip("This test only runs for engines that support external model freshness")
 
-    def _run_plan(
-        sqlmesh_context: Context, restate_models: t.Optional[t.List[str]] = None
-    ) -> PlanResults:
-        plan: Plan = sqlmesh_context.plan(
-            auto_apply=True, no_prompts=True, restate_models=restate_models
+    def _assert_snapshot_last_altered_ts(context: Context, snapshot_id: str, timestamp: datetime):
+        from sqlmesh.utils.date import to_datetime
+
+        snapshot = context.state_sync.get_snapshots([snapshot_id])[snapshot_id]
+        assert to_datetime(snapshot.last_altered_ts).replace(microsecond=0) == timestamp.replace(
+            microsecond=0
         )
-        return PlanResults.create(plan, ctx, schema)
 
     import sqlmesh
 
     spy = mocker.spy(sqlmesh.core.snapshot.evaluator.SnapshotEvaluator, "evaluate")
 
     def _assert_model_evaluation(lambda_func, was_evaluated, day_delta=0):
-        call_count_before = spy.call_count
-        logger = logging.getLogger("sqlmesh.core.scheduler")
-
-        with time_machine.travel(now(minute_floor=False) + timedelta(days=day_delta)):
-            with mock.patch.object(logger, "info") as mock_logger:
-                lambda_func()
-
-        evaluation_skipped_log = any(
-            "Skipping evaluation for snapshot" in call[0][0] for call in mock_logger.call_args_list
-        )
-
+        spy.reset_mock()
+        timestamp = now(minute_floor=False) + timedelta(days=day_delta)
+        with time_machine.travel(timestamp, tick=False):
+            with capture_output() as output:
+                plan_or_run_result = lambda_func()
+
+        evaluate_function_called = spy.call_count == 1
+        signal_was_checked = "Checking signals for" in output.stdout
+        restatement_plan = isinstance(plan_or_run_result, Plan) and plan_or_run_result.restatements
+        if restatement_plan:
+            # Restatement plans exclude this signal so we expect the actual evaluation
+            # to happen but not through the signal
+            assert evaluate_function_called
+            assert not signal_was_checked
+            return
+
+        # All other cases (e.g normal plans or runs) will check the freshness signal
+        assert signal_was_checked
         if was_evaluated:
-            assert not evaluation_skipped_log
-            assert spy.call_count == call_count_before + 1
+            assert "All ready" in output.stdout
+            assert evaluate_function_called
         else:
-            assert evaluation_skipped_log
-            assert spy.call_count == call_count_before
+            assert "None ready" in output.stdout
+            assert not evaluate_function_called
+
+        return timestamp, plan_or_run_result
 
     # Create & initialize schema
     schema = ctx.add_test_suffix(TEST_SCHEMA)
@@ -3912,7 +3923,10 @@ def _assert_model_evaluation(lambda_func, was_evaluated, day_delta=0):
         MODEL (
             name {model_name},
             start '2024-01-01',
-            kind FULL
+            kind FULL,
+            signals (
+                freshness(),
+            )
         );
 
          SELECT col1 * col2 AS col FROM {external_table1}, {external_table2};
@@ -3924,23 +3938,47 @@ def _set_config(gateway: str, config: Config) -> None:
 
     context = ctx.create_context(path=tmp_path, config_mutator=_set_config)
 
-    # Case 1: Model is evaluated on first insertion
-    _assert_model_evaluation(lambda: _run_plan(context), was_evaluated=True)
+    # Case 1: Model is evaluated for the first plan
+    prod_plan_ts, prod_plan = _assert_model_evaluation(
+        lambda: context.plan(auto_apply=True, no_prompts=True), was_evaluated=True
+    )
+
+    prod_snapshot_id = next(iter(prod_plan.context_diff.new_snapshots))
+    _assert_snapshot_last_altered_ts(context, prod_snapshot_id, prod_plan_ts)
 
     # Case 2: Model is NOT evaluated on run if external models are not fresh
-    _assert_model_evaluation(lambda: context.run(), was_evaluated=False, day_delta=2)
+    _assert_model_evaluation(lambda: context.run(), was_evaluated=False, day_delta=1)
 
-    # Case 3: Model is evaluated on run if any external model is fresh
-    adapter.execute(f"INSERT INTO {external_table2} (col2) VALUES (3)", quote_identifiers=False)
+    # Case 3: Differentiate last_altered_ts between snapshots with shared version
+    # For instance, creating a FORWARD_ONLY change in dev (reusing the version but creating a dev preview) should not cause
+    # the prod snapshot's last_altered_ts to be updated when fetched from the state sync
+    model_path.write_text(model_path.read_text().replace("col1 * col2", "col1 + col2"))
+    context.load()
+    dev_plan_ts = now(minute_floor=False) + timedelta(days=2)
+    with time_machine.travel(dev_plan_ts, tick=False):
+        dev_plan = context.plan(
+            environment="dev", forward_only=True, auto_apply=True, no_prompts=True
+        )
+
+    context.state_sync.clear_cache()
+    dev_snapshot_id = next(iter(dev_plan.context_diff.new_snapshots))
+    _assert_snapshot_last_altered_ts(context, dev_snapshot_id, dev_plan_ts)
+    _assert_snapshot_last_altered_ts(context, prod_snapshot_id, prod_plan_ts)
 
+    # Case 4: Model is evaluated on run if any external model is fresh
+    adapter.execute(f"INSERT INTO {external_table2} (col2) VALUES (3)", quote_identifiers=False)
     _assert_model_evaluation(lambda: context.run(), was_evaluated=True, day_delta=2)
 
-    # Case 4: Model is evaluated on a restatement plan even if the external model is not fresh
+    # Case 5: Model is evaluated if changed (case 3) even if the external model is not fresh
+    model_path.write_text(model_path.read_text().replace("col1 + col2", "col1 * col2 * 5"))
+    context.load()
     _assert_model_evaluation(
-        lambda: _run_plan(context, restate_models=[model_name]), was_evaluated=True, day_delta=3
+        lambda: context.plan(auto_apply=True, no_prompts=True), was_evaluated=True, day_delta=3
     )
 
-    # Case 5: Model is evaluated if changed even if the external model is not fresh
-    model_path.write_text(model_path.read_text().replace("col1 * col2", "col1 + col2"))
-    context.load()
-    _assert_model_evaluation(lambda: _run_plan(context), was_evaluated=True, day_delta=2)
+    # Case 6: Model is evaluated on a restatement plan even if the external model is not fresh
+    _assert_model_evaluation(
+        lambda: context.plan(restate_models=[model_name], auto_apply=True, no_prompts=True),
+        was_evaluated=True,
+        day_delta=4,
+    )