Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 17 additions & 1 deletion sqlmesh/utils/jinja.py
Original file line number Diff line number Diff line change
Expand Up @@ -200,6 +200,14 @@ def _extract(node: nodes.Node, parent: t.Optional[nodes.Node] = None) -> None:
return extracted


def is_variable_node(n: nodes.Node) -> bool:
return (
isinstance(n, nodes.Call)
and isinstance(n.node, nodes.Name)
and n.node.name in (c.VAR, c.BLUEPRINT_VAR)
)


def extract_macro_references_and_variables(
*jinja_strs: str, dbt_target_name: t.Optional[str] = None
) -> t.Tuple[t.Set[MacroReference], t.Set[str]]:
Expand Down Expand Up @@ -230,7 +238,15 @@ def extract_macro_references_and_variables(

for call_name, node in extract_call_names(jinja_str):
if call_name[0] in (c.VAR, c.BLUEPRINT_VAR):
assert isinstance(node, nodes.Call)
if not is_variable_node(node):
# Find the variable node which could be nested
for n in node.find_all(nodes.Call):
if is_variable_node(n):
node = n
break
else:
raise ValueError(f"Could not find variable name in {jinja_str}")
node = t.cast(nodes.Call, node)
args = [jinja_call_arg_name(arg) for arg in node.args]
if args and args[0]:
variable_name = args[0].lower()
Expand Down
13 changes: 12 additions & 1 deletion tests/dbt/converter/test_jinja.py
Original file line number Diff line number Diff line change
@@ -1,5 +1,9 @@
import pytest
from sqlmesh.utils.jinja import JinjaMacroRegistry, MacroExtractor
from sqlmesh.utils.jinja import (
JinjaMacroRegistry,
MacroExtractor,
extract_macro_references_and_variables,
)
from sqlmesh.dbt.converter.jinja import JinjaGenerator, convert_jinja_query, convert_jinja_macro
import sqlmesh.dbt.converter.jinja_transforms as jt
from pathlib import Path
Expand Down Expand Up @@ -437,3 +441,10 @@ def test_convert_jinja_macro(input: str, expected: str, sushi_dbt_context: Conte
result = convert_jinja_macro(sushi_dbt_context, input.strip())

assert " ".join(result.split()) == " ".join(expected.strip().split())


def test_extract_macro_references_and_variables() -> None:
input = """JINJA_QUERY('{%- set something = "'"~var("variable").split("|") -%}"""
_, variables = extract_macro_references_and_variables(input)
assert len(variables) == 1
assert variables == {"variable"}