Repository: TobikoData/sqlmesh
Pull request #2074 (2024-02-01 18:37:01)

Fix: Order projections during insert overwrite when there WHERE clause is not present


Patch:
diff --git a/sqlmesh/core/engine_adapter/base.py b/sqlmesh/core/engine_adapter/base.py
index 97573937..f0f60533 100644
--- a/sqlmesh/core/engine_adapter/base.py
+++ b/sqlmesh/core/engine_adapter/base.py
@@ -1022,10 +1022,8 @@ class EngineAdapter:
     ) -> None:
         if contains_json:
             query = self._escape_json(query)
-        if order_projections and query.named_selects != list(columns_to_types):
-            if isinstance(query, exp.Subqueryable):
-                query = query.subquery(alias="_ordered_projections")
-            query = self._select_columns(columns_to_types).from_(query)
+        if order_projections:
+            query = self._order_projections_and_filter(query, columns_to_types)
         self.execute(exp.insert(query, table_name, columns=list(columns_to_types)))
 
     def insert_overwrite_by_partition(
@@ -1105,7 +1103,7 @@ class EngineAdapter:
             columns_to_types = columns_to_types or self.columns(table_name)
             for i, source_query in enumerate(source_queries):
                 with source_query as query:
-                    query = self._add_where_to_query(query, where, columns_to_types)
+                    query = self._order_projections_and_filter(query, columns_to_types, where=where)
                     if i > 0 or insert_overwrite_strategy.is_delete_insert:
                         if i == 0:
                             self.delete_from(table_name, where=where or exp.true())
@@ -1764,22 +1762,24 @@ class EngineAdapter:
 
         return table
 
-    def _add_where_to_query(
+    def _order_projections_and_filter(
         self,
         query: Query,
-        where: t.Optional[exp.Expression],
         columns_to_types: t.Dict[str, exp.DataType],
+        where: t.Optional[exp.Expression] = None,
     ) -> Query:
-        if not where or not isinstance(query, exp.Subqueryable):
+        if not isinstance(query, exp.Subqueryable) or (
+            not where and query.named_selects == list(columns_to_types)
+        ):
             return query
 
         query = t.cast(exp.Subqueryable, query.copy())
         with_ = query.args.pop("with", None)
-        query = (
-            self._select_columns(columns_to_types)
-            .from_(query.subquery("_subquery", copy=False), copy=False)
-            .where(where, copy=False)
+        query = self._select_columns(columns_to_types).from_(
+            query.subquery("_subquery", copy=False), copy=False
         )
+        if where:
+            query = query.where(where, copy=False)
 
         if with_:
             query.set("with", with_)
diff --git a/sqlmesh/core/engine_adapter/mixins.py b/sqlmesh/core/engine_adapter/mixins.py
index b5014460..25d20000 100644
--- a/sqlmesh/core/engine_adapter/mixins.py
+++ b/sqlmesh/core/engine_adapter/mixins.py
@@ -101,7 +101,7 @@ class InsertOverwriteWithMergeMixin(EngineAdapter):
         columns_to_types = columns_to_types or self.columns(table_name)
         for source_query in source_queries:
             with source_query as query:
-                query = self._add_where_to_query(query, where, columns_to_types)
+                query = self._order_projections_and_filter(query, columns_to_types, where=where)
                 columns = [exp.to_column(col) for col in columns_to_types]
                 when_not_matched_by_source = exp.When(
                     matched=False,


Test patch:
diff --git a/tests/core/engine_adapter/test_base.py b/tests/core/engine_adapter/test_base.py
index 7e9fdb5d..2cbb6641 100644
--- a/tests/core/engine_adapter/test_base.py
+++ b/tests/core/engine_adapter/test_base.py
@@ -268,6 +268,30 @@ def test_insert_overwrite_by_time_partition_replace_where_pandas(
     ]
 
 
+def test_insert_overwrite_no_where(make_mocked_engine_adapter: t.Callable):
+    adapter = make_mocked_engine_adapter(EngineAdapter)
+
+    source_queries, columns_to_types = adapter._get_source_queries_and_columns_to_types(
+        parse_one("SELECT b, a FROM tbl"),
+        {"a": exp.DataType.build("INT"), "b": exp.DataType.build("STRING")},
+        target_table="test_table",
+    )
+
+    adapter._insert_overwrite_by_condition(
+        "test_table",
+        source_queries,
+        columns_to_types=columns_to_types,
+    )
+
+    adapter.cursor.begin.assert_called_once()
+    adapter.cursor.commit.assert_called_once()
+
+    assert to_sql_calls(adapter) == [
+        """DELETE FROM "test_table" WHERE TRUE""",
+        """INSERT INTO "test_table" ("a", "b") SELECT "a", "b" FROM (SELECT "b", "a" FROM "tbl") AS "_subquery" """.strip(),
+    ]
+
+
 def test_insert_append_query(make_mocked_engine_adapter: t.Callable):
     adapter = make_mocked_engine_adapter(EngineAdapter)
 
@@ -292,7 +316,7 @@ def test_insert_append_query_select_star(make_mocked_engine_adapter: t.Callable)
     )
 
     assert to_sql_calls(adapter) == [
-        'INSERT INTO "test_table" ("a", "b") SELECT "a", "b" FROM (SELECT 1 AS "a", * FROM "tbl") AS "_ordered_projections"',
+        'INSERT INTO "test_table" ("a", "b") SELECT "a", "b" FROM (SELECT 1 AS "a", * FROM "tbl") AS "_subquery"',
     ]
 
 
diff --git a/tests/core/engine_adapter/test_spark.py b/tests/core/engine_adapter/test_spark.py
index bfc3d473..5097a963 100644
--- a/tests/core/engine_adapter/test_spark.py
+++ b/tests/core/engine_adapter/test_spark.py
@@ -726,16 +726,25 @@ INSERT OVERWRITE TABLE `db`.`target` (
   `test_valid_to`
 )
 SELECT
-  *
-FROM `static`
-UNION ALL
-SELECT
-  *
-FROM `updated_rows`
-UNION ALL
-SELECT
-  *
-FROM `inserted_rows`
+  `id`,
+  `name`,
+  `price`,
+  `test_updated_at`,
+  `test_valid_from`,
+  `test_valid_to`
+FROM (
+  SELECT
+    *
+  FROM `static`
+  UNION ALL
+  SELECT
+    *
+  FROM `updated_rows`
+  UNION ALL
+  SELECT
+    *
+  FROM `inserted_rows`
+) AS `_subquery`
         """,
             dialect="spark",
         ).sql(dialect="spark"),