Fix: Changes to before_all/after_all alone now trigger virtual updates (#4788)

Author: themisvaltinos

sqlmesh/core/plan/builder.py

@@ -835,6 +835,7 @@ def _ensure_new_env_with_changes(self) -> None:
             and not self._include_unmodified
             and self._context_diff.is_new_environment
             and not self._context_diff.has_snapshot_changes
+            and not self._context_diff.has_environment_statements_changes
             and not self._backfill_models
         ):
             raise NoChangesPlanError(

tests/core/test_context.py

@@ -52,7 +52,13 @@
     to_timestamp,
     yesterday_ds,
 )
-from sqlmesh.utils.errors import ConfigError, SQLMeshError, LinterError, PlanError
+from sqlmesh.utils.errors import (
+    ConfigError,
+    SQLMeshError,
+    LinterError,
+    PlanError,
+    NoChangesPlanError,
+)
 from sqlmesh.utils.metaprogramming import Executable
 from tests.utils.test_helpers import use_terminal_console
 from tests.utils.test_filesystem import create_temp_file
@@ -2218,3 +2224,83 @@ def test_plan_explain_skips_tests(sushi_context: Context, mocker: MockerFixture)
     spy = mocker.spy(sushi_context, "_run_plan_tests")
     sushi_context.plan(environment="dev", explain=True, no_prompts=True, include_unmodified=True)
     spy.assert_called_once_with(skip_tests=True)
+
+
+def test_dev_environment_virtual_update_with_environment_statements(tmp_path: Path) -> None:
+    models_dir = tmp_path / "models"
+    models_dir.mkdir()
+    model_sql = """
+    MODEL (
+        name db.test_model,
+        kind FULL
+    );
+
+    SELECT 1 as id, 'test' as name
+    """
+
+    with open(models_dir / "test_model.sql", "w") as f:
+        f.write(model_sql)
+
+    # Create initial context without environment statements
+    config = Config(
+        model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+        gateways={"duckdb": GatewayConfig(connection=DuckDBConnectionConfig())},
+    )
+
+    context = Context(paths=tmp_path, config=config)
+
+    # First, apply to production
+    context.plan("prod", auto_apply=True, no_prompts=True)
+
+    # Try to create dev environment without changes (should fail)
+    with pytest.raises(NoChangesPlanError, match="Creating a new environment requires a change"):
+        context.plan("dev", auto_apply=True, no_prompts=True)
+
+    # Now create a new context with only new environment statements
+    config_with_statements = Config(
+        model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+        gateways={"duckdb": GatewayConfig(connection=DuckDBConnectionConfig())},
+        before_all=["CREATE TABLE IF NOT EXISTS audit_log (id INT, action VARCHAR(100))"],
+        after_all=["INSERT INTO audit_log VALUES (1, 'environment_created')"],
+    )
+
+    context_with_statements = Context(paths=tmp_path, config=config_with_statements)
+
+    # This should succeed because environment statements are different
+    context_with_statements.plan("dev", auto_apply=True, no_prompts=True)
+    env = context_with_statements.state_reader.get_environment("dev")
+    assert env is not None
+    assert env.name == "dev"
+
+    # Verify the environment statements were stored
+    stored_statements = context_with_statements.state_reader.get_environment_statements("dev")
+    assert len(stored_statements) == 1
+    assert stored_statements[0].before_all == [
+        "CREATE TABLE IF NOT EXISTS audit_log (id INT, action VARCHAR(100))"
+    ]
+    assert stored_statements[0].after_all == [
+        "INSERT INTO audit_log VALUES (1, 'environment_created')"
+    ]
+
+    # Update environment statements and plan again (should trigger another virtual update)
+    config_updated_statements = Config(
+        model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+        gateways={"duckdb": GatewayConfig(connection=DuckDBConnectionConfig())},
+        before_all=[
+            "CREATE TABLE IF NOT EXISTS audit_log (id INT, action VARCHAR(100))",
+            "CREATE TABLE IF NOT EXISTS metrics (metric_name VARCHAR(50), value INT)",
+        ],
+        after_all=["INSERT INTO audit_log VALUES (1, 'environment_created')"],
+    )
+
+    context_updated = Context(paths=tmp_path, config=config_updated_statements)
+    context_updated.plan("dev", auto_apply=True, no_prompts=True)
+
+    # Verify the updated statements were stored
+    updated_statements = context_updated.state_reader.get_environment_statements("dev")
+    assert len(updated_statements) == 1
+    assert len(updated_statements[0].before_all) == 2
+    assert (
+        updated_statements[0].before_all[1]
+        == "CREATE TABLE IF NOT EXISTS metrics (metric_name VARCHAR(50), value INT)"
+    )

tests/core/test_plan.py

@@ -41,7 +41,7 @@
     to_timestamp,
     yesterday_ds,
 )
-from sqlmesh.utils.errors import PlanError
+from sqlmesh.utils.errors import PlanError, NoChangesPlanError
 from sqlmesh.utils.rich import strip_ansi_codes
 
 
@@ -3204,3 +3204,88 @@ def _build_plan() -> Plan:
         assert to_datetime(plan.start) == to_datetime(output_start)
         assert to_datetime(plan.end) == to_datetime(output_end)
         assert to_datetime(plan.execution_time) == to_datetime(output_execution_time)
+
+
+def test_environment_statements_change_allows_dev_environment_creation(make_snapshot):
+    snapshot = make_snapshot(
+        SqlModel(
+            name="test_model",
+            dialect="duckdb",
+            query=parse_one("select 1, ds"),
+            kind=dict(name=ModelKindName.INCREMENTAL_BY_TIME_RANGE, time_column="ds"),
+        )
+    )
+
+    # First context diff of a new 'dev' environment without environment statements
+    context_diff_no_statements = ContextDiff(
+        environment="dev",
+        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.snapshot_id: snapshot},
+        new_snapshots={},
+        previous_plan_id=None,
+        previously_promoted_snapshot_ids={snapshot.snapshot_id},
+        previous_finalized_snapshots=None,
+        previous_gateway_managed_virtual_layer=False,
+        gateway_managed_virtual_layer=False,
+        environment_statements=[],
+        previous_environment_statements=[],
+    )
+
+    # Should fail because no changes
+    plan_builder = PlanBuilder(
+        context_diff_no_statements,
+        is_dev=True,
+    )
+
+    with pytest.raises(NoChangesPlanError, match="Creating a new environment requires a change"):
+        plan_builder.build()
+
+    # Now create context diff with environment statements
+    environment_statements = [
+        EnvironmentStatements(
+            before_all=["CREATE TABLE IF NOT EXISTS test_table (id INT)"],
+            after_all=[],
+            python_env={},
+            jinja_macros=None,
+        )
+    ]
+
+    context_diff_with_statements = ContextDiff(
+        environment="dev",
+        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.snapshot_id: snapshot},
+        new_snapshots={},
+        previous_plan_id=None,
+        previously_promoted_snapshot_ids={snapshot.snapshot_id},
+        previous_finalized_snapshots=None,
+        previous_gateway_managed_virtual_layer=False,
+        gateway_managed_virtual_layer=False,
+        environment_statements=environment_statements,
+        previous_environment_statements=[],
+    )
+
+    # Should succeed because there are environment statements changes
+    plan_builder_with_statements = PlanBuilder(
+        context_diff_with_statements,
+        is_dev=True,
+    )
+
+    # Test that allows creating a dev environment without other changes
+    plan = plan_builder_with_statements.build()
+    assert plan is not None
+    assert plan.context_diff.has_environment_statements_changes
+    assert plan.context_diff.environment_statements == environment_statements