Fix: Snapshots promoted in prod shouldn't be restated in dev (#3843)

Author: izeigerman

sqlmesh/core/plan/builder.py

@@ -329,9 +329,14 @@ def is_restateable_snapshot(snapshot: Snapshot) -> bool:
             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 or not snapshot.is_paused) 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(
+                if self._is_dev and not snapshot.is_paused:
+                    self._console.log_warning(
+                        f"Cannot restate model '{model_fqn}' because the current version is used in production. "
+                        "Run the restatement against the production environment instead to restate this model."
+                    )
+                    continue
+                elif (not self._is_dev or not snapshot.is_paused) and snapshot.disable_restatement:
+                    self._console.log_warning(
                         f"Cannot restate model '{model_fqn}'. "
                         "Restatement is disabled for this model to prevent possible data loss."
                         "If you want to restate this model, change the model's `disable_restatement` setting to `false`."

tests/core/test_integration.py

@@ -2983,6 +2983,23 @@ def test_prod_restatement_plan_missing_model_in_dev(
     )
 
 
+@time_machine.travel("2023-01-08 15:00:00 UTC")
+def test_dev_restatement_of_prod_model(init_and_plan_context: t.Callable):
+    context, plan = init_and_plan_context("examples/sushi")
+    context.apply(plan)
+
+    model = context.get_model("sushi.waiter_revenue_by_day")
+    context.upsert_model(add_projection_to_model(t.cast(SqlModel, model)))
+
+    context.plan("dev", auto_apply=True, no_prompts=True, skip_tests=True)
+
+    restatement_plan = context.plan_builder("dev", restate_models=["*"]).build()
+    assert set(restatement_plan.restatements) == {
+        context.get_snapshot("sushi.waiter_revenue_by_day").snapshot_id,
+        context.get_snapshot("sushi.top_waiters").snapshot_id,
+    }
+
+
 @time_machine.travel("2023-01-08 15:00:00 UTC")
 def test_plan_snapshot_table_exists_for_promoted_snapshot(init_and_plan_context: t.Callable):
     context, plan = init_and_plan_context("examples/sushi")

tests/core/test_plan.py

@@ -2697,3 +2697,59 @@ def test_unaligned_start_model_with_forward_only_preview(make_snapshot):
     assert set(plan.restatements) == {new_snapshot_a.snapshot_id, snapshot_b.snapshot_id}
     assert not plan.deployability_index.is_deployable(new_snapshot_a)
     assert not plan.deployability_index.is_deployable(snapshot_b)
+
+
+def test_restate_production_model_in_dev(make_snapshot, mocker: MockerFixture):
+    snapshot = make_snapshot(
+        SqlModel(
+            name="test_model_a",
+            dialect="duckdb",
+            query=parse_one("select 1, ds"),
+            kind=dict(name=ModelKindName.INCREMENTAL_BY_TIME_RANGE, time_column="ds"),
+        )
+    )
+
+    prod_snapshot = make_snapshot(
+        SqlModel(
+            name="test_model_b",
+            dialect="duckdb",
+            query=parse_one("select 2, ds"),
+            kind=dict(name=ModelKindName.INCREMENTAL_BY_TIME_RANGE, time_column="ds"),
+        )
+    )
+    prod_snapshot.unpaused_ts = 1
+
+    context_diff = ContextDiff(
+        environment="test_environment",
+        is_new_environment=False,
+        is_unfinalized_environment=True,
+        normalize_environment_name=True,
+        create_from="prod",
+        create_from_env_exists=True,
+        added=set(),
+        removed_snapshots={},
+        modified_snapshots={},
+        snapshots={snapshot.snapshot_id: snapshot, prod_snapshot.snapshot_id: prod_snapshot},
+        new_snapshots={},
+        previous_plan_id=None,
+        previously_promoted_snapshot_ids=set(),
+        previous_finalized_snapshots=None,
+    )
+
+    mock_console = mocker.Mock()
+
+    plan = PlanBuilder(
+        context_diff,
+        DuckDBEngineAdapter.SCHEMA_DIFFER,
+        is_dev=True,
+        restate_models={snapshot.name, prod_snapshot.name},
+        console=mock_console,
+    ).build()
+
+    assert len(plan.restatements) == 1
+    assert prod_snapshot.snapshot_id not in plan.restatements
+
+    mock_console.log_warning.assert_called_once_with(
+        "Cannot restate model '\"test_model_b\"' because the current version is used in production. "
+        "Run the restatement against the production environment instead to restate this model."
+    )