fix: sqlglot losses all CTE statements

Author: chelsea-lin

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -174,7 +174,7 @@ def from_query_string(
             alias=cte_name,
         )
         select_expr = sge.Select().select(sge.Star()).from_(sge.Table(this=cte_name))
-        select_expr.set("with", sge.With(expressions=[cte]))
+        select_expr.set("with_", sge.With(expressions=[cte]))
         return cls(expr=select_expr, uid_gen=uid_gen)
 
     @classmethod
@@ -197,7 +197,7 @@ def from_union(
             ), f"All provided expressions must be of type sge.Select, but got {type(select)}"
 
             select_expr = select.copy()
-            existing_ctes = [*existing_ctes, *select_expr.args.pop("with", [])]
+            existing_ctes = [*existing_ctes, *select_expr.args.pop("with_", [])]
 
             new_cte_name = sge.to_identifier(
                 next(uid_gen.get_uid_stream("bfcte_")), quoted=cls.quoted
@@ -229,7 +229,7 @@ def from_union(
             ),
         )
         final_select_expr = sge.Select().select(sge.Star()).from_(union_expr.subquery())
-        final_select_expr.set("with", sge.With(expressions=existing_ctes))
+        final_select_expr.set("with_", sge.With(expressions=existing_ctes))
         return cls(expr=final_select_expr, uid_gen=uid_gen)
 
     def select(
@@ -336,8 +336,8 @@ def join(
         left_select = _select_to_cte(self.expr, left_cte_name)
         right_select = _select_to_cte(right.expr, right_cte_name)
 
-        left_ctes = left_select.args.pop("with", [])
-        right_ctes = right_select.args.pop("with", [])
+        left_ctes = left_select.args.pop("with_", [])
+        right_ctes = right_select.args.pop("with_", [])
         merged_ctes = [*left_ctes, *right_ctes]
 
         join_on = _and(
@@ -353,7 +353,7 @@ def join(
             .from_(sge.Table(this=left_cte_name))
             .join(sge.Table(this=right_cte_name), on=join_on, join_type=join_type_str)
         )
-        new_expr.set("with", sge.With(expressions=merged_ctes))
+        new_expr.set("with_", sge.With(expressions=merged_ctes))
 
         return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
 
@@ -373,8 +373,8 @@ def isin_join(
         # Prefer subquery over CTE for the IN clause's right side to improve SQL readability.
         right_select = right.expr
 
-        left_ctes = left_select.args.pop("with", [])
-        right_ctes = right_select.args.pop("with", [])
+        left_ctes = left_select.args.pop("with_", [])
+        right_ctes = right_select.args.pop("with_", [])
         merged_ctes = [*left_ctes, *right_ctes]
 
         left_condition = typed_expr.TypedExpr(
@@ -415,7 +415,7 @@ def isin_join(
             .select(sge.Column(this=sge.Star(), table=left_cte_name), new_column)
             .from_(sge.Table(this=left_cte_name))
         )
-        new_expr.set("with", sge.With(expressions=merged_ctes))
+        new_expr.set("with_", sge.With(expressions=merged_ctes))
 
         return SQLGlotIR(expr=new_expr, uid_gen=self.uid_gen)
 
@@ -625,14 +625,14 @@ def _select_to_cte(expr: sge.Select, cte_name: sge.Identifier) -> sge.Select:
     into a new CTE and then generates a 'SELECT * FROM new_cte_name'
     for the new query."""
     select_expr = expr.copy()
-    existing_ctes = select_expr.args.pop("with", [])
+    existing_ctes = select_expr.args.pop("with_", [])
     new_cte = sge.CTE(
         this=select_expr,
         alias=cte_name,
     )
     new_with_clause = sge.With(expressions=[*existing_ctes, new_cte])
     new_select_expr = sge.Select().select(sge.Star()).from_(sge.Table(this=cte_name))
-    new_select_expr.set("with", new_with_clause)
+    new_select_expr.set("with_", new_with_clause)
     return new_select_expr