PR feedback

Author: erindru

sqlmesh/core/context.py

@@ -1656,7 +1656,7 @@ def plan_builder(
 
         # When handling prod restatements, only clear intervals from other model versions if we are using full virtual environments
         # If we are not, then there is no point, because none of the data in dev environments can be promoted by definition
-        clear_restated_intervals_across_model_versions = (
+        restate_all_snapshots = (
             expanded_restate_models is not None
             and not is_dev
             and self.config.virtual_environment_mode.is_full
@@ -1669,7 +1669,7 @@ def plan_builder(
             execution_time=execution_time,
             apply=self.apply,
             restate_models=expanded_restate_models,
-            clear_restated_intervals_across_model_versions=clear_restated_intervals_across_model_versions,
+            restate_all_snapshots=restate_all_snapshots,
             backfill_models=backfill_models,
             no_gaps=no_gaps,
             skip_backfill=skip_backfill,
@@ -1703,7 +1703,6 @@ def plan_builder(
             },
             explain=explain or False,
             ignore_cron=ignore_cron or False,
-            always_include_local_changes=always_include_local_changes,
         )
 
     def apply(

sqlmesh/core/plan/builder.py

@@ -65,11 +65,9 @@ class PlanBuilder:
         restate_models: A list of models for which the data should be restated for the time range
             specified in this plan. Note: models defined outside SQLMesh (external) won't be a part
             of the restatement.
-        clear_restated_intervals_across_model_versions: If restatements are present, this flag indicates whether or not the intervals
+        restate_all_snapshots: If restatements are present, this flag indicates whether or not the intervals
             being restated should be cleared from state for other versions of this model (typically, versions that are present in other environments).
             If set to None, the default behaviour is to not clear anything unless the target environment is prod.
-        always_include_local_changes: Usually when restatements are present, local changes in the filesystem are ignored.
-            However, it can be desirable to deploy changes + restatements in the same plan, so this flag overrides the default behaviour.
         backfill_models: A list of fully qualified model names for which the data should be backfilled as part of this plan.
         no_gaps:  Whether to ensure that new snapshots for nodes that are already a
             part of the target environment have no data gaps when compared against previous
@@ -108,8 +106,7 @@ def __init__(
         execution_time: t.Optional[TimeLike] = None,
         apply: t.Optional[t.Callable[[Plan], None]] = None,
         restate_models: t.Optional[t.Iterable[str]] = None,
-        clear_restated_intervals_across_model_versions: bool = False,
-        always_include_local_changes: t.Optional[bool] = None,
+        restate_all_snapshots: bool = False,
         backfill_models: t.Optional[t.Iterable[str]] = None,
         no_gaps: bool = False,
         skip_backfill: bool = False,
@@ -161,9 +158,7 @@ def __init__(
         self._auto_categorization_enabled = auto_categorization_enabled
         self._include_unmodified = include_unmodified
         self._restate_models = set(restate_models) if restate_models is not None else None
-        self._clear_restated_intervals_across_model_versions = (
-            clear_restated_intervals_across_model_versions
-        )
+        self._restate_all_snapshots = restate_all_snapshots
         self._effective_from = effective_from
 
         # note: this deliberately doesnt default to now() here.
@@ -182,7 +177,6 @@ def __init__(
         self._user_provided_flags = user_provided_flags
         self._selected_models = selected_models
         self._explain = explain
-        self._always_include_local_changes = always_include_local_changes
 
         self._start = start
         if not self._start and (
@@ -288,7 +282,6 @@ def build(self) -> Plan:
         if self._latest_plan:
             return self._latest_plan
 
-        self._ensure_no_new_snapshots_with_restatements()
         self._ensure_new_env_with_changes()
         self._ensure_valid_date_range()
         self._ensure_no_broken_references()
@@ -351,7 +344,7 @@ def build(self) -> Plan:
             deployability_index=deployability_index,
             selected_models_to_restate=self._restate_models,
             restatements=restatements,
-            clear_restated_intervals_across_model_versions=self._clear_restated_intervals_across_model_versions,
+            restate_all_snapshots=self._restate_all_snapshots,
             start_override_per_model=self._start_override_per_model,
             end_override_per_model=end_override_per_model,
             selected_models_to_backfill=self._backfill_models,
@@ -871,21 +864,6 @@ def _ensure_no_broken_references(self) -> None:
                     f"""Removed {broken_references_msg} are referenced in '{snapshot.name}'. Please remove broken references before proceeding."""
                 )
 
-    def _ensure_no_new_snapshots_with_restatements(self) -> None:
-        if self._always_include_local_changes:
-            # the sqlmesh_dbt cli sets "always include local changes" to deliberately allow changes and restatements
-            # to be deployed in the same plan. If this is set, "force_no_diff" is also turned off on the ContextDiff
-            # so that the user is shown the local changes that will be applied and must accept them in order to run the plan
-            return
-
-        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."
-            )
-
     def _ensure_new_env_with_changes(self) -> None:
         if (
             self._is_dev

sqlmesh/core/plan/definition.py

@@ -66,8 +66,8 @@ class Plan(PydanticModel, frozen=True):
     
     Note that dev previews are also considered restatements, so :selected_models_to_restate can be empty
     while :restatements is still populated with dev previews
-    """    
-    clear_restated_intervals_across_model_versions: bool
+    """
+    restate_all_snapshots: bool
     """Whether or not to clear intervals from state for other versions of the models listed in :restatements"""
 
     start_override_per_model: t.Optional[t.Dict[str, datetime]]
@@ -270,7 +270,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()},
-            clear_restated_intervals_across_model_versions=self.clear_restated_intervals_across_model_versions,
+            restate_all_snapshots=self.restate_all_snapshots,
             is_dev=self.is_dev,
             allow_destructive_models=self.allow_destructive_models,
             allow_additive_models=self.allow_additive_models,
@@ -315,7 +315,7 @@ class EvaluatablePlan(PydanticModel):
     skip_backfill: bool
     empty_backfill: bool
     restatements: t.Dict[str, Interval]
-    clear_restated_intervals_across_model_versions: bool
+    restate_all_snapshots: bool
     is_dev: bool
     allow_destructive_models: t.Set[str]
     allow_additive_models: t.Set[str]

sqlmesh/core/plan/stages.py

@@ -13,6 +13,7 @@
     SnapshotTableInfo,
     SnapshotId,
 )
+from sqlmesh.utils.errors import PlanError
 
 
 @dataclass
@@ -452,10 +453,16 @@ def _get_after_all_stage(
     def _get_restatement_stage(
         self, plan: EvaluatablePlan, snapshots_by_name: t.Dict[str, Snapshot]
     ) -> t.Optional[RestatementStage]:
-        if plan.restatements and plan.clear_restated_intervals_across_model_versions:
-            return RestatementStage(
-                all_snapshots=snapshots_by_name,
-            )
+        if plan.restate_all_snapshots:
+            if plan.is_dev:
+                raise PlanError(
+                    "Clearing intervals from state across dev model versions is only valid for prod plans"
+                )
+
+            if plan.restatements:
+                return RestatementStage(
+                    all_snapshots=snapshots_by_name,
+                )
 
         return None
 

tests/core/test_plan.py

@@ -826,7 +826,7 @@ def test_missing_intervals_lookback(make_snapshot, mocker: MockerFixture):
         indirectly_modified={},
         deployability_index=DeployabilityIndex.all_deployable(),
         restatements={},
-        clear_restated_intervals_across_model_versions=False,
+        restate_all_snapshots=False,
         end_bounded=False,
         ensure_finalized_snapshots=False,
         start_override_per_model=None,
@@ -1075,68 +1075,6 @@ def test_restate_missing_model(make_snapshot, mocker: MockerFixture):
         PlanBuilder(context_diff, restate_models=["missing"]).build()
 
 
-def test_new_snapshots_with_restatements(make_snapshot, mocker: MockerFixture):
-    snapshot_a = make_snapshot(SqlModel(name="a", query=parse_one("select 1, ds")))
-
-    context_diff = ContextDiff(
-        environment="test_environment",
-        is_new_environment=True,
-        is_unfinalized_environment=False,
-        normalize_environment_name=True,
-        create_from="prod",
-        create_from_env_exists=True,
-        added=set(),
-        removed_snapshots={},
-        modified_snapshots={},
-        snapshots={snapshot_a.snapshot_id: snapshot_a},
-        new_snapshots={snapshot_a.snapshot_id: snapshot_a},
-        previous_plan_id=None,
-        previously_promoted_snapshot_ids=set(),
-        previous_finalized_snapshots=None,
-        previous_gateway_managed_virtual_layer=False,
-        gateway_managed_virtual_layer=False,
-        environment_statements=[],
-    )
-
-    with pytest.raises(
-        PlanError,
-        match=r"Model changes and restatements can't be a part of the same plan.*",
-    ):
-        PlanBuilder(context_diff, restate_models=["a"]).build()
-
-
-def test_new_snapshots_with_restatements_allowed_if_flag_set(
-    make_snapshot: t.Callable[..., Snapshot],
-):
-    snapshot_a = make_snapshot(SqlModel(name="a", query=parse_one("select 1, ds")))
-
-    context_diff = ContextDiff(
-        environment="test_environment",
-        is_new_environment=True,
-        is_unfinalized_environment=False,
-        normalize_environment_name=True,
-        create_from="prod",
-        create_from_env_exists=True,
-        added=set(),
-        removed_snapshots={},
-        modified_snapshots={},
-        snapshots={snapshot_a.snapshot_id: snapshot_a},
-        new_snapshots={snapshot_a.snapshot_id: snapshot_a},
-        previous_plan_id=None,
-        previously_promoted_snapshot_ids=set(),
-        previous_finalized_snapshots=None,
-        previous_gateway_managed_virtual_layer=False,
-        gateway_managed_virtual_layer=False,
-        environment_statements=[],
-    )
-
-    plan = PlanBuilder(
-        context_diff, restate_models=[snapshot_a.name], always_include_local_changes=True
-    ).build()
-    assert plan.restatements
-    assert plan.new_snapshots
-
-
 def test_end_validation(make_snapshot, mocker: MockerFixture):
     snapshot_a = make_snapshot(
         SqlModel(