fix(jinja): route recursive macro calls through the registry wrapper

sqlmesh/utils/jinja.py

@@ -542,9 +542,12 @@ def _parse_macro(self, name: str, package: t.Optional[str]) -> Template:
             macro = self._get_macro(name, package)
 
             definition: nodes.Template = self._environment.parse(macro.definition)
-            if _is_private_macro(name):
-                # A workaround to expose private jinja macros.
-                definition = self._to_non_private_macro_def(name, definition)
+            # Define the macro under an internal name, so that references to it from its own body
+            # (e.g. recursive calls) resolve to its _MacroWrapper, which handles return(). This
+            # also exposes private macros, since names starting with "_" are not exported by Jinja.
+            for node in definition.find_all(nodes.Macro):
+                if node.name == name:
+                    node.name = _internal_macro_name(name)
 
             self._parser_cache[cache_key] = self._environment.from_string(definition)
         return self._parser_cache[cache_key]
@@ -600,15 +603,6 @@ def _macro_exists(self, name: str, package: t.Optional[str]) -> bool:
     def _get_macro(self, name: str, package: t.Optional[str]) -> MacroInfo:
         return self.packages[package][name] if package is not None else self.root_macros[name]
 
-    def _to_non_private_macro_def(self, name: str, template: nodes.Template) -> nodes.Template:
-        for node in template.find_all((nodes.Macro, nodes.Call)):
-            if isinstance(node, nodes.Macro):
-                node.name = _non_private_name(name)
-            elif isinstance(node, nodes.Call) and isinstance(node.node, nodes.Name):
-                node.node.name = _non_private_name(name)
-
-        return template
-
     def _create_builtin_globals(self, global_vars: t.Dict[str, t.Any]) -> t.Dict[str, t.Any]:
         """Creates Jinja builtin globals using a factory function defined in the provided module."""
         engine_adapter = global_vars.pop("engine_adapter", None)
@@ -647,20 +641,21 @@ def __call__(self, *args: t.Any, **kwargs: t.Any) -> t.Any:
 
             template = self.registry._parse_macro(self.name, self.package)
             macro_callable = getattr(
-                template.make_module(vars=context), _non_private_name(self.name)
+                template.make_module(vars=context), _internal_macro_name(self.name)
             )
+            # Jinja uses this name in its own error messages
+            macro_callable.name = self.name
             try:
                 return macro_callable(*args, **kwargs)
             except MacroReturnVal as ret:
                 return ret.value
 
 
-def _is_private_macro(name: str) -> bool:
-    return name.startswith("_")
+_INTERNAL_MACRO_NAME_PREFIX = "sqlmesh_macro__"
 
 
-def _non_private_name(name: str) -> str:
-    return name.lstrip("_")
+def _internal_macro_name(name: str) -> str:
+    return f"{_INTERNAL_MACRO_NAME_PREFIX}{name}"
 
 
 JINJA_REGEX = re.compile(r"({{|{%)")

tests/utils/test_jinja.py

@@ -2,6 +2,9 @@
 
 from base64 import b64encode
 
+import pytest
+from jinja2 import UndefinedError
+
 from sqlmesh.utils import AttributeDict, yaml
 from sqlmesh.utils.jinja import (
     ENVIRONMENT,
@@ -10,6 +13,7 @@
     MacroReference,
     MacroReturnVal,
     call_name,
+    extract_error_details,
     nodes,
 )
 
@@ -90,6 +94,24 @@ def test_macro_registry_render_private_macros():
     assert rendered == "macro_a_a"
 
 
+def test_macro_registry_render_private_macro_calling_other_functions():
+    package_a = """
+{% macro macro_a_a(value) %}a{{ value }}{% endmacro %}
+
+{% macro _macro_a_b() %}{{ macro_a_a(1) }}-{{ range(3) | join(',') }}{% endmacro %}
+
+{% macro macro_a_c() %}{{ package_a._macro_a_b() }}{% endmacro %}
+"""
+
+    extractor = MacroExtractor()
+    registry = JinjaMacroRegistry()
+
+    registry.add_macros(extractor.extract(package_a), package="package_a")
+
+    rendered = registry.build_environment().from_string("{{ package_a.macro_a_c() }}").render()
+    assert rendered == "a1-0,1,2"
+
+
 def test_macro_registry_render_different_vars():
     package_a = "{% macro macro_a_a() %}{{ external() }}{% endmacro %}"
 
@@ -196,6 +218,62 @@ def macro_return(val):
     assert rendered == "[1, 2, 3]"
 
 
+def test_macro_return_recursive():
+    macros = """
+{% macro flatten(items) %}
+  {% set result = [] %}
+  {% for item in items %}
+    {% if item is iterable and item is not string %}
+      {% do result.extend(flatten(item)) %}
+    {% else %}
+      {% do result.append(item) %}
+    {% endif %}
+  {% endfor %}
+  {{ macro_return(result) }}
+{% endmacro %}
+"""
+
+    def macro_return(val):
+        raise MacroReturnVal(val)
+
+    extractor = MacroExtractor()
+    registry = JinjaMacroRegistry()
+
+    registry.add_macros(extractor.extract(macros))
+
+    rendered = (
+        registry.build_environment(macro_return=macro_return)
+        .from_string("{{ flatten([1, [2, 3], 4, [5, [6]]]) }}")
+        .render()
+    )
+    assert rendered == "[1, 2, 3, 4, 5, 6]"
+
+
+def test_macro_error_uses_macro_name():
+    registry = JinjaMacroRegistry()
+    registry.add_macros(MacroExtractor().extract("{% macro foo(a) %}{{ a }}{% endmacro %}"))
+
+    with pytest.raises(TypeError, match="macro 'foo' takes not more than 1 argument"):
+        registry.build_environment().from_string("{{ foo(1, 2) }}").render()
+
+
+def test_extract_error_details_private_macro():
+    package_a = "{% macro _macro_a() %}{{ undefined_macro() }}{% endmacro %}"
+
+    registry = JinjaMacroRegistry()
+    registry.add_macros(MacroExtractor().extract(package_a), package="package_a")
+
+    # extract_error_details reads the traceback of the exception being handled
+    try:
+        registry.build_environment().from_string("{{ package_a._macro_a() }}").render()
+    except UndefinedError as ex:
+        details = extract_error_details(ex)
+    else:
+        pytest.fail("Expected UndefinedError")
+
+    assert "in macro: '_macro_a'" in details
+
+
 def test_global_objs():
     original_registry = JinjaMacroRegistry(global_objs={"target": AttributeDict({"test": "value"})})