From 5150dc31f6086d58e8977cb6541b73cb0c81355c Mon Sep 17 00:00:00 2001
From: Toby Mao <toby.mao@gmail.com>
Date: Thu, 4 Apr 2024 19:59:15 -0700
Subject: [PATCH] fix: allow macros to recursively resolve closes #2399

---
 sqlmesh/core/macros.py    | 31 +++++++++++++++++++++++++++++--
 tests/core/test_audit.py  | 12 ++++++++++++
 tests/core/test_macros.py |  7 +++++++
 3 files changed, 48 insertions(+), 2 deletions(-)

diff --git a/sqlmesh/core/macros.py b/sqlmesh/core/macros.py
index f8edb22dc7..4fdffe676b 100644
--- a/sqlmesh/core/macros.py
+++ b/sqlmesh/core/macros.py
@@ -3,6 +3,7 @@
 import typing as t
 from enum import Enum
 from functools import reduce
+from itertools import chain
 from string import Template
 
 import sqlglot
@@ -177,7 +178,15 @@ def evaluate_macros(
                     return node
 
                 value = self.locals[node.name]
-                return exp.convert(tuple(value) if isinstance(value, list) else value)
+                if isinstance(value, list):
+                    return exp.convert(
+                        tuple(
+                            self.transform(v) if isinstance(v, exp.Expression) else v for v in value
+                        )
+                    )
+                return exp.convert(
+                    self.transform(value) if isinstance(value, exp.Expression) else value
+                )
             if isinstance(node, exp.Identifier):
                 text = self.template(node.this, self.locals)
                 if node.this != text:
@@ -222,7 +231,25 @@ def template(self, text: t.Any, local_variables: t.Dict[str, t.Any]) -> str:
         Returns:
            The rendered string.
         """
-        return MacroStrTemplate(str(text)).safe_substitute(self.locals, **local_variables)
+        mapping = {}
+
+        for k, v in chain(self.locals.items(), local_variables.items()):
+            # try to convert all variables into sqlglot expressions
+            # because they're going to be converted into strings in sql
+            # we use bare Exception instead of ValueError because there's
+            # a recursive error with MagicMock.
+            # we don't convert strings because that would result in adding quotes
+            if not isinstance(v, str):
+                try:
+                    v = exp.convert(v)
+                except Exception:
+                    pass
+
+            if isinstance(v, exp.Expression):
+                v = v.sql(dialect=self.dialect)
+            mapping[k] = v
+
+        return MacroStrTemplate(str(text)).safe_substitute(mapping)
 
     def evaluate(self, node: MacroFunc) -> exp.Expression | t.List[exp.Expression] | None:
         if isinstance(node, MacroDef):
diff --git a/tests/core/test_audit.py b/tests/core/test_audit.py
index 25ab79fd33..f418e0c31a 100644
--- a/tests/core/test_audit.py
+++ b/tests/core/test_audit.py
@@ -686,3 +686,15 @@ def test_not_constant_audit(model: Model):
         rendered_query.sql()
         == """SELECT 1 AS "1" FROM (SELECT COUNT(DISTINCT "x") AS "t_cardinality" FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_q_0" WHERE "x" > 1) AS "r" WHERE "r"."t_cardinality" <= 1"""
     )
+
+
+def test_condition_with_macro_var(model: Model):
+    rendered_query = builtin.not_null_audit.render_query(
+        model,
+        columns=[exp.column("x")],
+        condition=exp.condition("dt BETWEEN @start_dt AND @end_dt"),
+    )
+    assert (
+        rendered_query.sql(dialect="duckdb")
+        == """SELECT * FROM (SELECT * FROM "db"."test_model" AS "test_model" WHERE "ds" BETWEEN '1970-01-01' AND '1970-01-01') AS "_q_0" WHERE "x" IS NULL AND "dt" BETWEEN CAST('1970-01-01 00:00:00+00:00' AS TIMESTAMP) AND CAST('1970-01-01 23:59:59.999999+00:00' AS TIMESTAMP)"""
+    )
diff --git a/tests/core/test_macros.py b/tests/core/test_macros.py
index c05b67bfa2..0c57f48756 100644
--- a/tests/core/test_macros.py
+++ b/tests/core/test_macros.py
@@ -332,6 +332,13 @@ def test_ast_correctness(macro_evaluator):
             "SELECT * FROM x ORDER BY a, b",
             {"x": [exp.column("a"), exp.column("b")]},
         ),
+        (
+            """select @SQL('@x')""",
+            "SELECT VAR_MAP('a', 1)",
+            {
+                "x": {"a": 1},
+            },
+        ),
     ],
 )
 def test_macro_functions(macro_evaluator, assert_exp_eq, sql, expected, args):