Fix: Place LIMIT before ORDER BY for UNION queries when generating a CTAS query (#2277)

Author: izeigerman

sqlmesh/core/model/definition.py

@@ -424,13 +424,15 @@ def ctas_query(self, **render_kwarg: t.Any) -> exp.Query:
         """
         query = self.render_query_or_raise(**render_kwarg).copy()
 
-        for select in query.find_all(exp.Select):
-            if select.args.get("from"):
-                select.where(exp.false(), copy=False)
-                if not isinstance(select.parent, exp.Union) or (
-                    select.parent.parent is None and select.arg_key == "expression"
-                ):
-                    select.limit(0, copy=False)
+        for select_or_union in query.find_all(exp.Select, exp.Union):
+            if isinstance(select_or_union, exp.Select) and select_or_union.args.get("from"):
+                select_or_union.where(exp.false(), copy=False)
+                if not isinstance(select_or_union.parent, exp.Union):
+                    select_or_union.limit(0, copy=False)
+            elif isinstance(select_or_union, exp.Union) and not isinstance(
+                select_or_union.parent, exp.Union
+            ):
+                select_or_union.set("limit", exp.Limit(expression=exp.Literal.number(0)))
 
         if self.managed_columns:
             query.select(

tests/core/test_model.py

@@ -1799,6 +1799,18 @@ def test_model_ctas_query():
         == 'SELECT 1 AS "a" FROM "t" AS "t" WHERE FALSE UNION ALL SELECT 2 AS "a" FROM "t" AS "t" WHERE FALSE LIMIT 0'
     )
 
+    expressions = d.parse(
+        """
+        MODEL (name `a-b-c.table`, kind FULL, dialect bigquery);
+        SELECT 1 AS a FROM t UNION ALL SELECT 2 AS a FROM t ORDER BY 1
+    """
+    )
+
+    assert (
+        load_sql_based_model(expressions, dialect="bigquery").ctas_query().sql()
+        == 'SELECT 1 AS "a" FROM "t" AS "t" WHERE FALSE UNION ALL SELECT 2 AS "a" FROM "t" AS "t" WHERE FALSE ORDER BY 1 LIMIT 0'
+    )
+
 
 def test_is_breaking_change():
     model = create_external_model("a", columns={"a": "int", "limit": "int"})