revise logic to extend current union; add unit tests

Author: themisvaltinos

sqlmesh/core/macros.py

@@ -970,93 +970,68 @@ 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'
-    """
-    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
-    ]
-
-    return reduce(
-        lambda a, b: a.union(b, distinct=kind == "DISTINCT"),  # type: ignore
-        [exp.select(*projections).from_(t) for t in tables],
-    )
-
-
-@macro()
-def union_if(
-    evaluator: MacroEvaluator,
-    condition: exp.Expression,
-    type_: exp.Literal = exp.Literal.string("ALL"),
-    *tables: exp.Table | exp.Query,
-) -> exp.Query:
-    """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
-        >>> from sqlglot.schema import MappingSchema
-        >>> from sqlmesh.core.macros import MacroEvaluator
-        >>> sql = "@UNION_IF(True, 'distinct', foo, 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 UNION.")
+
+    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 condition.")
+
+    # Check for union type
+    type_ = exp.Literal.string("ALL")
+    if arg_idx < len(args) and 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'.")
 
-    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")
+    # Remaining args should be tables
+    tables = args[arg_idx:]
 
     columns = {
         column
         for column, _ in reduce(
             lambda a, b: a & b,  # type: ignore
-            (
-                evaluator.columns_to_types(table).items()
-                for table in tables
-                if isinstance(table, exp.Table)  # for mypy
-            ),
+            (evaluator.columns_to_types(table).items() for table in tables),  # type: ignore
         )
     }
 
     projections = [
         exp.cast(column, type_, dialect=evaluator.dialect).as_(column)
-        for column, type_ in evaluator.columns_to_types(tables[0]).items()
+        for column, type_ in evaluator.columns_to_types(tables[0]).items()  # type: ignore
         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
+    # Skip the union if condition is False
+    if condition == False:
         return exp.select(*projections).from_(tables[0])
 
     return reduce(

tests/core/test_model.py

@@ -267,103 +267,98 @@ 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_table(sushi_context, assert_exp_eq):
+@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",
+        ),
+    ],
+)
+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()}'"
 
-    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
+    expected_select = """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"."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"
-        """,
-    )
-
-
-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"))
+    # Create tables argument list based on table_count
+    tables = ", ".join(["sushi.marketing"] * table_count)
 
-    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"
-        """,
-    )
+    # Handle the missing union_type case
+    union_type_arg = f", {union_type}" if union_type else ""
 
     expressions = d.parse(
-        """
+        f"""
         MODEL (
-            name sushi.test_query,
+            name sushi.{test_id},
             kind FULL,
         );
 
-        @union_if(False, 'all', 'select 1 as c', 'select 2 as c', 'select 3 as c')
+        @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("sushi.test_query").render_query(),
-        'SELECT 1 AS "c"',
+        sushi_context.get_model(f"sushi.{test_id}").render_query(),
+        expected_result(expected_select),
     )