Fix: Support coercion of SQL literal expressions into literal types for macro arguments (#4643)

Author: izeigerman

sqlmesh/core/macros.py

@@ -1452,6 +1452,26 @@ def _coerce(
         if base is SQL and isinstance(expr, exp.Expression):
             return expr.sql(dialect)
 
+        if base is t.Literal:
+            if not isinstance(expr, (exp.Literal, exp.Boolean)):
+                raise SQLMeshError(
+                    f"{base_err_msg} Coercion to {base} requires a literal expression."
+                )
+            literal_type_args = t.get_args(typ)
+            try:
+                for literal_type_arg in literal_type_args:
+                    expr_is_bool = isinstance(expr.this, bool)
+                    literal_is_bool = isinstance(literal_type_arg, bool)
+                    if (expr_is_bool and literal_is_bool and literal_type_arg == expr.this) or (
+                        not expr_is_bool
+                        and not literal_is_bool
+                        and str(literal_type_arg) == str(expr.this)
+                    ):
+                        return type(literal_type_arg)(expr.this)
+            except Exception:
+                raise SQLMeshError(base_err_msg)
+            raise SQLMeshError(base_err_msg)
+
         if isinstance(expr, base):
             return expr
         if issubclass(base, exp.Expression):

tests/core/test_macros.py

@@ -96,6 +96,12 @@ def test_default_arg_coercion(
     def test_select_macro(evaluator):
         return "SELECT 1 AS col"
 
+    @macro()
+    def test_literal_type(evaluator, a: t.Literal["test_literal_a", "test_literal_b", 1, True]):
+        if isinstance(a, exp.Expression):
+            raise SQLMeshError("Coercion failed")
+        return f"'{a}'"
+
     return MacroEvaluator(
         "hive",
         {"test": Executable(name="test", payload="def test(_):\n    return 'test'")},
@@ -1087,3 +1093,33 @@ def test_macro_with_spaces():
         ("d.@z", 'd.a."b c"'),
     ):
         assert evaluator.transform(parse_one(sql)).sql() == expected
+
+
+def test_macro_coerce_literal_type(macro_evaluator):
+    expression = d.parse_one("@TEST_LITERAL_TYPE('test_literal_a')")
+    assert macro_evaluator.transform(expression).sql() == "'test_literal_a'"
+
+    expression = d.parse_one("@TEST_LITERAL_TYPE('test_literal_b')")
+    assert macro_evaluator.transform(expression).sql() == "'test_literal_b'"
+
+    expression = d.parse_one("@TEST_LITERAL_TYPE(1)")
+    assert macro_evaluator.transform(expression).sql() == "'1'"
+
+    expression = d.parse_one("@TEST_LITERAL_TYPE(True)")
+    assert macro_evaluator.transform(expression).sql() == "'True'"
+
+    expression = d.parse_one("@TEST_LITERAL_TYPE('test_literal_c')")
+    with pytest.raises(MacroEvalError, match=".*Coercion failed"):
+        macro_evaluator.transform(expression)
+
+    expression = d.parse_one("@TEST_LITERAL_TYPE(2)")
+    with pytest.raises(MacroEvalError, match=".*Coercion failed"):
+        macro_evaluator.transform(expression)
+
+    expression = d.parse_one("@TEST_LITERAL_TYPE(False)")
+    with pytest.raises(MacroEvalError, match=".*Coercion failed"):
+        macro_evaluator.transform(expression)
+
+    expression = d.parse_one("@TEST_LITERAL_TYPE(1.0)")
+    with pytest.raises(MacroEvalError, match=".*Coercion failed"):
+        macro_evaluator.transform(expression)