feat: execute and re-write plain CTEs

integration/rust/tests/integration/subqueries/tests/insert.sql

@@ -133,3 +133,19 @@ INSERT INTO subpackages (id, user_id) SELECT id, id FROM subusers ORDER BY id DE
 -- route insert: 0
 -- verify: SELECT * FROM subpackages ORDER BY id
 INSERT INTO subpackages (id, user_id, order_id) SELECT suborders.id, suborders.user_id, suborders.id FROM suborders JOIN subusers ON suborders.user_id = subusers.id
+
+-- case
+-- what: WITH a AS (..), b AS (SELECT .. FROM a) INSERT .. FROM a JOIN b .., one CTE over another
+-- route select: 0,1
+-- route select: any
+-- route insert: 0
+-- verify: SELECT * FROM subpackages ORDER BY id
+WITH a AS (SELECT id, user_id FROM suborders), b AS (SELECT id FROM a WHERE user_id = 1) INSERT INTO subpackages (id, user_id, order_id) SELECT a.id, a.user_id, a.id FROM a JOIN b ON a.id = b.id
+
+-- case
+-- what: WITH a AS (..), b AS (SELECT .. FROM a) INSERT .. FROM a JOIN b .., same CTE referenced more than once
+-- route select: 0,1
+-- route select: any
+-- route insert: 0
+-- verify: SELECT * FROM subpackages ORDER BY id
+WITH a AS (SELECT id, user_id FROM suborders) INSERT INTO subpackages (id, user_id, order_id) SELECT x.id, x.user_id, x.id FROM a x JOIN a y ON x.id = y.id

integration/rust/tests/integration/subqueries/tests/select.sql

