diff --git a/sqlmesh/utils/jinja.py b/sqlmesh/utils/jinja.py index 829981db1c..494f331794 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,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"({{|{%)") diff --git a/tests/utils/test_jinja.py b/tests/utils/test_jinja.py index 01eeb47412..c1c7e69802 100644 --- a/tests/utils/test_jinja.py +++ b/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"})})