Feat: Add union if macro for conditional union of tables

Author: themisvaltinos

sqlmesh/core/macros.py

@@ -1007,6 +1007,53 @@ def union(
     )
 
 
+@macro()
+def union_if(
+    evaluator: MacroEvaluator,
+    condition: exp.Expression,
+    type_: exp.Literal = exp.Literal.string("ALL"),
+    *tables: exp.Table,
+) -> 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.
+
+    Example:
+        >>> from sqlglot import parse_one
+        >>> from sqlglot.schema import MappingSchema
+        >>> from sqlmesh.core.macros import MacroEvaluator
+        >>> sql = "@UNION_IF(@start_ts = '2025-01-01 00:00:00', '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'
+    """
+    kind = type_.name.upper()
+    if kind not in ("ALL", "DISTINCT"):
+        raise SQLMeshError(f"Invalid type '{type_}'. Expected 'ALL' or 'DISTINCT'.")
+
+    columns = {
+        column
+        for column, _ in reduce(
+            lambda a, b: a & b,  # type: ignore
+            (evaluator.columns_to_types(table).items() for table in tables),
+        )
+    }
+
+    projections = [
+        exp.cast(column, type_, dialect=evaluator.dialect).as_(column)
+        for column, type_ in evaluator.columns_to_types(tables[0]).items()
+        if column in columns
+    ]
+
+    result = evaluator.eval_expression(condition)
+    if not result:
+        # If condition is false, return just the first table with proper column casting
+        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],
+    )
+
+
 @macro()
 def haversine_distance(
     _: MacroEvaluator,

tests/core/test_model.py

@@ -266,6 +266,70 @@ 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):
+    @macro()
+    def get_date(evaluator):
+        from sqlmesh.utils.date import now
+
+        return f"'{now().date()}'"
+
+    expressions = d.parse(
+        """
+        MODEL (
+            name sushi.test_1,
+            kind FULL,
+        );
+
+        @union_if(@get_date() == '1996-02-10', 'all', sushi.marketing, sushi.marketing)
+        """
+    )
+    sushi_context.upsert_model(load_sql_based_model(expressions, default_catalog="memory"))
+    assert_exp_eq(
+        sushi_context.get_model("sushi.test_1").render_query(),
+        """SELECT
+  CAST("marketing"."customer_id" AS INT) AS "customer_id",
+  CAST("marketing"."status" AS TEXT) AS "status",
+  CAST("marketing"."updated_at" AS TIMESTAMP) 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"
+UNION ALL
+SELECT
+  CAST("marketing"."customer_id" AS INT) AS "customer_id",
+  CAST("marketing"."status" AS TEXT) AS "status",
+  CAST("marketing"."updated_at" AS TIMESTAMP) 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"
+        """,
+    )
+
+    expressions = d.parse(
+        """
+        MODEL (
+            name sushi.test_2,
+            kind FULL,
+        );
+
+        @union_if(@get_date() > '1996-02-10', 'all', sushi.marketing, sushi.marketing)
+        """
+    )
+    sushi_context.upsert_model(load_sql_based_model(expressions, default_catalog="memory"))
+    assert_exp_eq(
+        sushi_context.get_model("sushi.test_2").render_query(),
+        """
+SELECT
+  CAST("marketing"."customer_id" AS INT) AS "customer_id",
+  CAST("marketing"."status" AS TEXT) AS "status",
+  CAST("marketing"."updated_at" AS TIMESTAMP) 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"
+        """,
+    )
+
+
 def test_model_validation_union_query():
     expressions = d.parse(
         """