add support for queries

Author: themisvaltinos

sqlmesh/core/macros.py

@@ -1012,10 +1012,9 @@ def union_if(
     evaluator: MacroEvaluator,
     condition: exp.Expression,
     type_: exp.Literal = exp.Literal.string("ALL"),
-    *tables: exp.Table,
+    *tables: exp.Table | exp.Query,
 ) -> exp.Query:
-    """Returns a UNION of the given tables if the condition is true, otherwise returns just the first table.
-       The behaviour remains of only choosing columns that have the same name and type.
+    """Returns a UNION of the given tables or queries if the condition is true, otherwise returns just the first.
 
     Example:
         >>> from sqlglot import parse_one
@@ -1029,11 +1028,23 @@ def union_if(
     if kind not in ("ALL", "DISTINCT"):
         raise SQLMeshError(f"Invalid type '{type_}'. Expected 'ALL' or 'DISTINCT'.")
 
+    result = evaluator.eval_expression(condition)
+
+    if isinstance(tables[0], exp.Query):
+        if not result:
+            # If condition is false, return just the first query
+            return tables[0]
+        return tables[0].union(*tables[1:], distinct=kind == "DISTINCT")
+
     columns = {
         column
         for column, _ in reduce(
             lambda a, b: a & b,  # type: ignore
-            (evaluator.columns_to_types(table).items() for table in tables),
+            (
+                evaluator.columns_to_types(table).items()
+                for table in tables
+                if isinstance(table, exp.Table)  # for mypy
+            ),
         )
     }
 

tests/core/test_model.py

@@ -267,7 +267,7 @@ def test_model_union_query(sushi_context, assert_exp_eq):
 
 
 @time_machine.travel("1996-02-10 00:00:00 UTC")
-def test_model_union_if_query(sushi_context, assert_exp_eq):
+def test_model_union_if_table(sushi_context, assert_exp_eq):
     @macro()
     def get_date(evaluator):
         from sqlmesh.utils.date import now
@@ -330,6 +330,43 @@ def get_date(evaluator):
     )
 
 
+def test_model_union_if_query(sushi_context, assert_exp_eq):
+    expressions = d.parse(
+        """
+        MODEL (
+            name sushi.test_query,
+            kind FULL,
+        );
+
+        @union_if(True, 'all', 'select 1 as c', 'select 2 as c', 'select 3 as c')
+        """
+    )
+    sushi_context.upsert_model(load_sql_based_model(expressions, default_catalog="memory"))
+
+    assert_exp_eq(
+        sushi_context.get_model("sushi.test_query").render_query(),
+        """SELECT 1 AS "c" UNION ALL SELECT 2 AS "c" UNION ALL SELECT 3 AS "c"
+        """,
+    )
+
+    expressions = d.parse(
+        """
+        MODEL (
+            name sushi.test_query,
+            kind FULL,
+        );
+
+        @union_if(False, 'all', 'select 1 as c', 'select 2 as c', 'select 3 as c')
+        """
+    )
+    sushi_context.upsert_model(load_sql_based_model(expressions, default_catalog="memory"))
+
+    assert_exp_eq(
+        sushi_context.get_model("sushi.test_query").render_query(),
+        'SELECT 1 AS "c"',
+    )
+
+
 def test_model_validation_union_query():
     expressions = d.parse(
         """