Feat!: First-class support for calling python macros from jinja context (#4211)

Author: izeigerman

sqlmesh/core/macros.py

@@ -150,10 +150,9 @@ def __init__(
         self,
         dialect: DialectType = "",
         python_env: t.Optional[t.Dict[str, Executable]] = None,
-        jinja_env: t.Optional[Environment] = None,
         schema: t.Optional[MappingSchema] = None,
         runtime_stage: RuntimeStage = RuntimeStage.LOADING,
-        resolve_table: t.Optional[t.Callable[[str | exp.Expression], str]] = None,
+        resolve_table: t.Optional[t.Callable[[str | exp.Table], str]] = None,
         resolve_tables: t.Optional[t.Callable[[exp.Expression], exp.Expression]] = None,
         snapshots: t.Optional[t.Dict[str, Snapshot]] = None,
         default_catalog: t.Optional[str] = None,
@@ -177,7 +176,7 @@ def __init__(
         self.columns_to_types_called = False
         self.default_catalog = default_catalog
 
-        self._jinja_env: t.Optional[Environment] = jinja_env
+        self._jinja_env: t.Optional[Environment] = None
         self._schema = schema
         self._resolve_table = resolve_table
         self._resolve_tables = resolve_tables
@@ -431,7 +430,7 @@ def get_snapshot(self, model_name: TableName | exp.Column) -> t.Optional[Snapsho
             )
         )
 
-    def resolve_table(self, table: str | exp.Expression) -> str:
+    def resolve_table(self, table: str | exp.Table) -> str:
         """Gets the physical table name for a given model."""
         if not self._resolve_table:
             raise SQLMeshError(

sqlmesh/core/model/definition.py

@@ -2476,9 +2476,7 @@ def _create_model(
     for _, kwargs in model.signals:
         statements.extend((signal_kwarg, True) for signal_kwarg in kwargs.values())
 
-    python_env = python_env or {}
-
-    make_python_env(
+    python_env = make_python_env(
         statements,
         jinja_macro_references,
         module_path,

sqlmesh/core/renderer.py

@@ -3,6 +3,7 @@
 import logging
 import typing as t
 from contextlib import contextmanager
+from functools import partial
 from pathlib import Path
 
 from sqlglot import exp, parse
@@ -134,7 +135,35 @@ def _render(
         if this_snapshot and (kind := this_snapshot.model_kind_name):
             kwargs["model_kind_name"] = kind.name
 
-        expressions = [self._expression]
+        def _resolve_table(table: str | exp.Table) -> str:
+            return self._resolve_table(
+                d.normalize_model_name(table, self._default_catalog, self._dialect),
+                snapshots=snapshots,
+                table_mapping=table_mapping,
+                deployability_index=deployability_index,
+            ).sql(dialect=self._dialect, identify=True, comments=False)
+
+        macro_evaluator = MacroEvaluator(
+            self._dialect,
+            python_env=self._python_env,
+            schema=self.schema,
+            runtime_stage=runtime_stage,
+            resolve_table=_resolve_table,
+            resolve_tables=lambda e: self._resolve_tables(
+                e,
+                snapshots=snapshots,
+                table_mapping=table_mapping,
+                deployability_index=deployability_index,
+                start=start,
+                end=end,
+                execution_time=execution_time,
+                runtime_stage=runtime_stage,
+            ),
+            snapshots=snapshots,
+            default_catalog=self._default_catalog,
+            path=self._path,
+            environment_naming_info=environment_naming_info,
+        )
 
         start_time, end_time = (
             make_inclusive(start or c.EPOCH, end or c.EPOCH, self._dialect)
@@ -153,18 +182,17 @@ def _render(
 
         variables = kwargs.pop("variables", {})
         jinja_env_kwargs = {
-            **{**render_kwargs, **prepare_env(self._python_env), **variables},
+            **{
+                **render_kwargs,
+                **_prepare_python_env_for_jinja(macro_evaluator, self._python_env),
+                **variables,
+            },
             "snapshots": snapshots or {},
             "table_mapping": table_mapping,
             "deployability_index": deployability_index,
             "default_catalog": self._default_catalog,
             "runtime_stage": runtime_stage.value,
-            "resolve_table": lambda table: self._resolve_table(
-                d.normalize_model_name(table, self._default_catalog, self._dialect),
-                snapshots=snapshots,
-                table_mapping=table_mapping,
-                deployability_index=deployability_index,
-            ).sql(dialect=self._dialect, identify=True, comments=False),
+            "resolve_table": _resolve_table,
         }
         if this_model:
             render_kwargs["this_model"] = this_model
@@ -174,6 +202,7 @@ def _render(
 
         jinja_env = self._jinja_macro_registry.build_environment(**jinja_env_kwargs)
 
+        expressions = [self._expression]
         if isinstance(self._expression, d.Jinja):
             try:
                 expressions = []
@@ -190,29 +219,6 @@ def _render(
                     f"Could not render or parse jinja at '{self._path}'.\n{ex}"
                 ) from ex
 
-        macro_evaluator = MacroEvaluator(
-            self._dialect,
-            python_env=self._python_env,
-            jinja_env=jinja_env,
-            schema=self.schema,
-            runtime_stage=runtime_stage,
-            resolve_table=jinja_env.globals["resolve_table"],  # type: ignore
-            resolve_tables=lambda e: self._resolve_tables(
-                e,
-                snapshots=snapshots,
-                table_mapping=table_mapping,
-                deployability_index=deployability_index,
-                start=start,
-                end=end,
-                execution_time=execution_time,
-                runtime_stage=runtime_stage,
-            ),
-            snapshots=snapshots,
-            default_catalog=self._default_catalog,
-            path=self._path,
-            environment_naming_info=environment_naming_info,
-        )
-
         macro_evaluator.locals.update(render_kwargs)
 
         if variables:
@@ -637,3 +643,15 @@ def _optimize_query(self, query: exp.Query, all_deps: t.Set[str]) -> exp.Query:
                 annotate_types(select)
 
         return query
+
+
+def _prepare_python_env_for_jinja(
+    evaluator: MacroEvaluator,
+    python_env: t.Dict[str, Executable],
+) -> t.Dict[str, t.Any]:
+    prepared_env = prepare_env(python_env)
+    # Pass the evaluator to all macro functions
+    return {
+        key: partial(value, evaluator) if callable(value) else value
+        for key, value in prepared_env.items()
+    }

tests/core/test_model.py

@@ -2803,7 +2803,7 @@ def test_parse_expression_list_with_jinja():
 
 def test_no_depends_on_runtime_jinja_query():
     @macro()
-    def runtime_macro(**kwargs) -> None:
+    def runtime_macro(evaluator, **kwargs) -> None:
         from sqlmesh.utils.errors import ParsetimeAdapterCallError
 
         raise ParsetimeAdapterCallError("")
@@ -9051,3 +9051,40 @@ def test_formatting_flag_serde():
 
     deserialized_model = SqlModel.parse_raw(model_json)
     assert deserialized_model.dict() == model.dict()
+
+
+def test_call_python_macro_from_jinja():
+    def noop() -> None:
+        print("noop")
+
+    @macro()
+    def test_runtime_stage(evaluator):
+        noop()
+        return evaluator.runtime_stage
+
+    expressions = d.parse(
+        """
+        MODEL (
+            name db.table,
+            dialect spark,
+            owner owner_name,
+        );
+
+        JINJA_QUERY_BEGIN;
+        SELECT '{{ test_runtime_stage() }}' AS a, '{{ test_runtime_stage_jinja('bla') }}' AS b;
+        JINJA_END;
+    """
+    )
+
+    jinja_macros = JinjaMacroRegistry(
+        root_macros={
+            "test_runtime_stage_jinja": MacroInfo(
+                definition="{% macro test_runtime_stage_jinja(value) %}{{ test_runtime_stage() }}_{{ value }}{% endmacro %}",
+                depends_on=[],
+            )
+        }
+    )
+
+    model = load_sql_based_model(expressions, jinja_macros=jinja_macros)
+    assert model.render_query().sql() == "SELECT 'loading' AS a, 'loading_bla' AS b"
+    assert set(model.python_env) == {"noop", "test_runtime_stage"}