From 850b4979b32643899ac6aa994e8587396d876d4f Mon Sep 17 00:00:00 2001 From: JanakiRamChittineni <341781521+raams95@users.noreply.github.com> Date: Tue, 13 Oct 2026 15:33:47 -0500 Subject: [PATCH] fix(jinja): route recursive macro calls through the registry wrapper Macros were parsed under their own name, so a call to a macro from its own body resolved to the raw Jinja macro instead of the registry's _MacroWrapper. A return() inside the inner call therefore raised MacroReturnVal past the outer macro, which ended the outer call early and returned the inner call's value. Recursive dbt macros that collect results with return() (e.g. walking nested struct fields) only produced the result of the first recursive call. Define each macro under an internal name and call it by that name, so that references to the macro from within its body go through the wrapper, like calls from any other macro. The internal name is mapped back to the macro's own name in undefined variable error messages. This also replaces the private macro workaround, which renamed every call in a private macro's body to the macro itself and broke private macros that call other macros or functions. Fixes #5341 Signed-off-by: JanakiRamChittineni <341781521+raams95@users.noreply.github.com> --- sqlmesh/utils/jinja.py | 30 ++++++++++-------------- tests/utils/test_jinja.py | 49 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 61 insertions(+), 18 deletions(-) diff --git a/sqlmesh/utils/jinja.py b/sqlmesh/utils/jinja.py index 829981db1c..0df80bf681 100644 --- a/sqlmesh/utils/jinja.py +++ b/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,7 +641,7 @@ 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) ) try: return macro_callable(*args, **kwargs) @@ -655,12 +649,11 @@ def __call__(self, *args: t.Any, **kwargs: t.Any) -> t.Any: 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"({{|{%)") @@ -741,7 +734,8 @@ def extract_error_details(ex: Exception) -> str: if frame.f_code.co_name == "_invoke": macro = frame.f_locals.get("self") if isinstance(macro, Macro): - error_details += f" in macro: '{macro.name}'\n" + macro_name = macro.name.removeprefix(_INTERNAL_MACRO_NAME_PREFIX) + error_details += f" in macro: '{macro_name}'\n" break except: # to fall back to the generic error message if frame analysis fails diff --git a/tests/utils/test_jinja.py b/tests/utils/test_jinja.py index 01eeb47412..0b352cbad0 100644 --- a/tests/utils/test_jinja.py +++ b/tests/utils/test_jinja.py @@ -90,6 +90,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 +214,37 @@ 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_global_objs(): original_registry = JinjaMacroRegistry(global_objs={"target": AttributeDict({"test": "value"})})