Fix!: Have Spark put CTE at front of insert (#2086)

* have spark put cte at front of insert

* only copy if needed

sqlglot/dialects/spark2.py

@@ -114,6 +114,13 @@ def _unqualify_pivot_columns(expression: exp.Expression) -> exp.Expression:
     return expression
 
 
+def _insert_sql(self: Hive.Generator, expression: exp.Insert) -> str:
+    if expression.expression.args.get("with"):
+        expression = expression.copy()
+        expression.set("with", expression.expression.args.pop("with"))
+    return self.insert_sql(expression)
+
+
 class Spark2(Hive):
     class Parser(Hive.Parser):
         FUNCTIONS = {
@@ -202,6 +209,7 @@ class Generator(Hive.Generator):
             exp.DayOfYear: rename_func("DAYOFYEAR"),
             exp.FileFormatProperty: lambda self, e: f"USING {e.name.upper()}",
             exp.From: transforms.preprocess([_unalias_pivot]),
+            exp.Insert: _insert_sql,
             exp.LogicalAnd: rename_func("BOOL_AND"),
             exp.LogicalOr: rename_func("BOOL_OR"),
             exp.Map: _map_sql,

tests/dialects/test_spark.py

@@ -559,3 +559,13 @@ def test_transform_query(self):
         self.validate_identity(
             "SELECT TRANSFORM(zip_code, name, age) USING 'cat' FROM person WHERE zip_code > 94500"
         )
+
+    def test_insert_cte(self):
+        self.validate_all(
+            "INSERT OVERWRITE TABLE table WITH cte AS (SELECT cola FROM other_table) SELECT cola FROM cte",
+            write={
+                "spark": "WITH cte AS (SELECT cola FROM other_table) INSERT OVERWRITE TABLE table SELECT cola FROM cte",
+                "spark2": "WITH cte AS (SELECT cola FROM other_table) INSERT OVERWRITE TABLE table SELECT cola FROM cte",
+                "databricks": "WITH cte AS (SELECT cola FROM other_table) INSERT OVERWRITE TABLE table SELECT cola FROM cte",
+            },
+        )