fix cte emitter

Author: TrevorBergeron

bigframes/core/compile/sqlglot/sqlglot_ir.py

@@ -510,7 +510,11 @@ def _as_select(self) -> sge.Select:
         if isinstance(self.expr, sge.Select):
             return self.expr
         else:  # table
-            return sge.Select().from_(self.expr)
+            return (
+                sge.Select()
+                .select(sge.Column(this=sge.Star(), table=self.expr))
+                .from_(self.expr)
+            )
 
     def _as_subquery(self) -> sge.Subquery:
         return self._as_select().subquery()

bigframes/core/rewrite/as_sql.py

@@ -237,7 +237,9 @@ def _extract_ctes(root: nodes.BigFrameNode) -> nodes.BigFrameNode:
     return sql_nodes.SqlWithCtesNode(
         root.top_down(lambda x: mapping.get(x, x)),
         cte_names,
-        tuple(cte.top_down(lambda x: mapping.get(x, x)) for cte in topological_ctes),
+        tuple(
+            cte.child.top_down(lambda x: mapping.get(x, x)) for cte in topological_ctes  # type: ignore
+        ),
     )
 
 

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_binary_compiler/test_corr/out.sql

@@ -1,4 +1,9 @@
 SELECT
+  (
+    SELECT
+      `bfcol_2` AS `corr_col`
+    FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
+  ).*
 FROM (
   SELECT
     `bfcol_2` AS `corr_col`

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_binary_compiler/test_cov/out.sql

@@ -1,4 +1,9 @@
 SELECT
+  (
+    SELECT
+      `bfcol_2` AS `cov_col`
+    FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
+  ).*
 FROM (
   SELECT
     `bfcol_2` AS `cov_col`

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_nullary_compiler/test_size/out.sql

@@ -1,4 +1,9 @@
 SELECT
+  (
+    SELECT
+      `bfcol_32` AS `size`
+    FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
+  ).*
 FROM (
   SELECT
     `bfcol_32` AS `size`

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_ordered_unary_compiler/test_array_agg/out.sql

@@ -1,4 +1,9 @@
 SELECT
+  (
+    SELECT
+      `bfcol_1` AS `int64_col`
+    FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
+  ).*
 FROM (
   SELECT
     `bfcol_1` AS `int64_col`

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_ordered_unary_compiler/test_string_agg/out.sql

@@ -1,4 +1,9 @@
 SELECT
+  (
+    SELECT
+      `bfcol_1` AS `string_col`
+    FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
+  ).*
 FROM (
   SELECT
     `bfcol_1` AS `string_col`

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_all/out.sql

@@ -1,4 +1,10 @@
 SELECT
+  (
+    SELECT
+      `bfcol_2` AS `bool_col`,
+      `bfcol_3` AS `int64_col`
+    FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
+  ).*
 FROM (
   SELECT
     `bfcol_2` AS `bool_col`,

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_any/out.sql

@@ -1,4 +1,10 @@
 SELECT
+  (
+    SELECT
+      `bfcol_2` AS `bool_col`,
+      `bfcol_3` AS `int64_col`
+    FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
+  ).*
 FROM (
   SELECT
     `bfcol_2` AS `bool_col`,

tests/unit/core/compile/sqlglot/aggregations/snapshots/test_unary_compiler/test_any_value/out.sql

@@ -1,4 +1,9 @@
 SELECT
+  (
+    SELECT
+      `bfcol_1` AS `int64_col`
+    FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
+  ).*
 FROM (
   SELECT
     `bfcol_1` AS `int64_col`