Fix: Gracefully handle execution errors in before after all (#4333)

Author: themisvaltinos

sqlmesh/core/environment.py

@@ -14,6 +14,7 @@
 from sqlmesh.core.snapshot import SnapshotId, SnapshotTableInfo, Snapshot
 from sqlmesh.utils import word_characters_only
 from sqlmesh.utils.date import TimeLike, now_timestamp
+from sqlmesh.utils.errors import SQLMeshError
 from sqlmesh.utils.jinja import JinjaMacroRegistry
 from sqlmesh.utils.metaprogramming import Executable
 from sqlmesh.utils.pydantic import PydanticModel, field_validator, ValidationInfo
@@ -261,25 +262,35 @@ def execute_environment_statements(
     end: t.Optional[TimeLike] = None,
     execution_time: t.Optional[TimeLike] = None,
 ) -> None:
-    rendered_expressions = [
-        expr
-        for statements in environment_statements
-        for expr in render_statements(
-            statements=getattr(statements, runtime_stage.value),
-            dialect=adapter.dialect,
-            default_catalog=default_catalog,
-            python_env=statements.python_env,
-            jinja_macros=statements.jinja_macros,
-            snapshots=snapshots,
-            start=start,
-            end=end,
-            execution_time=execution_time,
-            environment_naming_info=environment_naming_info,
-            runtime_stage=runtime_stage,
-            engine_adapter=adapter,
+    try:
+        rendered_expressions = [
+            expr
+            for statements in environment_statements
+            for expr in render_statements(
+                statements=getattr(statements, runtime_stage.value),
+                dialect=adapter.dialect,
+                default_catalog=default_catalog,
+                python_env=statements.python_env,
+                jinja_macros=statements.jinja_macros,
+                snapshots=snapshots,
+                start=start,
+                end=end,
+                execution_time=execution_time,
+                environment_naming_info=environment_naming_info,
+                runtime_stage=runtime_stage,
+                engine_adapter=adapter,
+            )
+        ]
+    except Exception as e:
+        raise SQLMeshError(
+            f"An error occurred during rendering of the '{runtime_stage.value}' statements:\n\n{e}"
         )
-    ]
     if rendered_expressions:
         with adapter.transaction():
             for expr in rendered_expressions:
-                adapter.execute(expr)
+                try:
+                    adapter.execute(expr)
+                except Exception as e:
+                    raise SQLMeshError(
+                        f"An error occurred during execution of the following '{runtime_stage.value}' statement:\n\n{expr}\n\n{e}"
+                    )

tests/core/test_integration.py

@@ -18,6 +18,7 @@
 from pytest_mock.plugin import MockerFixture
 from sqlglot import exp
 from sqlglot.expressions import DataType
+import re
 
 from sqlmesh import CustomMaterialization
 from sqlmesh.cli.example_project import init_example_project
@@ -5501,6 +5502,69 @@ def test_plan_production_environment_statements(tmp_path: Path):
         ctx.fetchdf("select * from not_create")
 
 
+def test_environment_statements_error_handling(tmp_path: Path):
+    model_a = """
+    MODEL (
+        name test_schema.a,
+        kind FULL,
+    );
+
+    SELECT 1 AS account_id
+    """
+
+    models_dir = tmp_path / "models"
+    models_dir.mkdir()
+
+    for path, defn in {"a.sql": model_a}.items():
+        with open(models_dir / path, "w") as f:
+            f.write(defn)
+
+    before_all = [
+        "CREATE TABLE identical_table (physical_schema_name VARCHAR)",
+        "CREATE TABLE identical_table (physical_schema_name VARCHAR)",
+    ]
+
+    config = Config(
+        model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+        before_all=before_all,
+    )
+    ctx = Context(paths=[tmp_path], config=config)
+
+    expected_error_message = re.escape(
+        """An error occurred during execution of the following 'before_all' statement:
+
+CREATE TABLE identical_table (physical_schema_name TEXT)
+
+Catalog Error: Table with name "identical_table" already exists!"""
+    )
+
+    with pytest.raises(SQLMeshError, match=expected_error_message):
+        ctx.plan(auto_apply=True, no_prompts=True)
+
+    after_all = [
+        "@bad_macro()",
+    ]
+
+    config = Config(
+        model_defaults=ModelDefaultsConfig(dialect="duckdb"),
+        after_all=after_all,
+    )
+    ctx = Context(paths=[tmp_path], config=config)
+
+    expected_error_message = re.escape(
+        """An error occurred during rendering of the 'after_all' statements:
+
+Failed to resolve macros for
+
+@bad_macro()
+
+Macro 'bad_macro' does not exist."""
+    )
+
+    with pytest.raises(SQLMeshError, match=expected_error_message):
+        ctx.plan(auto_apply=True, no_prompts=True)
+
+
 @time_machine.travel("2025-03-08 00:00:00 UTC")
 def test_tz(init_and_plan_context):
     context, _ = init_and_plan_context("examples/sushi")

tests/integrations/jupyter/test_magics.py

@@ -325,8 +325,8 @@ def test_run_dag(
     )
     assert not output.stderr
 
-    # At least 6 outputs expected as the number of models in the particular batch might vary
-    assert len(output.outputs) >= 6
+    # At least 4 outputs expected as the number of models in the particular batch might vary
+    assert len(output.outputs) >= 4
 
     html_text_actual = convert_all_html_output_to_text(output)