Fix macro func variable extraction & add tests

Author: georgesittas

sqlmesh/core/model/common.py

@@ -160,16 +160,22 @@ def make_python_env(
 def _extract_macro_func_variable_references(macro_func: exp.Expression) -> t.Set[str]:
     references = set()
 
-    for n in macro_func.walk():
-        if n is macro_func:
-            continue
+    # Don't descend into nested MacroFunc nodes besides @VAR() and @BLUEPRINT_VAR(), because
+    # they will be handled in a separate call of _extract_macro_func_variable_references.
+    def _prune_nested_macro_func(expression: exp.Expression) -> bool:
+        return (
+            type(n) is d.MacroFunc
+            and n is not macro_func
+            and n.this.name.lower() not in (c.VAR, c.BLUEPRINT_VAR)
+        )
 
-        # Don't descend into nested MacroFunc nodes besides @VAR() and @BLUEPRINT_VAR(), because
-        # they will be handled in a separate call of _extract_macro_func_variable_references.
-        if isinstance(n, d.MacroFunc):
+    for n in macro_func.walk(prune=_prune_nested_macro_func):
+        if type(n) is d.MacroFunc:
             this = n.this
-            if this.name.lower() in (c.VAR, c.BLUEPRINT_VAR) and this.expressions:
-                references.add(this.expressions[0].this.lower())
+            args = this.expressions
+
+            if this.name.lower() in (c.VAR, c.BLUEPRINT_VAR) and args and args[0].is_string:
+                references.add(args[0].this.lower())
         elif isinstance(n, d.MacroVar):
             references.add(n.name.lower())
         elif isinstance(n, (exp.Identifier, d.MacroStrReplace, d.MacroSQL)) and "@" in n.name:

tests/core/test_model.py

@@ -11117,3 +11117,22 @@ def entrypoint(context, **kwargs):
 
     assert model_daily is not None
     assert model_daily.cron == "@daily"
+
+
+@pytest.mark.parametrize(
+    "macro_func, variables",
+    [
+        ("@M(@v1)", {"v1"}),
+        ("@M(@{v1})", {"v1"}),
+        ("@M(@SQL('@v1'))", {"v1"}),
+        ("@M(@'@{v1}_foo')", {"v1"}),
+        ("@M1(@VAR('v1'))", {"v1"}),
+        ("@M1(@v1, @M2(@v2), @BLUEPRINT_VAR('v3'))", {"v1", "v3"}),
+        ("@M1(@BLUEPRINT_VAR(@VAR('v1')))", {"v1"}),
+    ],
+)
+def test_extract_macro_func_variable_references(macro_func: str, variables: t.Set[str]) -> None:
+    from sqlmesh.core.model.common import _extract_macro_func_variable_references
+
+    macro_func_ast = parse_one(macro_func)
+    assert _extract_macro_func_variable_references(macro_func_ast) == variables