Feat: Add condition in the union macro for conditional union of tables (#4337)

Author: themisvaltinos

docs/concepts/macros/sqlmesh_macros.md

@@ -855,7 +855,9 @@ FROM foo
 
 `@UNION` returns a `UNION` query that selects all columns with matching names and data types from the tables.
 
-Its first argument is the `UNION` "type", `'DISTINCT'` (removing duplicated rows) or `'ALL'` (returning all rows). Subsequent arguments are the tables to be combined.
+Its first argument can be either a condition or the `UNION` "type". If the first argument evaluates to a boolean (`TRUE` or `FALSE`), it's treated as a condition. If the condition is `FALSE`, only the first table is returned. If it's `TRUE`, the union operation is performed.
+
+If the first argument is not a boolean condition, it's treated as the `UNION` "type": either `'DISTINCT'` (removing duplicated rows) or `'ALL'` (returning all rows). Subsequent arguments are the tables to be combined.
 
 Let's assume that:
 
@@ -882,6 +884,47 @@ SELECT
 FROM bar
 ```
 
+If the union type is omitted, `'ALL'` is used as the default. So the following expression:
+
+```sql linenums="1"
+@UNION(foo, bar)
+```
+
+would be rendered as:
+
+```sql linenums="1"
+SELECT
+  CAST(a AS INT) AS a,
+  CAST(c AS TEXT) AS c
+FROM foo
+UNION ALL
+SELECT
+  CAST(a AS INT) AS a,
+  CAST(c AS TEXT) AS c
+FROM bar
+```
+
+You can also use a condition to control whether the union happens:
+
+```sql linenums="1"
+@UNION(1 > 0, 'all', foo, bar)
+```
+
+This would render the same as above. However, if the condition is `FALSE`:
+
+```sql linenums="1"
+@UNION(1 > 2, 'all', foo, bar)
+```
+
+Only the first table would be selected:
+
+```sql linenums="1"
+SELECT
+  CAST(a AS INT) AS a,
+  CAST(c AS TEXT) AS c
+FROM foo
+```
+
 ### @HAVERSINE_DISTANCE
 
 `@HAVERSINE_DISTANCE` returns the [haversine distance](https://en.wikipedia.org/wiki/Haversine_formula) between two geographic points.

sqlmesh/core/macros.py

@@ -970,23 +970,54 @@ def safe_div(_: MacroEvaluator, numerator: exp.Expression, denominator: exp.Expr
 @macro()
 def union(
     evaluator: MacroEvaluator,
-    type_: exp.Literal = exp.Literal.string("ALL"),
-    *tables: exp.Table,
+    *args: exp.Expression,
 ) -> exp.Query:
     """Returns a UNION of the given tables. Only choosing columns that have the same name and type.
 
+    Args:
+        evaluator: MacroEvaluator that invoked the macro
+        args: Variable arguments that can be:
+            - First argument can be a condition (exp.Condition)
+            - A union type ('ALL' or 'DISTINCT') as exp.Literal
+            - Tables (exp.Table)
+
     Example:
         >>> from sqlglot import parse_one
         >>> from sqlglot.schema import MappingSchema
         >>> from sqlmesh.core.macros import MacroEvaluator
         >>> sql = "@UNION('distinct', foo, bar)"
         >>> MacroEvaluator(schema=MappingSchema({"foo": {"a": "int", "b": "string", "c": "string"}, "bar": {"c": "string", "a": "int", "b": "int"}})).transform(parse_one(sql)).sql()
         'SELECT CAST(a AS INT) AS a, CAST(c AS TEXT) AS c FROM foo UNION SELECT CAST(a AS INT) AS a, CAST(c AS TEXT) AS c FROM bar'
+        >>> sql = "@UNION(True, 'distinct', foo, bar)"
+        >>> MacroEvaluator(schema=MappingSchema({"foo": {"a": "int", "b": "string", "c": "string"}, "bar": {"c": "string", "a": "int", "b": "int"}})).transform(parse_one(sql)).sql()
+        'SELECT CAST(a AS INT) AS a, CAST(c AS TEXT) AS c FROM foo UNION SELECT CAST(a AS INT) AS a, CAST(c AS TEXT) AS c FROM bar'
     """
+
+    if not args:
+        raise SQLMeshError("At least one table is required for the @UNION macro.")
+
+    arg_idx = 0
+    # Check for condition
+    condition = evaluator.eval_expression(args[arg_idx])
+    if isinstance(condition, bool):
+        arg_idx += 1
+        if arg_idx >= len(args):
+            raise SQLMeshError("Expected more arguments after the condition of the `@UNION` macro.")
+
+    # Check for union type
+    type_ = exp.Literal.string("ALL")
+    if isinstance(args[arg_idx], exp.Literal):
+        type_ = args[arg_idx]  # type: ignore
+        arg_idx += 1
     kind = type_.name.upper()
     if kind not in ("ALL", "DISTINCT"):
         raise SQLMeshError(f"Invalid type '{type_}'. Expected 'ALL' or 'DISTINCT'.")
 
+    # Remaining args should be tables
+    tables = [
+        exp.to_table(e.sql(evaluator.dialect), dialect=evaluator.dialect) for e in args[arg_idx:]
+    ]
+
     columns = {
         column
         for column, _ in reduce(
@@ -1001,6 +1032,10 @@ def union(
         if column in columns
     ]
 
+    # Skip the union if condition is False
+    if condition == False:
+        return exp.select(*projections).from_(tables[0])
+
     return reduce(
         lambda a, b: a.union(b, distinct=kind == "DISTINCT"),  # type: ignore
         [exp.select(*projections).from_(t) for t in tables],

tests/core/test_model.py

@@ -266,6 +266,142 @@ def test_model_union_query(sushi_context, assert_exp_eq):
     )
 
 
+@time_machine.travel("1996-02-10 00:00:00 UTC")
+@pytest.mark.parametrize(
+    "test_id, condition, union_type, table_count, expected_result",
+    [
+        # Test case 1: Basic conditional union - True condition
+        (
+            "test_1",
+            "@get_date() == '1996-02-10'",
+            "'all'",
+            2,
+            lambda expected_select: f"{expected_select}\nUNION ALL\n{expected_select}\n",
+        ),
+        # Test case 2: False condition - should return just first table
+        (
+            "test_2",
+            "@get_date() > '1996-02-10'",
+            "'all'",
+            2,
+            lambda expected_select: f"{expected_select}\n",
+        ),
+        # Test case 3: Multiple tables in union
+        (
+            "test_3",
+            "@get_date() == '1996-02-10'",
+            "'all'",
+            3,
+            lambda expected_select: f"{expected_select}\nUNION ALL\n{expected_select}\nUNION ALL\n{expected_select}\n",
+        ),
+        # Test case 4: DISTINCT type
+        (
+            "test_4",
+            "@get_date() == '1996-02-10'",
+            "'distinct'",
+            2,
+            lambda expected_select: f"{expected_select}\nUNION\n{expected_select}\n",
+        ),
+        # Test case 5: Complex condition
+        (
+            "test_5",
+            "@get_date() = '1996-02-10' and 1=1 or @get_date() > '1996-02-10'",
+            "'distinct'",
+            2,
+            lambda expected_select: f"{expected_select}\nUNION\n{expected_select}\n",
+        ),
+        # Test case 6: Missing union type (defaults to ALL)
+        (
+            "test_6",
+            "@get_date() == '1996-02-10'",
+            "",
+            2,
+            lambda expected_select: f"{expected_select}\nUNION ALL\n{expected_select}\n",
+        ),
+        # Test case 7: Missing union type AND condition
+        (
+            "test_7",
+            "",
+            "",
+            2,
+            lambda expected_select: f"{expected_select}\nUNION ALL\n{expected_select}\n",
+        ),
+        # Test case 8: Missing union type AND condition multiple tables
+        (
+            "test_8",
+            "",
+            "",
+            3,
+            lambda expected_select: f"{expected_select}\nUNION ALL\n{expected_select}\n\nUNION ALL\n{expected_select}\n",
+        ),
+        # Test case 9: Missing union type AND condition one table
+        (
+            "test_9",
+            "",
+            "",
+            1,
+            lambda expected_select: f"{expected_select}",
+        ),
+        # Test case 10: Union type with one table
+        (
+            "test_10",
+            "",
+            "'distinct'",
+            1,
+            lambda expected_select: f"{expected_select}",
+        ),
+        # Test case 11: Condition with one table
+        (
+            "test_9",
+            "True",
+            "",
+            1,
+            lambda expected_select: f"{expected_select}",
+        ),
+    ],
+)
+def test_model_union_conditional(
+    sushi_context, assert_exp_eq, test_id, condition, union_type, table_count, expected_result
+):
+    @macro()
+    def get_date(evaluator):
+        from sqlmesh.utils.date import now
+
+        return f"'{now().date()}'"
+
+    expected_select = """SELECT
+  CAST("marketing"."customer_id" AS INT) AS "customer_id",
+  CAST("marketing"."status" AS TEXT) AS "status",
+  CAST("marketing"."updated_at" AS TIMESTAMPNTZ) AS "updated_at",
+  CAST("marketing"."valid_from" AS TIMESTAMP) AS "valid_from",
+  CAST("marketing"."valid_to" AS TIMESTAMP) AS "valid_to"
+FROM "memory"."sushi"."marketing" AS "marketing"
+"""
+
+    # Create tables argument list based on table_count
+    tables = ", ".join(["sushi.marketing"] * table_count)
+
+    # Handle the missing union_type case
+    union_type_arg = f", {union_type}" if union_type else ""
+
+    expressions = d.parse(
+        f"""
+        MODEL (
+            name sushi.{test_id},
+            kind FULL,
+        );
+
+        @union({condition}{union_type_arg}, {tables})
+        """
+    )
+    sushi_context.upsert_model(load_sql_based_model(expressions, default_catalog="memory"))
+
+    assert_exp_eq(
+        sushi_context.get_model(f"sushi.{test_id}").render_query(),
+        expected_result(expected_select),
+    )
+
+
 def test_model_validation_union_query():
     expressions = d.parse(
         """