Fix: Pass runtime_stage and engine_adapter in environment_statements (#4021)

Author: themisvaltinos

sqlmesh/core/environment.py

@@ -244,6 +244,8 @@ def execute_environment_statements(
             end=end,
             execution_time=execution_time,
             environment_naming_info=environment_naming_info,
+            runtime_stage=runtime_stage,
+            engine_adapter=adapter,
         )
     ]:
         with adapter.transaction():

tests/core/test_context.py

@@ -31,7 +31,7 @@
 from sqlmesh.core.dialect import parse, schema_
 from sqlmesh.core.engine_adapter.duckdb import DuckDBEngineAdapter
 from sqlmesh.core.environment import Environment, EnvironmentNamingInfo, EnvironmentStatements
-from sqlmesh.core.macros import MacroEvaluator
+from sqlmesh.core.macros import MacroEvaluator, RuntimeStage
 from sqlmesh.core.model import load_sql_based_model, model, SqlModel, Model
 from sqlmesh.core.model.cache import OptimizedQueryCache
 from sqlmesh.core.renderer import render_statements
@@ -1444,7 +1444,7 @@ def test_environment_statements(tmp_path: pathlib.Path):
 
 @macro()
 def grant_select_privileges(evaluator):
-    if evaluator._environment_naming_info:
+    if evaluator._environment_naming_info and evaluator.runtime_stage == 'before_all':
         mapping = to_view_mapping(
             evaluator._snapshots.values(), evaluator._environment_naming_info
         )
@@ -1493,7 +1493,10 @@ def grant_schema_usage(evaluator):
     assert isinstance(python_env["grant_select_privileges"], Executable)
 
     before_all_rendered = render_statements(
-        statements=before_all, dialect=dialect, python_env=python_env
+        statements=before_all,
+        dialect=dialect,
+        python_env=python_env,
+        runtime_stage=RuntimeStage.BEFORE_ALL,
     )
 
     assert before_all_rendered == [
@@ -1506,6 +1509,7 @@ def grant_schema_usage(evaluator):
         python_env=python_env,
         snapshots=snapshots,
         environment_naming_info=EnvironmentNamingInfo(name="prod"),
+        runtime_stage=RuntimeStage.BEFORE_ALL,
     )
 
     assert after_all_rendered == [
@@ -1519,6 +1523,7 @@ def grant_schema_usage(evaluator):
         python_env=python_env,
         snapshots=snapshots,
         environment_naming_info=EnvironmentNamingInfo(name="dev"),
+        runtime_stage=RuntimeStage.BEFORE_ALL,
     )
 
     assert after_all_rendered_dev == [
@@ -1534,7 +1539,7 @@ def test_plan_environment_statements(tmp_path: pathlib.Path):
 
     config = Config(
         model_defaults=ModelDefaultsConfig(dialect=dialect),
-        before_all=["@create_stats_table()"],
+        before_all=["@create_stats_table()", "@access_adapter()"],
         after_all=["CREATE TABLE IF NOT EXISTS after_table AS SELECT @some_var"],
         variables={"some_var": 5},
     )
@@ -1578,9 +1583,34 @@ def create_stats_table(evaluator):
 """,
     )
 
+    create_temp_file(
+        tmp_path,
+        pathlib.Path(macros_dir, "access_adapter.py"),
+        """
+from sqlmesh.core.macros import macro
+
+@macro()
+def access_adapter(evaluator):
+    if evaluator.runtime_stage == 'before_all':
+        engine_adapter = evaluator.engine_adapter
+        for i in range(10):
+            try:
+                sql_inside_macro = f"CREATE TABLE IF NOT EXISTS db_connect AS SELECT {i} as 'access_attempt'"
+                engine_adapter.execute(sql_inside_macro)
+                return None
+            except Exception as e:
+                sleep(10)
+        raise Exception(f"Failed to connect to the database")
+    """,
+    )
+
     context = Context(paths=tmp_path, config=config)
 
-    assert context._environment_statements[0].before_all == ["@create_stats_table()"]
+    assert context._environment_statements[0].before_all == [
+        "@create_stats_table()",
+        "@access_adapter()",
+    ]
+
     assert context._environment_statements[0].after_all == [
         "CREATE TABLE IF NOT EXISTS after_table AS SELECT @some_var"
     ]
@@ -1619,6 +1649,11 @@ def create_stats_table(evaluator):
     assert state_table[0].after_all == context._environment_statements[0].after_all
     assert state_table[0].python_env == context._environment_statements[0].python_env
 
+    # This table will be created inside the macro by accessing the engine_adapter directly
+    inside_macro_execute = context.fetchdf("select * from memory.db_connect").to_dict()
+    assert (attempt_column := inside_macro_execute.get("access_attempt"))
+    assert isinstance(attempt_column, dict) and attempt_column[0] < 10
+
 
 def test_environment_statements_dialect(tmp_path: Path):
     before_all = [

tests/core/test_integration.py

@@ -4924,9 +4924,13 @@ def test_plan_production_environment_statements(tmp_path: Path):
             f.write(defn)
 
     before_all = [
-        "CREATE TABLE IF NOT EXISTS schema_names_for_@this_env (physical_schema_name VARCHAR)"
+        "CREATE TABLE IF NOT EXISTS schema_names_for_@this_env (physical_schema_name VARCHAR)",
+        "@IF(@runtime_stage = 'before_all', CREATE TABLE IF NOT EXISTS should_create AS SELECT @runtime_stage)",
+    ]
+    after_all = [
+        "@IF(@this_env = 'prod', CREATE TABLE IF NOT EXISTS after_t AS SELECT @var_5)",
+        "@IF(@runtime_stage = 'before_all', CREATE TABLE IF NOT EXISTS not_create AS SELECT @runtime_stage)",
     ]
-    after_all = ["@IF(@this_env = 'prod', CREATE TABLE IF NOT EXISTS after_t AS SELECT @var_5)"]
     config = Config(
         model_defaults=ModelDefaultsConfig(dialect="duckdb"),
         before_all=before_all,
@@ -4948,6 +4952,14 @@ def test_plan_production_environment_statements(tmp_path: Path):
     assert environment_statements[0].python_env.keys() == {"__sqlmesh__vars__"}
     assert environment_statements[0].python_env["__sqlmesh__vars__"].payload == "{'var_5': 5}"
 
+    should_create = ctx.fetchdf("select * from should_create").to_dict()
+    assert should_create["before_all"][0] == "before_all"
+
+    with pytest.raises(
+        Exception, match=r"Catalog Error: Table with name not_create does not exist!"
+    ):
+        ctx.fetchdf("select * from not_create")
+
 
 @time_machine.travel("2025-03-08 00:00:00 UTC")
 def test_tz(init_and_plan_context):