Skip to content

Commit cd11ba5

Browse files
committed
fix: find variable node if nested
1 parent a7feeb7 commit cd11ba5

File tree

2 files changed

+28
-2
lines changed

2 files changed

+28
-2
lines changed

sqlmesh/utils/jinja.py

Lines changed: 16 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -228,9 +228,24 @@ def extract_macro_references_and_variables(
228228
)
229229
)
230230

231+
def is_variable_node(n: nodes.Node) -> bool:
232+
return (
233+
isinstance(n, nodes.Call)
234+
and isinstance(n.node, nodes.Name)
235+
and n.node.name in (c.VAR, c.BLUEPRINT_VAR)
236+
)
237+
231238
for call_name, node in extract_call_names(jinja_str):
232239
if call_name[0] in (c.VAR, c.BLUEPRINT_VAR):
233-
assert isinstance(node, nodes.Call)
240+
if not is_variable_node(node):
241+
# Find the variable node which could be nested
242+
for n in node.find_all(nodes.Call):
243+
if is_variable_node(n):
244+
node = n
245+
break
246+
else:
247+
raise ValueError(f"Could not find variable name in {jinja_str}")
248+
node = t.cast(nodes.Call, node)
234249
args = [jinja_call_arg_name(arg) for arg in node.args]
235250
if args and args[0]:
236251
variable_name = args[0].lower()

tests/dbt/converter/test_jinja.py

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,9 @@
11
import pytest
2-
from sqlmesh.utils.jinja import JinjaMacroRegistry, MacroExtractor
2+
from sqlmesh.utils.jinja import (
3+
JinjaMacroRegistry,
4+
MacroExtractor,
5+
extract_macro_references_and_variables,
6+
)
37
from sqlmesh.dbt.converter.jinja import JinjaGenerator, convert_jinja_query, convert_jinja_macro
48
import sqlmesh.dbt.converter.jinja_transforms as jt
59
from pathlib import Path
@@ -437,3 +441,10 @@ def test_convert_jinja_macro(input: str, expected: str, sushi_dbt_context: Conte
437441
result = convert_jinja_macro(sushi_dbt_context, input.strip())
438442

439443
assert " ".join(result.split()) == " ".join(expected.strip().split())
444+
445+
446+
def test_extract_macro_references_and_variables() -> None:
447+
input = """JINJA_QUERY('{%- set something = "'"~var("variable").split("|") -%}"""
448+
_, variables = extract_macro_references_and_variables(input)
449+
assert len(variables) == 1
450+
assert variables == {"variable"}

0 commit comments

Comments
 (0)