@@ -186,3 +186,47 @@ SELECT s.id, s.status FROM (
 SELECT s.id, s.status FROM (
     SELECT id, status FROM subusers WHERE status = 'some random status which doesnt exist'
 ) s ORDER BY s.id DESC LIMIT 2;
+
+-- case
+-- what: WITH c AS (..) SELECT .. WHERE <col> IN (SELECT .. FROM c)
+-- route select: 0,1
+-- route select: 0,1
+WITH c AS (
+    SELECT id FROM subusers WHERE status = 'VIP'
+) SELECT * FROM suborders WHERE id IN (
+    SELECT id FROM c
+) ORDER BY id;
+
+-- case
+-- what: WITH c (a, b) AS (..) SELECT .. FROM c
+-- route select: 0,1
+-- route select: any
+WITH c (a, b) AS (
+    SELECT id, status FROM subusers
+) SELECT a, b FROM c ORDER BY a;
+
+-- case
+-- what: WITH c AS (SELECT .. FROM <st1>) SELECT <col>, <agg>(..) FROM c GROUP BY <col>
+-- route select: 0,1
+-- route select: any
+WITH c AS (
+    SELECT id, status FROM subusers
+) SELECT status, count(*) FROM c GROUP BY status ORDER BY status;
+
+-- case
+-- what: WITH a AS (..), b AS (SELECT .. FROM a) .. FROM a JOIN b .., one CTE over another
+-- route select: 0,1
+-- route select: any
+WITH a AS (
+    SELECT id, status FROM subusers
+), b AS (
+    SELECT id FROM a WHERE status = 'VIP'
+) SELECT a.id, a.status FROM a JOIN b ON a.id = b.id ORDER BY a.id;
+
+-- case
+-- what: WITH a AS (..), b AS (SELECT .. FROM a) .. FROM a JOIN b .., same CTE referenced more than once
+-- route select: 0,1
+-- route select: any
+WITH a AS (
+    SELECT id, status FROM subusers
+) SELECT x.id, y.id FROM a x JOIN a y ON x.status = y.status ORDER BY x.id, y.id;

integration/rust/tests/integration/subqueries/todo.sql

@@ -95,12 +95,12 @@
 -- [DONE] [A1 A2 A3 A4 A5 A6 O9 Q2 Q3] [update] UPDATE <st1> SET .. FROM (<sq>) s WHERE ..
 
 -- Plain CTEs
---   [A1 A2 A3 A4 A5 A6] [select] WITH c AS (..) SELECT .. WHERE <col> IN (SELECT .. FROM c)
---   [A1 A2 A4 A6] [select] WITH c (a, b) AS (..) SELECT .. FROM c
---   [A1 A2 A3 A4 A5 A6] [select] WITH c AS (SELECT .. FROM <st1>) SELECT <col>, <agg>(..) FROM c GROUP BY <col>
---   [A1 A2 A3 A4 A5 A6] [select, insert] WITH a AS (..), b AS (SELECT .. FROM a) .. FROM a JOIN b ..
---       case: one CTE over another
---       case: same CTE referenced more than once
+-- [DONE] [A1 A2 A3 A4 A5 A6] [select] WITH c AS (..) SELECT .. WHERE <col> IN (SELECT .. FROM c)
+-- [DONE] [A1 A2 A4 A6] [select] WITH c (a, b) AS (..) SELECT .. FROM c
+-- [DONE] [A1 A2 A3 A4 A5 A6] [select] WITH c AS (SELECT .. FROM <st1>) SELECT <col>, <agg>(..) FROM c GROUP BY <col>
+-- [DONE] [A1 A2 A3 A4 A5 A6] [select, insert] WITH a AS (..), b AS (SELECT .. FROM a) .. FROM a JOIN b ..
+--   case: one CTE over another
+--   case: same CTE referenced more than once
 
 -- WHERE follow-ups
 --   [A1 A2 A3 A4 A5 A6] [select] .. WHERE <col> IN (<sq>) <op> EXISTS (<sq>)

pgdog/src/frontend/client/query_engine/multi_step/error.rs

@@ -27,6 +27,9 @@ pub(crate) enum Error {
 
     #[error("net: {0}")]
     Net(#[from] crate::net::Error),
+
+    #[error("sql parser: {0}")]
+    SqlParser(#[from] pg_raw_parse::error::Error),
 }
 
 #[derive(Debug, Error)]

pgdog/src/frontend/client/query_engine/multi_step/subquery/from_clause.rs

@@ -1,117 +1,102 @@
 use std::ops::Deref;
 
 use pg_raw_parse::{
-    Node, NodeMut, deparse,
-    make::{MemoryToken, Unique, owned},
+    Node, NodeMut,
+    make::{MemoryToken, Unique},
     nodes::{self, Alias, SelectStmt},
     transform,
 };
 use pgdog_postgres_types::{DataType, Datum};
 
 use crate::{
-    frontend::client::query_engine::multi_step::{Error, SubqueryMulti},
+    frontend::client::query_engine::multi_step::SubqueryMulti,
     net::{DataRow, Decoder, RowDescription},
 };
 
 impl<'a> SubqueryMulti<'a> {
-    pub fn handle_from_clause(
+    pub fn handle_from_clause<'mem>(
         &self,
-        statement: Node<'a>,
+        mem: MemoryToken<'mem>,
+        mut statement: Unique<'mem, Node<'mem>>,
         row_description: Option<&RowDescription>,
         data_rows: &[DataRow],
-    ) -> Result<Option<String>, Error> {
+    ) -> Unique<'mem, Node<'mem>> {
         let Some(row_description) = row_description else {
-            return Ok(None);
+            return statement;
         };
 
-        let new_statement = owned(|mem| {
-            let owned_stmt = mem.make_unique(statement);
-            let mut owned_raw_stmt = mem.make_raw_stmt(owned_stmt);
-            let mutable_raw_stmt = owned_raw_stmt.as_mut();
+        let new_from_clause = mem.make_unique(match statement.as_ref() {
+            Node::SelectStmt(stmt) => stmt.from_clause(),
+            Node::UpdateStmt(stmt) => stmt.from_clause(),
+            _ => unreachable!("The parser only allows through SELECT and UPDATE"),
+        });
 
-            let new_from_clause = mem.make_unique(match mutable_raw_stmt.stmt() {
-                Node::SelectStmt(stmt) => stmt.from_clause(),
-                Node::UpdateStmt(stmt) => stmt.from_clause(),
-                _ => unreachable!("The parser only allows through SELECT and UPDATE"),
-            });
-
-            let mut new_from_clause = mem.make_unique(new_from_clause.first().unwrap());
-
-            transform::transform(&mut new_from_clause, |node| match &*node {
-                NodeMut::RangeSubselect(range_subselect) => {
-                    let mut range_subselect = mem.make_unique(range_subselect.deref());
-                    // Replace the subquery in the statement with the DataRows (via VALUES)
-                    {
-                        let explicit_values_list = Self::new_select_with_explicit_values(
-                            mem,
-                            row_description.clone(),
-                            data_rows,
-                        );
-
-                        range_subselect
-                            .as_mut()
-                            .set_subquery(explicit_values_list.uncast());
-                    }
+        let mut new_from_clause = mem.make_unique(new_from_clause.first().unwrap());
+        transform::transform(&mut new_from_clause, |node| match &*node {
+            NodeMut::RangeSubselect(range_subselect) => {
+                let mut range_subselect = mem.make_unique(range_subselect.deref());
+                // Replace the subquery in the statement with the DataRows (via VALUES)
+                {
+                    let explicit_values_list = Self::new_select_with_explicit_values(
+                        mem,
+                        row_description.clone(),
+                        data_rows,
+                    );
+
+                    range_subselect
+                        .as_mut()
+                        .set_subquery(explicit_values_list.uncast());
+                }
 
-                    {
-                        let existing_alias = range_subselect
-                            .alias()
-                            .expect("alias was checked in the parser");
+                {
+                    let existing_alias = range_subselect
+                        .alias()
+                        .expect("alias was checked in the parser");
 
-                        if existing_alias.colnames().is_empty() {
-                            // If no explicit columns are provided in the alias, then we added them
-                            // to correctly resolve against alias columns referenced in the outer query.
+                    if existing_alias.colnames().is_empty() {
+                        // If no explicit columns are provided in the alias, then we added them
+                        // to correctly resolve against alias columns referenced in the outer query.
 
-                            let new_alias =
-                                Self::make_alias(mem, existing_alias, row_description.clone());
+                        let new_alias =
+                            Self::make_alias(mem, existing_alias, row_description.clone());
 
-                            range_subselect.as_mut().set_alias(new_alias.as_option());
-                        }
+                        range_subselect.as_mut().set_alias(new_alias.as_option());
                     }
-
-                    node.replace(range_subselect.uncast());
-                    None
                 }
-                _ => Some(node),
-            });
 
-            let new_from_clause = mem.make_list(&[new_from_clause]);
-
-            match owned_raw_stmt.as_mut().stmt_mut() {
-                NodeMut::SelectStmt(mut stmt) => stmt.set_from_clause(new_from_clause),
-                NodeMut::UpdateStmt(mut stmt) => stmt.set_from_clause(new_from_clause),
-                _ => unreachable!("The parser only allows through SELECT and UPDATE"),
-            };
-
-            if data_rows.is_empty() {
-                // If rows are empty, now set WHERE to false to cancel out the NULL row we created
-                // with `explicit_values_list` to properly return type OIDs in the RowDescription
-                // when we have an empty derived table response.
-
-                match owned_raw_stmt.as_mut().stmt_mut() {
-                    NodeMut::SelectStmt(mut stmt) => stmt.set_where_clause(
-                        mem.make_a_const(pg_raw_parse::ConstValue::Boolean(false))
-                            .uncast(),
-                    ),
-                    NodeMut::UpdateStmt(mut stmt) => stmt.set_where_clause(
-                        mem.make_a_const(pg_raw_parse::ConstValue::Boolean(false))
-                            .uncast(),
-                    ),
-                    _ => unreachable!("The parser only allows through SELECT and UPDATE"),
-                }
+                node.replace(range_subselect.uncast());
+                None
             }
-
-            owned_raw_stmt
+            _ => Some(node),
         });
 
-        let new_sql = {
-            let new_sql = deparse(new_statement.deref()).unwrap();
-            let new_sql = new_sql.as_str();
+        let new_from_clause = mem.make_list(&[new_from_clause]);
 
-            new_sql.to_owned()
+        match statement.as_mut() {
+            NodeMut::SelectStmt(mut stmt) => stmt.set_from_clause(new_from_clause),
+            NodeMut::UpdateStmt(mut stmt) => stmt.set_from_clause(new_from_clause),
+            _ => unreachable!("The parser only allows through SELECT and UPDATE"),
         };
 
-        Ok(Some(new_sql))
+        if data_rows.is_empty() {
+            // If rows are empty, now set WHERE to false to cancel out the NULL row we created
+            // with `explicit_values_list` to properly return type OIDs in the RowDescription
+            // when we have an empty derived table response.
+
+            match statement.as_mut() {
+                NodeMut::SelectStmt(mut stmt) => stmt.set_where_clause(
+                    mem.make_a_const(pg_raw_parse::ConstValue::Boolean(false))
+                        .uncast(),
+                ),
+                NodeMut::UpdateStmt(mut stmt) => stmt.set_where_clause(
+                    mem.make_a_const(pg_raw_parse::ConstValue::Boolean(false))
+                        .uncast(),
+                ),
+                _ => unreachable!("The parser only allows through SELECT and UPDATE"),
+            }
+        }
+
+        statement
     }
 
     /// Construct a new `SelectStmt` to be set as the `subselect` of our new `RangeSubselect` node,
@@ -121,7 +106,7 @@ impl<'a> SubqueryMulti<'a> {
     ///
     /// TODO(jkaczman): we need to account for too much data being present, and split
     /// apart the query into multiple when that occurs. This TODO also exists elsewhere for subqueries.
-    fn new_select_with_explicit_values<'mem>(
+    pub(crate) fn new_select_with_explicit_values<'mem>(
         mem: MemoryToken<'mem>,
         row_description: RowDescription,
         data_rows: &[DataRow],
@@ -223,7 +208,7 @@ impl<'a> SubqueryMulti<'a> {
     /// Assemble a new alias to be used within the `RangeSubselect` node we're re-writing,
     /// to include the colnames if they're not already present. This allows for explicit
     /// <ALIAS>.<col> use within the outer-query (whom references the `RangeSubselect` results)
-    fn make_alias<'mem>(
+    pub(crate) fn make_alias<'mem>(
         mem: MemoryToken<'mem>,
         existing_alias: &Alias,
         row_description: RowDescription,

pgdog/src/frontend/client/query_engine/multi_step/subquery/insert_select.rs

@@ -1,65 +1,55 @@
-use std::ops::Deref;
-
-use pg_raw_parse::{Node, deparse, make::owned, nodes::SelectStmt};
+use pg_raw_parse::{
+    Node, NodeMut,
+    make::{MemoryToken, Unique},
+    nodes::{SelectStmt, WithClause},
+};
 
 use crate::{
-    frontend::client::query_engine::multi_step::{Error, SubqueryMulti},
-    frontend::router::parser::Error as ParserError,
+    frontend::client::query_engine::multi_step::SubqueryMulti,
     net::{DataRow, Decoder, RowDescription},
 };
 
 impl<'a> SubqueryMulti<'a> {
     /// Re-writes the top-level INSERT (with SELECT FROM subquery) with the evaluated subquery's RowDescription and DataRows,
     /// and returns the re-written SQL.
-    pub fn handle_insert_select(
+    pub fn handle_insert_select<'mem>(
         &self,
-        stmt: Node<'a>,
+        mem: MemoryToken<'mem>,
+        mut stmt: Unique<'mem, Node<'mem>>,
         row_description: Option<&RowDescription>,
-        data_rows: &Vec<DataRow>,
-    ) -> Result<Option<String>, Error> {
-        let Node::InsertStmt(insert) = stmt else {
-            return Ok(None);
+        data_rows: &[DataRow],
+    ) -> Unique<'mem, Node<'mem>> {
+        let NodeMut::InsertStmt(mut insert) = stmt.as_mut() else {
+            return stmt;
         };
 
         // Replace with SelectStmt with new VALUES
-        let new_insert = owned(|mem| {
-            let mut new_select = mem.make_node::<SelectStmt>();
-            let mut new_select_mut = new_select.as_mut();
-
-            let mut decoder = Decoder::default();
-            decoder.set_row_description(row_description.unwrap().clone());
-
-            // Build the VALUES list from the DataRows
-            let mut lists = Vec::new();
-            for row in data_rows {
-                let mut elems = Vec::new();
+        let mut new_select = mem.make_node::<SelectStmt>();
+        let mut new_select_mut = new_select.as_mut();
 
-                for col in 0..row.len() {
-                    let datum = row.get_column(col, &decoder).unwrap().unwrap().value;
+        let mut decoder = Decoder::default();
+        decoder.set_row_description(row_description.unwrap().clone());
 
-                    elems.push(Self::convert_datum_to_a_const(datum, mem).uncast());
-                }
+        // Build the VALUES list from the DataRows
+        let mut lists = Vec::new();
+        for row in data_rows {
+            let mut elems = Vec::new();
 
-                let tuple = mem.make_list(elems.as_ref()).uncast();
-                lists.push(tuple);
+            for col in 0..row.len() {
+                let datum = row.get_column(col, &decoder).unwrap().unwrap().value;
+                elems.push(Self::convert_datum_to_a_const(datum, mem).uncast());
             }
 
-            let values_lists = mem.make_list(lists.as_ref());
-            new_select_mut.set_values_lists(values_lists);
+            let tuple = mem.make_list(elems.as_ref()).uncast();
+            lists.push(tuple);
+        }
 
-            let mut cloned_insert = mem.make_unique(insert);
-            cloned_insert.as_mut().set_select_stmt(new_select.uncast());
+        let values_lists = mem.make_list(lists.as_ref());
+        new_select_mut.set_values_lists(values_lists);
 
-            cloned_insert
-        });
-
-        let new_sql: String = {
-            let new_insert_sql = deparse(new_insert.deref()).map_err(ParserError::from)?;
-            let new_insert_sql = new_insert_sql.as_str();
-
-            new_insert_sql.to_string()
-        };
+        insert.set_select_stmt(new_select.uncast());
+        insert.set_with_clause(mem.make_unique(None::<&WithClause>));
 
-        Ok(Some(new_sql))
+        stmt
     }
 }