Feat: allow UNION macros as model queries (#2361)

Author: georgesittas

sqlmesh/core/model/definition.py

@@ -1477,8 +1477,12 @@ def load_sql_based_model(
         **meta_fields,
     )
 
-    if query_or_seed_insert is not None and isinstance(
-        query_or_seed_insert, (exp.Query, d.JinjaQuery)
+    if query_or_seed_insert is not None and (
+        isinstance(query_or_seed_insert, (exp.Query, d.JinjaQuery))
+        or (
+            isinstance(query_or_seed_insert, d.MacroFunc)
+            and query_or_seed_insert.this.name.lower() == "union"
+        )
     ):
         jinja_macro_references.update(
             extract_macro_references(query_or_seed_insert.sql(dialect=dialect))
@@ -1792,6 +1796,7 @@ def _split_sql_model_statements(
         if (
             isinstance(expression, (exp.Query, d.JinjaQuery))
             or expression == INSERT_SEED_MACRO_CALL
+            or (isinstance(expression, d.MacroFunc) and expression.this.name.lower() == "union")
         ):
             query_positions.append((expression, idx))
 

tests/core/test_model.py

@@ -190,7 +190,7 @@ def test_model_validation(query, error):
     assert error in str(ex.value)
 
 
-def test_model_union_query():
+def test_model_union_query(sushi_context, assert_exp_eq):
     expressions = d.parse(
         """
         MODEL (
@@ -204,6 +204,37 @@ def test_model_union_query():
 
     load_sql_based_model(expressions)
 
+    expressions = d.parse(
+        """
+        MODEL (
+            name sushi.test,
+            kind FULL,
+        );
+
+        @union('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").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"
+        """,
+    )
+
 
 def test_model_validation_union_query():
     expressions = d.parse(