Fix: Ignore non-restateable models and print out a warning instead of failing the plan (#2605)

Author: izeigerman

sqlmesh/core/context.py

@@ -1004,10 +1004,6 @@ def plan_builder(
         expanded_restate_models = None
         if restate_models is not None:
             expanded_restate_models = model_selector.expand_model_selections(restate_models)
-            if not expanded_restate_models:
-                self.console.log_error(
-                    f"Provided restated models do not match any models. No models will be included in plan. Provided: {', '.join(restate_models)}"
-                )
 
         snapshots = self._snapshots(models_override)
         context_diff = self._context_diff(
@@ -1083,6 +1079,7 @@ def plan_builder(
             ),
             end_bounded=not run,
             ensure_finalized_snapshots=self.config.plan.use_finalized_state,
+            console=self.console,
         )
 
     def apply(

sqlmesh/core/plan/builder.py

@@ -3,10 +3,12 @@
 import re
 import sys
 import typing as t
+import logging
 from collections import defaultdict
 from datetime import datetime
 from functools import cached_property
 
+from sqlmesh.core.console import Console, get_console
 from sqlmesh.core.config import (
     AutoCategorizationMode,
     CategorizerConfig,
@@ -28,6 +30,9 @@
 from sqlmesh.utils.errors import NoChangesPlanError, PlanError, SQLMeshError
 
 
+logger = logging.getLogger(__name__)
+
+
 class PlanBuilder:
     """Plan Builder constructs a Plan based on user choices for how they want to backfill, preview, etc. their changes.
 
@@ -88,6 +93,7 @@ def __init__(
         enable_preview: bool = False,
         end_bounded: bool = False,
         ensure_finalized_snapshots: bool = False,
+        console: t.Optional[Console] = None,
     ):
         self._context_diff = context_diff
         self._no_gaps = no_gaps
@@ -101,12 +107,13 @@ def __init__(
         self._categorizer_config = categorizer_config or CategorizerConfig()
         self._auto_categorization_enabled = auto_categorization_enabled
         self._include_unmodified = include_unmodified
-        self._restate_models = set(restate_models or [])
+        self._restate_models = set(restate_models) if restate_models is not None else None
         self._effective_from = effective_from
         self._execution_time = execution_time
         self._backfill_models = backfill_models
         self._end = end or default_end
         self._apply = apply
+        self._console = console or get_console()
 
         self._start = start
         if not self._start and self._forward_only_preview_needed:
@@ -286,9 +293,15 @@ def is_restateable_snapshot(snapshot: Snapshot) -> bool:
                 return False
             return not snapshot.is_symbolic and not snapshot.is_seed
 
-        restatements: t.Dict[SnapshotId, Interval] = {}
-        dummy_interval = (sys.maxsize, -sys.maxsize)
         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 {}
+
+        restatements: t.Dict[SnapshotId, Interval] = {}
         forward_only_preview_needed = self._forward_only_preview_needed
         if not restate_models and forward_only_preview_needed:
             # Add model names for new forward-only snapshots to the restatement list
@@ -298,18 +311,27 @@ def is_restateable_snapshot(snapshot: Snapshot) -> bool:
                 for s in self._context_diff.new_snapshots.values()
                 if s.is_materialized and (self._forward_only or s.model.forward_only)
             }
+
         if not restate_models:
-            return restatements
+            return {}
 
         # Add restate snapshots and their downstream snapshots
+        dummy_interval = (sys.maxsize, -sys.maxsize)
         for model_fqn in restate_models:
             snapshot = self._model_fqn_to_snapshot.get(model_fqn)
-            if not snapshot or (
-                not forward_only_preview_needed and not is_restateable_snapshot(snapshot)
-            ):
-                raise PlanError(
-                    f"Cannot restate from '{model_fqn}'. Either such model doesn't exist, no other materialized model references it, or restatement was disabled for this model."
-                )
+            if not snapshot:
+                raise PlanError(f"Cannot restate model '{model_fqn}'. Model does not exist.")
+            if not forward_only_preview_needed:
+                if not self._is_dev and snapshot.disable_restatement:
+                    # This is a warning but we print this as error since the Console is lacking API for warnings.
+                    self._console.log_error(
+                        f"Cannot restate model '{model_fqn}'. Restatement is disabled for this model."
+                    )
+                    continue
+                elif snapshot.is_symbolic or snapshot.is_seed:
+                    logger.info("Skipping restatement for model '%s'", model_fqn)
+                    continue
+
             restatements[snapshot.snapshot_id] = dummy_interval
             for downstream_s_id in dag.downstream(snapshot.snapshot_id):
                 if is_restateable_snapshot(self._context_diff.snapshots[downstream_s_id]):
@@ -632,7 +654,9 @@ def _ensure_no_broken_references(self) -> None:
                 )
 
     def _ensure_no_new_snapshots_with_restatements(self) -> None:
-        if self._restate_models and self._context_diff.new_snapshots:
+        if self._restate_models is not None and (
+            self._context_diff.new_snapshots or self._context_diff.modified_snapshots
+        ):
             raise PlanError(
                 "Model changes and restatements can't be a part of the same plan. "
                 "Revert or apply changes before proceeding with restatements."

tests/core/test_plan.py

@@ -357,7 +357,7 @@ def test_restate_models_with_existing_missing_intervals(sushi_context: Context):
     assert plan.requires_backfill
 
 
-def test_restate_model_with_merge_strategy(make_snapshot, mocker: MockerFixture):
+def test_restate_symbolic_model(make_snapshot, mocker: MockerFixture):
     snapshot_a = make_snapshot(
         SqlModel(
             name="a",
@@ -381,11 +381,61 @@ def test_restate_model_with_merge_strategy(make_snapshot, mocker: MockerFixture)
         previous_finalized_snapshots=None,
     )
 
+    plan = PlanBuilder(context_diff, restate_models=[snapshot_a.name]).build()
+    assert not plan.restatements
+
+
+def test_restate_seed_model(make_snapshot, mocker: MockerFixture):
+    snapshot_a = make_snapshot(
+        SeedModel(
+            name="a",
+            kind=SeedKind(path="./path/to/seed"),
+            seed=Seed(content="new_content"),
+            column_hashes={"col": "hash2"},
+            depends_on=set(),
+        )
+    )
+
+    context_diff = ContextDiff(
+        environment="test_environment",
+        is_new_environment=True,
+        is_unfinalized_environment=False,
+        create_from="prod",
+        added=set(),
+        removed_snapshots={},
+        modified_snapshots={},
+        snapshots={snapshot_a.snapshot_id: snapshot_a},
+        new_snapshots={},
+        previous_plan_id=None,
+        previously_promoted_snapshot_ids=set(),
+        previous_finalized_snapshots=None,
+    )
+
+    plan = PlanBuilder(context_diff, restate_models=[snapshot_a.name]).build()
+    assert not plan.restatements
+
+
+def test_restate_missing_model(make_snapshot, mocker: MockerFixture):
+    context_diff = ContextDiff(
+        environment="test_environment",
+        is_new_environment=True,
+        is_unfinalized_environment=False,
+        create_from="prod",
+        added=set(),
+        removed_snapshots={},
+        modified_snapshots={},
+        snapshots={},
+        new_snapshots={},
+        previous_plan_id=None,
+        previously_promoted_snapshot_ids=set(),
+        previous_finalized_snapshots=None,
+    )
+
     with pytest.raises(
         PlanError,
-        match="Cannot restate from 'a'. Either such model doesn't exist, no other materialized model references it.*",
+        match=r"Cannot restate model 'missing'. Model does not exist.",
     ):
-        PlanBuilder(context_diff, restate_models=["a"]).build()
+        PlanBuilder(context_diff, restate_models=["missing"]).build()
 
 
 def test_new_snapshots_with_restatements(make_snapshot, mocker: MockerFixture):
@@ -1111,8 +1161,8 @@ def test_disable_restatement(make_snapshot, mocker: MockerFixture):
         previous_finalized_snapshots=None,
     )
 
-    with pytest.raises(PlanError, match="""Cannot restate from '"a"'.*"""):
-        PlanBuilder(context_diff, restate_models=['"a"']).build()
+    plan = PlanBuilder(context_diff, restate_models=['"a"']).build()
+    assert not plan.restatements
 
     # Effective from doesn't apply to snapshots for which restatements are disabled.
     plan = PlanBuilder(context_diff, forward_only=True, effective_from="2023-01-01").build()