Fix: A restatement plan should not override environment statements (#4880)

Author: izeigerman

sqlmesh/core/context_diff.py

@@ -84,7 +84,7 @@ class ContextDiff(PydanticModel):
     """Python dependencies."""
     previous_environment_statements: t.List[EnvironmentStatements] = []
     """Previous environment statements."""
-    environment_statements: t.List[EnvironmentStatements] = []
+    environment_statements: t.List[EnvironmentStatements]
     """Environment statements."""
     diff_rendered: bool = False
     """Whether the diff should compare raw vs rendered models"""
@@ -268,6 +268,7 @@ def create_no_diff(cls, environment: str, state_reader: StateReader) -> ContextD
         if not env:
             raise SQLMeshError(f"Environment '{environment}' must exist for this operation.")
 
+        environment_statements = state_reader.get_environment_statements(environment)
         snapshots = state_reader.get_snapshots(env.snapshots)
 
         return ContextDiff(
@@ -288,6 +289,7 @@ def create_no_diff(cls, environment: str, state_reader: StateReader) -> ContextD
             previous_requirements=env.requirements,
             requirements=env.requirements,
             previous_environment_statements=[],
+            environment_statements=environment_statements,
             previous_gateway_managed_virtual_layer=env.gateway_managed,
             gateway_managed_virtual_layer=env.gateway_managed,
         )

tests/core/test_integration.py

@@ -5652,6 +5652,28 @@ def test_restatement_of_full_model_with_start(init_and_plan_context: t.Callable)
     assert waiter_by_day_interval == (to_timestamp("2023-01-07"), to_timestamp("2023-01-08"))
 
 
+@time_machine.travel("2023-01-08 15:00:00 UTC")
+def test_restatement_should_not_override_environment_statements(init_and_plan_context: t.Callable):
+    context, _ = init_and_plan_context("examples/sushi")
+    context.config.before_all = ["SELECT 'test_before_all';"]
+    context.load()
+
+    context.plan("prod", auto_apply=True, no_prompts=True, skip_tests=True)
+
+    prod_env_statements = context.state_reader.get_environment_statements(c.PROD)
+    assert prod_env_statements[0].before_all == ["SELECT 'test_before_all';"]
+
+    context.plan(
+        restate_models=["sushi.waiter_revenue_by_day"],
+        start="2023-01-07",
+        auto_apply=True,
+        no_prompts=True,
+    )
+
+    prod_env_statements = context.state_reader.get_environment_statements(c.PROD)
+    assert prod_env_statements[0].before_all == ["SELECT 'test_before_all';"]
+
+
 @time_machine.travel("2023-01-08 15:00:00 UTC")
 def test_restatement_shouldnt_backfill_beyond_prod_intervals(init_and_plan_context: t.Callable):
     context, _ = init_and_plan_context("examples/sushi")