Feature: add the if macro. (#796)

Author: tobymao

Makefile

@@ -28,7 +28,7 @@ it-test: core-it-test airflow-it-test-with-env
 
 it-test-docker: core-it-test airflow-it-test-docker-with-env
 
-test: unit-test it-test doc-test
+test: unit-test doc-test it-test
 
 package:
 	pip3 install wheel && python3 setup.py sdist bdist_wheel

sqlmesh/core/macros.py

@@ -188,14 +188,18 @@ def evaluate(self, node: MacroFunc) -> exp.Expression | t.List[exp.Expression] |
             return [exp.convert(item) for item in result if item is not None]
         return exp.convert(result)
 
-    def eval_expression(self, node: exp.Expression) -> t.Any:
+    def eval_expression(self, node: t.Any) -> t.Any:
         """Converts a SQLGlot expression into executable Python code and evals it.
 
+        If the node is not an expression, it will simply be returned.
+
         Args:
             node: expression
         Returns:
             The return value of the evaled Python Code.
         """
+        if not isinstance(node, exp.Expression):
+            return node
         code = node.sql()
         try:
             code = self.generator.generate(node)
@@ -333,6 +337,31 @@ def each(
     return [item for item in map(func, ensure_collection(items)) if item is not None]
 
 
+@macro("IF")
+def if_(
+    evaluator: MacroEvaluator,
+    condition: t.Any,
+    true: t.Any,
+    false: t.Any = None,
+) -> t.Any:
+    """Evaluates a given condition and returns the second argument if true or else the third argument.
+
+    If false is not passed in, the default return value will be None.
+
+    Example:
+        >>> from sqlglot import parse_one
+        >>> from sqlmesh.core.macros import MacroEvaluator
+        >>> MacroEvaluator().transform(parse_one("@IF('a' = 1, a, b)")).sql()
+        'b'
+
+        >>> MacroEvaluator().transform(parse_one("@IF('a' = 1, a)"))
+    """
+
+    if evaluator.eval_expression(condition):
+        return true
+    return false
+
+
 @macro("REDUCE")
 def reduce_(evaluator: MacroEvaluator, *args: t.Any) -> t.Any:
     """Iterates through items applying provided function that takes two arguments

tests/core/test_macros.py

@@ -168,6 +168,26 @@ def test_ast_correctness(macro_evaluator):
             "SELECT * FROM city",
             {"do_order": False},
         ),
+        (
+            """select @if(TRUE, 1, 0)""",
+            "SELECT 1",
+            {},
+        ),
+        (
+            """select @if(FALSE, 1, 0)""",
+            "SELECT 0",
+            {},
+        ),
+        (
+            """select @if(1 > 0, 1, 0)""",
+            "SELECT 1",
+            {},
+        ),
+        (
+            """select @if('a' = 'b', c), d""",
+            "SELECT d",
+            {},
+        ),
     ],
 )
 def test_macro_functions(macro_evaluator, assert_exp_eq, sql, expected, args):