diff --git a/integration/rust/tests/integration/subqueries/tests/insert.sql b/integration/rust/tests/integration/subqueries/tests/insert.sql index 7bcd4195d..b9d76d8e0 100644 --- a/integration/rust/tests/integration/subqueries/tests/insert.sql +++ b/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 diff --git a/integration/rust/tests/integration/subqueries/tests/select.sql b/integration/rust/tests/integration/subqueries/tests/select.sql index d11713f93..6ca917a5d 100644 --- a/integration/rust/tests/integration/subqueries/tests/select.sql +++ b/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 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 ) SELECT , (..) FROM c GROUP BY +-- 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; diff --git a/integration/rust/tests/integration/subqueries/todo.sql b/integration/rust/tests/integration/subqueries/todo.sql index 8e2859cc8..9aac88539 100644 --- a/integration/rust/tests/integration/subqueries/todo.sql +++ b/integration/rust/tests/integration/subqueries/todo.sql @@ -95,12 +95,12 @@ -- [DONE] [A1 A2 A3 A4 A5 A6 O9 Q2 Q3] [update] UPDATE SET .. FROM () s WHERE .. -- Plain CTEs --- [A1 A2 A3 A4 A5 A6] [select] WITH c AS (..) SELECT .. WHERE 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 ) SELECT , (..) FROM c GROUP BY --- [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 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 ) SELECT , (..) FROM c GROUP BY +-- [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 IN () EXISTS () diff --git a/pgdog/src/frontend/client/query_engine/multi_step/error.rs b/pgdog/src/frontend/client/query_engine/multi_step/error.rs index 37bf3d50d..8f82c2db4 100644 --- a/pgdog/src/frontend/client/query_engine/multi_step/error.rs +++ b/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)] diff --git a/pgdog/src/frontend/client/query_engine/multi_step/subquery/from_clause.rs b/pgdog/src/frontend/client/query_engine/multi_step/subquery/from_clause.rs index 61dbf116e..d82cdbaf7 100644 --- a/pgdog/src/frontend/client/query_engine/multi_step/subquery/from_clause.rs +++ b/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, 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 /// . 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, diff --git a/pgdog/src/frontend/client/query_engine/multi_step/subquery/insert_select.rs b/pgdog/src/frontend/client/query_engine/multi_step/subquery/insert_select.rs index a52ee99e8..c27a8649e 100644 --- a/pgdog/src/frontend/client/query_engine/multi_step/subquery/insert_select.rs +++ b/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, - ) -> Result, 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::(); - 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::(); + 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 } } diff --git a/pgdog/src/frontend/client/query_engine/multi_step/subquery/mod.rs b/pgdog/src/frontend/client/query_engine/multi_step/subquery/mod.rs index 9025b5a5d..488b2451a 100644 --- a/pgdog/src/frontend/client/query_engine/multi_step/subquery/mod.rs +++ b/pgdog/src/frontend/client/query_engine/multi_step/subquery/mod.rs @@ -2,9 +2,11 @@ //! //! Note: This is still in development. Do not use in production. //! Please do not work on without discussing with @jkaczman first. + use pg_raw_parse::{ - ConstValue, Node, - make::{MemoryToken, Unique}, + ConstValue, Node, Owned, deparse, + make::{MemoryToken, Unique, owned}, + nodes::{CommonTableExpr, RawStmt, SelectStmt}, }; use pgdog_postgres_types::Datum; use rust_decimal::prelude::ToPrimitive; @@ -14,40 +16,105 @@ mod from_clause; mod insert_select; mod request; mod where_clause; +mod with_clause; use super::Error; use crate::{ frontend::{ ClientRequest, client::query_engine::{QueryEngine, QueryEngineContext}, - router::parser::rewrite::statement::insert::{Subquery, SubqueryType}, + router::parser::rewrite::statement::{ + insert::Subquery, + subquery::{ + from_clause::parse_from_clause, insert_select::parse_insert_select, + where_clause::parse_where_clause, with_clause::parse_with_clause, + }, + }, }, - net::{Bind, DataRow, RowDescription}, + net::{Bind, DataRow, Decoder, ProtocolMessage, RowDescription}, }; /// Evaluates the Subquery and injects the results into the top-level statement. -pub(crate) struct SubqueryMulti<'a> { - pub(super) subquery: &'a Subquery, - pub(super) engine: &'a mut QueryEngine, +pub(crate) struct SubqueryMulti<'ast> { + pub(super) subqueries: Vec, // Store the very outer statement's Bind message. // As in, if we're 3 levels deep, it'd be the 1st level's Bind. // // This allows us to find the relevant parameters for each subquery level // within the client's Bind message. - pub(super) top_level_bind: Option<&'a Bind>, + pub(super) top_level_bind: Option<&'ast Bind>, + + // Store the results here, so that can reference previously executed subqueries (e.g., if we need to + // use earlier resolved CTEs to resolve a CTE occurring after those) + pub(super) subquery_results: Vec<(Vec, Option)>, + + // Was the original statement from the Client sent from simple or extended protocol? + // This determines how we need to send the response back. + pub(super) is_extended_protocol: bool, + + // The Ast of the top-level statement for us to reference for subquery planning + pub(super) top_level_statement: Node<'ast>, } -impl<'a> SubqueryMulti<'a> { - pub(crate) fn new( - engine: &'a mut QueryEngine, - subquery: &'a Subquery, - top_level_bind: Option<&'a Bind>, - ) -> Self { +impl<'ast> SubqueryMulti<'ast> { + pub(crate) fn new(subqueries: Vec, client_request: &'ast ClientRequest) -> Self { + let ast = client_request + .ast + .as_ref() + .expect("The parser had to have been on to reach this code branch"); + + let top_level_statement: Node<'ast> = ast + .ast + .stmts() + .next() + .expect("The parser had to have been on to reach this code branch"); + + // TODO(jkaczman): There's likely a different way for us to determine if this is extended protocol or not, + // not to mention... this is a redundant loop (btwn. is_extended_protocol and top_level_bind) + let is_extended_protocol = client_request + .messages + .iter() + .find(|message| { + matches!( + message, + ProtocolMessage::Parse(_) | ProtocolMessage::Bind(_) + ) + }) + .is_some(); + + let top_level_bind: Option<&Bind> = if is_extended_protocol { + client_request.messages.iter().find_map(|message| { + if let ProtocolMessage::Bind(bind) = message { + Some(bind) + } else { + None + } + }) + } else { + None + }; + Self { - subquery, - engine, + subqueries, top_level_bind, + subquery_results: Vec::new(), + is_extended_protocol, + top_level_statement, + } + } + + pub fn new_sub_engine<'b>( + parent_engine: &'b SubqueryMulti<'ast>, + new_parent_ast: Node<'ast>, + subqueries: Vec, + ) -> Self { + Self { + subqueries, + top_level_bind: parent_engine.top_level_bind, + subquery_results: parent_engine.subquery_results.clone(), + is_extended_protocol: parent_engine.is_extended_protocol, + top_level_statement: new_parent_ast, } } @@ -57,18 +124,23 @@ impl<'a> SubqueryMulti<'a> { /// to the Client. pub(crate) async fn execute( &mut self, + engine: &mut QueryEngine, context: &mut QueryEngineContext<'_>, - client_request: &ClientRequest, ) -> Result<(), Error> { // Send the inner-subquery and re-write the outer statement with its response. - let rewritten_outer_statement: String = self - .evaluate_and_rewrite_outer_statement(context, client_request) + let rewritten_outer_statement = self + .evaluate_and_rewrite_outer_statement(engine, context) .await?; + let rewritten_outer_statement = deparse(rewritten_outer_statement.stmt())? + .as_str() + .to_string(); + // Send the re-written statement self.send_query_string( - client_request, - rewritten_outer_statement, + engine, + self.top_level_statement, + rewritten_outer_statement.to_string(), context, // `is_subquery` = false will send the response right back to the client false, @@ -83,60 +155,240 @@ impl<'a> SubqueryMulti<'a> { /// Returns the SQL of the re-written outer statement. pub(crate) async fn evaluate_and_rewrite_outer_statement( &mut self, + engine: &mut QueryEngine, context: &mut QueryEngineContext<'_>, - client_request: &ClientRequest, - ) -> Result { - // Send the subquery that we parsed through the QueryEngine, and fetch its response back - let (data_rows, row_description) = self - .send_query_string( - client_request, - self.subquery.inner_stmt.clone(), - context, - // 'is_subquery' will allow us to get the response back - true, - ) - .await? - .expect("send_query_string is only None when `is_subquery` is false"); + ) -> Result, Error> { + // TODO(jkaczman): we could execute some of these in parallel (in groups) e.g., if CTEs are not referencing earlier ones - // Re-write outer statement with the response - let new_sql: Option = self.rewrite_outer_statement_with_subquery_response( - client_request, - row_description.as_ref(), - &data_rows, - )?; + let current_statement: Node<'ast> = self.top_level_statement; + + let subqueries = self.subqueries.clone(); + for subquery in &subqueries { + // Send the subquery that we parsed through the QueryEngine, and fetch its response back + + let rewritten_subquery_ast: Owned = + self.parse_and_rewrite_subquery_ast(subquery, current_statement)?; + + let subquery_sql = deparse(&*rewritten_subquery_ast)?.as_str().to_string(); + + let results = self + .send_query_string( + engine, + Node::SelectStmt(&rewritten_subquery_ast), + subquery_sql.to_string(), + context, + // 'is_subquery' will allow us to get the response back + true, + ) + .await? + .expect("send_query_string is only None when `is_subquery` is false"); + + self.subquery_results.push(results); + } - let Some(new_sql) = new_sql else { + if self.subquery_results.is_empty() { unreachable!( "Subqueries: the parser check for a re-writable subquery should match our logic here." ) + } + + // Re-write outer statement with the response + // TODO(jkaczman): have this return the AST as well ig + let new_ast: Owned = self.rewrite_outer_statement_with_subquery_response()?; + Ok(new_ast) + } + + fn rewrite_subquery_ast_with_earlier_subquery_responses( + &self, + subquery: &Subquery, + subquery_ast: Node<'ast>, + ) -> Owned { + let Node::SelectStmt(select) = subquery_ast else { + unreachable!("We only allow subselects to be subqueries"); }; - Ok(new_sql) + owned(|mem| { + let mut subquery_ast = mem.make_unique(select); + + if matches!(subquery, Subquery::InsertSelect) + && let Node::InsertStmt(insert) = self.top_level_statement + && let Some(with_clause) = insert.with_clause() + { + subquery_ast + .as_mut() + .set_with_clause(mem.make_unique(with_clause).as_option()); + } + + let mut previously_executed_ctes: Vec> = Vec::new(); + + for (i, (rows, row_description)) in self.subquery_results.iter().enumerate() { + let subquery = self + .subqueries + .get(i) + .expect("Subquery 'i' must have existed to create the matching result"); + + match &subquery { + Subquery::Cte(alias) | Subquery::WhereClauseCte(alias) => { + let row_description = row_description + .as_ref() + .expect("RowDescription must be Some(..) in order to reach here"); + + let mut decoder = Decoder::default(); + decoder.set_row_description(row_description.clone()); + + // We need to set this within a WITH clause of the subselect, + // so that the subselect can reference previously ran CTEs + + let mut new_cte_subselect = mem.make_node::(); + let mut values_lists = Vec::new(); + + for row in rows { + let mut elems = Vec::new(); + + 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 tuple = mem.make_list(elems.as_ref()).uncast(); + values_lists.push(tuple) + } + + let values_lists = mem.make_list(values_lists.as_ref()); + new_cte_subselect.as_mut().set_values_lists(values_lists); + + // Need to form the alias for the CTE to be referenced within the other CTE, + // so that it can reference . + let mut alias_col_names = Vec::new(); + for i in 0..row_description.len() { + alias_col_names.push( + mem.make_string(Some( + row_description.fields.get(i).unwrap().name.as_str(), + )) + .uncast(), + ); + } + + let new_cte = mem.make_common_table_expr( + alias, + mem.make_list(alias_col_names.as_slice()), + new_cte_subselect.uncast(), + ); + + previously_executed_ctes.push(new_cte); + } + + // For now, we'll only consider re-writing if it was previously a CTE. + _ => continue, + } + } + + if !previously_executed_ctes.is_empty() { + let new_with_clause = + mem.make_with_clause(mem.make_list(previously_executed_ctes.as_slice()), false); + + subquery_ast + .as_mut() + .set_with_clause(new_with_clause.as_option()); + } + + subquery_ast + }) } - fn rewrite_outer_statement_with_subquery_response( + fn parse_and_rewrite_subquery_ast( &self, - client_request: &ClientRequest, + subquery: &Subquery, + entire_ast: Node<'ast>, + ) -> Result, Error> { + let subquery_ast: Node<'ast> = self.parse_out_relevant_query_to_send(subquery, entire_ast); + + // We might have to reference earlier statements and perform a re-write. E.g., if we have CTE b referencing CTE a, then we need to include the re-written CTE (w/ VALUES) here + // Only if it's actually relevant, but we need to either (a) figure that out somehow or (b) always include earlier ones + + let rewritten_subquery_ast = + self.rewrite_subquery_ast_with_earlier_subquery_responses(subquery, subquery_ast); + + Ok(rewritten_subquery_ast) + } + + /// This is the relevant part of what we need to send, e.g., the subselect of a specific CTE. + fn parse_out_relevant_query_to_send( + &self, + subquery: &Subquery, + entire_ast: Node<'ast>, + ) -> Node<'ast> { + match &subquery { + Subquery::FromClause => { + parse_from_clause(entire_ast) + } + Subquery::WhereClause => { + parse_where_clause(entire_ast) + } + Subquery::Cte(alias) => { + parse_with_clause(entire_ast, alias) + } + Subquery::InsertSelect => { + parse_insert_select(entire_ast) + } + Subquery::WhereClauseCte(alias) => { + parse_where_clause(entire_ast).and_then(|subselect| parse_with_clause(subselect, alias)) + } + }.expect("We already ran these methods in the parser originally. No reason they shouldn't pass here.") + } + + fn rewrite_outer_statement_with_subquery_response(&self) -> Result, Error> { + // TODO(jkaczman): if we're always passing these three, might as well make a new type or + // hold these in the existing one + + let new_statement: Owned = owned(|mem| { + let mut statement: Unique<'_, Node<'_>> = mem.make_unique(self.top_level_statement); + + // TODO(jkaczman): We'll either have to save the AST of the re-written node that we're + // constantly rewriting in this outer loop, or we can save it in the struct itself + // for easy reference.. however that'd mean we need a mutable reference to the struct, + // could be tricky when we try to parallelize this. + for (subquery, (data_rows, row_description)) in + self.subqueries.iter().zip(self.subquery_results.iter()) + { + statement = self.rewrite_with_subquery_response( + mem, + statement, + subquery, + row_description.as_ref(), + data_rows.as_slice(), + ); + } + + mem.make_raw_stmt(statement) + }); + + Ok(new_statement) + } + + fn rewrite_with_subquery_response<'mem>( + &self, + mem: MemoryToken<'mem>, + stmt: Unique<'mem, Node<'mem>>, + subquery: &Subquery, row_description: Option<&RowDescription>, - data_rows: &Vec, - ) -> Result, Error> { - if let Some(ast) = client_request.ast.as_ref() - && let Some(stmt) = ast.ast.stmts().next() - { - return match self.subquery.identifier { - SubqueryType::FromClause => { - self.handle_from_clause(stmt, row_description, data_rows) - } - SubqueryType::WhereClause => { - self.handle_where_clause(stmt, row_description, data_rows) - } - SubqueryType::InsertSelect => { - self.handle_insert_select(stmt, row_description, data_rows) - } - }; + data_rows: &[DataRow], + ) -> Unique<'mem, Node<'mem>> { + match subquery { + Subquery::FromClause => self.handle_from_clause(mem, stmt, row_description, data_rows), + Subquery::WhereClause => { + self.handle_where_clause(mem, stmt, row_description, data_rows) + } + Subquery::InsertSelect => { + self.handle_insert_select(mem, stmt, row_description, data_rows) + } + Subquery::WhereClauseCte(alias) => { + self.handle_where_clause_cte(alias, mem, stmt, row_description, data_rows) + } + Subquery::Cte(alias) => { + self.handle_with_clause(alias, mem, stmt, row_description, data_rows) + } } - - Ok(None) } /// Helper method to turn `Datum` (from DataRow response) into an A_Const node we can inject into the subquery diff --git a/pgdog/src/frontend/client/query_engine/multi_step/subquery/request.rs b/pgdog/src/frontend/client/query_engine/multi_step/subquery/request.rs index bc170bbe2..4a9606053 100644 --- a/pgdog/src/frontend/client/query_engine/multi_step/subquery/request.rs +++ b/pgdog/src/frontend/client/query_engine/multi_step/subquery/request.rs @@ -1,10 +1,12 @@ //! Sending and receiving logic for subqueries +use pg_raw_parse::{Node, deparse}; + use crate::{ frontend::{ ClientRequest, client::query_engine::{ - QueryEngineContext, + QueryEngine, QueryEngineContext, multi_step::{self, Error, SubqueryMulti}, }, router::parser::{ @@ -20,26 +22,25 @@ impl<'a> SubqueryMulti<'a> { /// `is_subquery` is true. Else, return back directly to the client. pub(crate) async fn send_query_string( &mut self, - original_request: &ClientRequest, + engine: &mut QueryEngine, + query_ast: Node<'_>, query_string: String, context: &mut QueryEngineContext<'_>, is_subquery: bool, ) -> Result, Option)>, Error> { let mut client_request: ClientRequest = self - .form_client_request(query_string, original_request, context) + .form_client_request(engine, query_string, context) .await?; - multi_step::update::route(self.engine, &mut client_request, context)?; + multi_step::update::route(engine, &mut client_request, context)?; if is_subquery { - let rewrite_result: Option = self - .engine - .rewrite_request(context, &mut client_request) - .await?; + let rewrite_result: Option = + engine.rewrite_request(context, &mut client_request).await?; match rewrite_result { Some(rewrite_result @ RewriteResult::InPlace { .. }) => { - let cluster = self.engine.backend.cluster()?; + let cluster = engine.backend.cluster()?; projection::finalize_after_route( &mut client_request, &cluster.schema(), @@ -47,29 +48,32 @@ impl<'a> SubqueryMulti<'a> { )?; rewrite_result.apply_after_route(&mut client_request)?; - Ok(Some(self.send_request(&mut client_request).await?)) + Ok(Some(multi_step::fetch_rows(engine, &client_request).await?)) } - Some(RewriteResult::Subquery(subquery)) => { + Some(RewriteResult::Subquery(subqueries)) => { // Pass the `client_request` we just made back through multi_step/subquery.rs's internal // SQL response method (return new "outer" SQL to us instead of a response to client). // In this case, the "outer" SQL is the evaluated subquery within this statement. // // Then, we send the re-written subquery here, and return the relevant DataRows / RD - let mut subquery_multi = SubqueryMulti { - subquery: &subquery, - engine: self.engine, - top_level_bind: self.top_level_bind, + + let rewritten_subquery_ast = { + let mut subquery_multi = + SubqueryMulti::new_sub_engine(self, query_ast, subqueries); + + Box::pin( + subquery_multi.evaluate_and_rewrite_outer_statement(engine, context), + ) + .await? }; - let rewritten_subquery_sql = Box::pin( - subquery_multi - .evaluate_and_rewrite_outer_statement(context, &client_request), - ) - .await?; + let rewritten_subquery_sql = + deparse(rewritten_subquery_ast.stmt())?.as_str().to_string(); Box::pin(self.send_query_string( - original_request, + engine, + rewritten_subquery_ast.stmt(), rewritten_subquery_sql, context, true, @@ -86,16 +90,12 @@ impl<'a> SubqueryMulti<'a> { // extracting a full helper out of QueryEngine that doesn't return the response // to the Client when a flag is false? - let rewrite_result = self - .engine - .rewrite_request(context, &mut client_request) - .await?; + let rewrite_result = engine.rewrite_request(context, &mut client_request).await?; // Need this to set route for stuff like INSERT split, where it relies on // "rewrite_result.apply_after_route(client_request)?;" // which sets the route of the outer statement if all tuples converge to the same shard. - if !self - .engine + if !engine .route_query(context, &mut client_request, rewrite_result.as_ref()) .await? { @@ -108,12 +108,9 @@ impl<'a> SubqueryMulti<'a> { // enabled, then INSERTs will still fail with // INSERT SELECT FROM. Maybe we override INSERT split // if subqueries are enabled in that case - Box::pin( - self.engine - .execute(context, &mut client_request, rewrite_result), - ) - .await - .unwrap(); + Box::pin(engine.execute(context, &mut client_request, rewrite_result)) + .await + .unwrap(); Ok(None) } @@ -121,22 +118,11 @@ impl<'a> SubqueryMulti<'a> { async fn form_client_request( &self, + engine: &mut QueryEngine, mut query_string: String, - original_request: &ClientRequest, context: &mut QueryEngineContext<'_>, ) -> Result { - let is_extended_protocol = original_request - .messages - .iter() - .find(|message| { - matches!( - message, - ProtocolMessage::Parse(_) | ProtocolMessage::Bind(_) - ) - }) - .is_some(); - - let mut client_request = if is_extended_protocol { + let mut client_request = if self.is_extended_protocol { let mut bind = Bind::default(); if let Some(top_level_bind) = self.top_level_bind { let mut actual_param_index = 1; @@ -179,7 +165,7 @@ impl<'a> SubqueryMulti<'a> { ClientRequest::from(vec![ProtocolMessage::Query(Query::new(&query_string))]) }; - let cluster = self.engine.backend.cluster()?; + let cluster = engine.backend.cluster()?; let ast_ctx = AstContext::from_cluster(cluster, context.params, context.timestamps()); let query = crate::frontend::BufferedQuery::Query(Query::new(&query_string)); @@ -191,12 +177,4 @@ impl<'a> SubqueryMulti<'a> { Ok(client_request) } - - /// Send `client_request` and return the rows and RowDescription (if present) back. - pub(crate) async fn send_request( - &mut self, - client_request: &mut ClientRequest, - ) -> Result<(Vec, Option), Error> { - multi_step::fetch_rows(self.engine, client_request).await - } } diff --git a/pgdog/src/frontend/client/query_engine/multi_step/subquery/where_clause.rs b/pgdog/src/frontend/client/query_engine/multi_step/subquery/where_clause.rs index 8436cc076..7b9b2e7e3 100644 --- a/pgdog/src/frontend/client/query_engine/multi_step/subquery/where_clause.rs +++ b/pgdog/src/frontend/client/query_engine/multi_step/subquery/where_clause.rs @@ -1,10 +1,8 @@ //! Handles re-writing the WHERE clause of a Select / Update / Delete statement when a subquery is present. -use std::ops::Deref; - use pg_raw_parse::{ - ConstValue, Node, NodeMut, deparse, - make::{MemoryToken, Unique, owned}, + ConstValue, Node, NodeMut, + make::{MemoryToken, Unique}, nodes::SubLink, raw::SubLinkType, transform, @@ -13,84 +11,67 @@ use pgdog_postgres_types::Datum; use tracing::debug; use crate::{ - frontend::{ - client::query_engine::multi_step::{Error, SubqueryMulti}, - router::parser::{Error as ParserError, rewrite}, - }, + frontend::{client::query_engine::multi_step::SubqueryMulti, router::parser::rewrite}, net::{DataRow, Decoder, RowDescription}, }; impl<'a> SubqueryMulti<'a> { - pub(crate) fn handle_where_clause( + pub(crate) fn handle_where_clause<'mem>( &self, - statement: Node<'a>, + mem: MemoryToken<'mem>, + mut statement: Unique<'mem, Node<'mem>>, row_description: Option<&RowDescription>, - data_rows: &Vec, - ) -> Result, Error> { + data_rows: &[DataRow], + ) -> Unique<'mem, Node<'mem>> { let Some((sublink, lexpr)) = - rewrite::statement::subquery::where_clause::has_sub_link_in_where(statement) + rewrite::statement::subquery::where_clause::has_sub_link_in_where(statement.as_ref()) else { - return Ok(None); + return statement; }; - let new_stmt = owned(|mem| { - let owned_stmt = mem.make_unique(statement); - let mut owned_raw_stmt = mem.make_raw_stmt(owned_stmt); - let mut mutable_raw_stmt = owned_raw_stmt.as_mut(); - - // Fetch the where clause. - let mut where_clause = mem.make_unique(match mutable_raw_stmt.stmt() { - Node::SelectStmt(stmt) => stmt.where_clause(), - Node::UpdateStmt(stmt) => stmt.where_clause(), - Node::DeleteStmt(stmt) => stmt.where_clause(), - _ => unreachable!("Parser should prevent other statements from reaching here."), - }); - - // Replace the SubLink node present with our own re-written version. - transform::transform(&mut where_clause, |node| match &*node { - NodeMut::SubLink(_) => { - let sublink_replacement = self.rewrite_sublink_in_where_clause( - sublink, - data_rows, - row_description.as_ref().unwrap(), - mem, - lexpr, - ); - - node.replace(sublink_replacement); - - // We shouldn't recurse here. Even if we have nested subqueries, - // that'll be handled later through routing through the QueryEngine again. - None - } - _ => Some(node), - }); - - // Put back the modified where clause. - match mutable_raw_stmt.stmt_mut() { - NodeMut::SelectStmt(mut stmt) => stmt.set_where_clause(where_clause), - NodeMut::UpdateStmt(mut stmt) => stmt.set_where_clause(where_clause), - NodeMut::DeleteStmt(mut stmt) => stmt.set_where_clause(where_clause), - _ => unreachable!("Parser should prevent other statements from reaching here."), - }; - - owned_raw_stmt + // Fetch the where clause. + let mut where_clause = mem.make_unique(match statement.as_ref() { + Node::SelectStmt(stmt) => stmt.where_clause(), + Node::UpdateStmt(stmt) => stmt.where_clause(), + Node::DeleteStmt(stmt) => stmt.where_clause(), + _ => unreachable!("Parser should prevent other statements from reaching here."), }); - let new_sql: String = { - let new_insert_sql = deparse(new_stmt.deref()).map_err(ParserError::from)?; - let new_insert_sql = new_insert_sql.as_str(); + // Replace the SubLink node present with our own re-written version. + transform::transform(&mut where_clause, |node| match &*node { + NodeMut::SubLink(_) => { + let sublink_replacement = self.rewrite_sublink_in_where_clause( + sublink, + data_rows, + row_description.as_ref().unwrap(), + mem, + lexpr, + ); + + node.replace(sublink_replacement); + + // We shouldn't recurse here. Even if we have nested subqueries, + // that'll be handled later through routing through the QueryEngine again. + None + } + _ => Some(node), + }); - new_insert_sql.to_string() + // Put back the modified where clause. + match statement.as_mut() { + NodeMut::SelectStmt(mut stmt) => stmt.set_where_clause(where_clause), + NodeMut::UpdateStmt(mut stmt) => stmt.set_where_clause(where_clause), + NodeMut::DeleteStmt(mut stmt) => stmt.set_where_clause(where_clause), + _ => unreachable!("Parser should prevent other statements from reaching here."), }; - Ok(Some(new_sql)) + statement } fn rewrite_sublink_in_where_clause<'mem>( &self, sublink: &SubLink, - data_rows: &Vec, + data_rows: &[DataRow], row_description: &RowDescription, mem: MemoryToken<'mem>, lexpr: Node<'_>, diff --git a/pgdog/src/frontend/client/query_engine/multi_step/subquery/with_clause.rs b/pgdog/src/frontend/client/query_engine/multi_step/subquery/with_clause.rs new file mode 100644 index 000000000..ba329ff90 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/multi_step/subquery/with_clause.rs @@ -0,0 +1,156 @@ +use pg_raw_parse::{ + Node, NodeMut, + make::{MemoryToken, Unique}, + transform, +}; + +use crate::{ + frontend::{ + client::query_engine::multi_step::SubqueryMulti, + router::parser::rewrite::statement::subquery::where_clause::has_sub_link_in_where, + }, + net::{DataRow, RowDescription}, +}; + +impl<'a> SubqueryMulti<'a> { + pub fn handle_where_clause_cte<'mem>( + &self, + alias: &String, + mem: MemoryToken<'mem>, + mut statement: Unique<'mem, Node<'mem>>, + row_description: Option<&'a RowDescription>, + data_rows: &'a [DataRow], + ) -> Unique<'mem, Node<'mem>> { + let Some((sublink, _)) = has_sub_link_in_where(statement.as_ref()) else { + return statement; + }; + + let rewritten_subselect = self.handle_with_clause( + alias, + mem, + mem.make_unique(sublink.subselect()), + row_description, + data_rows, + ); + + let mut rewritten_sublink = mem.make_unique(sublink); + rewritten_sublink + .as_mut() + .set_subselect(rewritten_subselect); + + let mut rewritten_sublink = Some(rewritten_sublink.uncast()); + + let mut where_clause = mem.make_unique(match statement.as_ref() { + Node::SelectStmt(stmt) => stmt.where_clause(), + Node::UpdateStmt(stmt) => stmt.where_clause(), + Node::DeleteStmt(stmt) => stmt.where_clause(), + _ => unreachable!("Parser should prevent other statements from reaching here."), + }); + + transform::transform(&mut where_clause, |node| match &*node { + NodeMut::SubLink(_) => { + if let Some(rewritten_sublink) = rewritten_sublink.take() { + node.replace(rewritten_sublink); + } + + None + } + + _ => Some(node), + }); + + match statement.as_mut() { + NodeMut::SelectStmt(mut stmt) => stmt.set_where_clause(where_clause), + NodeMut::UpdateStmt(mut stmt) => stmt.set_where_clause(where_clause), + NodeMut::DeleteStmt(mut stmt) => stmt.set_where_clause(where_clause), + _ => unreachable!("Parser should prevent other statements from reaching here."), + }; + + statement + } + + pub fn handle_with_clause<'mem>( + &self, + alias: &String, + mem: MemoryToken<'mem>, + mut statement: Unique<'mem, Node<'mem>>, + row_description: Option<&'a RowDescription>, + data_rows: &'a [DataRow], + ) -> Unique<'mem, Node<'mem>> { + let Some(row_description) = row_description else { + return statement; + }; + + let with_clause = match statement.as_ref() { + Node::SelectStmt(select) => select.with_clause(), + Node::UpdateStmt(update) => update.with_clause(), + _ => unreachable!("Parser wouldn't have let through otherwise"), + }; + + let Some(with_clause) = with_clause else { + return statement; + }; + + // TODO(jkaczman): create shared structure between derived tables + // and CTEs as both are defining the columns explicitly in the alias, + // as well as re-writing to use explicit VALUES + + let mut new_with_clause = mem.make_unique(with_clause); + + let mut new_cte_list = Vec::new(); + for cte in with_clause.ctes() { + if let Some(name) = cte.ctename() + && !name.eq(alias) + { + // If we're not modifying this CTE, we still want it in our new CTE list + // Especially in the same order. + new_cte_list.push(mem.make_unique(cte)); + + continue; + } + + // TODO: This should be based on the list that's already + // fed to us. Although, what's the argument for parsing all + // these beforehand vs. just knowing we have a CTE? + // + // Any memory implications from storing it? Is it good + // as it's reusable? What about the length of the strings? + let new_cte_query = + Self::new_select_with_explicit_values(mem, row_description.clone(), data_rows); + + let mut alias_col_names = mem.make_unique(cte.aliascolnames()); + if alias_col_names.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. + for i in 0..row_description.len() { + alias_col_names.as_mut().push( + mem, + mem.make_string(Some(row_description.fields.get(i).unwrap().name.as_str())) + .uncast(), + ); + } + } + + let new_cte_node = mem.make_common_table_expr( + cte.ctename().unwrap(), + alias_col_names, + new_cte_query.uncast(), + ); + + new_cte_list.push(new_cte_node); + } + + new_with_clause + .as_mut() + .set_ctes(mem.make_list(new_cte_list.as_slice())); + + // Now we set the new WITH onto the statement + match statement.as_mut() { + NodeMut::SelectStmt(mut select) => select.set_with_clause(new_with_clause.as_option()), + NodeMut::UpdateStmt(mut update) => update.set_with_clause(new_with_clause.as_option()), + _ => unreachable!("Parser doesn't allow this."), + } + + statement + } +} diff --git a/pgdog/src/frontend/client/query_engine/query.rs b/pgdog/src/frontend/client/query_engine/query.rs index 9261c8983..0e3abc8ce 100644 --- a/pgdog/src/frontend/client/query_engine/query.rs +++ b/pgdog/src/frontend/client/query_engine/query.rs @@ -10,7 +10,7 @@ use crate::{ }, }, net::{ - Bind, DataRow, FromBytes, Message, Protocol, ProtocolMessage, Query, ReadyForQuery, + DataRow, FromBytes, Message, Protocol, ProtocolMessage, Query, ReadyForQuery, RowDescription, ToBytes, TransactionState, }, state::State, @@ -139,20 +139,9 @@ impl QueryEngine { // We get to here when detecting a Subquery in a top-level statement; this isn't recursed. // Thus, we should now fetch the Bind, if present, and pass it in (at the top-level), // so that we can supply the very top level's Bind to all sub-levels. - // - // TODO(jkaczman): don't need to do this if it's not extended protocol; do we have an easy way - // to find that out via a pre-computed flag? - let top_level_bind: Option<&Bind> = - client_request.messages.iter().find_map(|message| { - if let ProtocolMessage::Bind(bind) = message { - Some(bind) - } else { - None - } - }); - multi_step::SubqueryMulti::new(self, &subquery, top_level_bind) - .execute(context, client_request) + multi_step::SubqueryMulti::new(subquery, client_request) + .execute(self, context) .await?; } diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs b/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs index 67754d84f..19f01335c 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs @@ -10,27 +10,24 @@ use crate::net::{Bind, Parse, ProtocolMessage, Query}; use super::{Error, InsertSplitRewriteResult, RewritePlan, StatementRewrite}; #[derive(Debug, Clone)] -pub(crate) struct Subquery { - /// The inner-SQL of the outer statement (the Subquery's SQL) that we wish to execute - /// and re-write in our multi-step code. - pub(crate) inner_stmt: String, - - /// What kind of subquery is it? - /// E.g., FROM clause - pub(crate) identifier: SubqueryType, -} - -#[derive(Debug, Clone)] -pub(crate) enum SubqueryType { - // Determine if there's a RangeSubselect within a FROM clause of a SELECT / UPDATE statement. +pub(crate) enum Subquery { + /// Determine if there's a RangeSubselect within a FROM clause of a SELECT / UPDATE statement. FromClause, - // SELECT / UPDATE / DELETE with a subquery in a WHERE clause. + /// SELECT / UPDATE / DELETE with a subquery in a WHERE clause. WhereClause, - // Select subquery within an INSERT statement with a SELECT FROM - // e.g., INSERT INTO .. SELECT FROM generate_series(..) + /// Select subquery within an INSERT statement with a SELECT FROM + /// e.g., INSERT INTO .. SELECT FROM generate_series(..) InsertSelect, + + /// CTE present in query + /// e.g., + /// Inner param `String` is the alias + Cte(String), + + /// CTE present in WHERE clause + WhereClauseCte(String), } #[derive(Debug, Clone)] diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs index fd7657ab1..46a6262da 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs @@ -237,7 +237,7 @@ impl<'a> StatementRewrite<'a> { // Handle planning subqueries if config().config.rewrite.subqueries == RewriteMode::Rewrite { - plan.subquery = self.plan_subquery(stmt.stmt())?; + plan.subqueries = self.plan_subquery(stmt.stmt())?; } if let Node::UpdateStmt(stmt) = stmt.stmt() { diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs index 2e77d4005..1990dc15c 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs @@ -121,7 +121,7 @@ where /// /// Executed in order of fields in this struct. /// -#[derive(Default, Clone, Debug)] +#[derive(Default, Debug)] pub(crate) struct RewritePlan { /// One-based parameter indexes and ID sources in allocation order. /// Simple protocol records sequence calls here without using the indexes. @@ -141,7 +141,7 @@ pub(crate) struct RewritePlan { /// Executing an inner subquery separately to ensure proper routing /// Involves injecting the results of the inner query into the outer query - pub(crate) subquery: Option, + pub(crate) subqueries: Vec, /// Sharding key is being updated, we need to execute /// a multi-step plan. @@ -151,11 +151,11 @@ pub(crate) struct RewritePlan { pub(crate) offset: Option, } -#[derive(Debug, Clone)] +#[derive(Debug)] pub(crate) enum RewriteResult { InPlace { offset: Option }, InsertSplit(InsertSplitRewriteResult), - Subquery(Subquery), + Subquery(Vec), ShardingKeyUpdate(ShardingKeyUpdate), } @@ -402,10 +402,8 @@ impl RewritePlan { )); } - if let Some(subqueries) = &self.subquery - && request.is_executable() - { - return Ok(RewriteResult::Subquery(subqueries.clone())); + if !self.subqueries.is_empty() && request.is_executable() { + return Ok(RewriteResult::Subquery(self.subqueries.clone())); } Ok(RewriteResult::InPlace { diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/subquery/from_clause.rs b/pgdog/src/frontend/router/parser/rewrite/statement/subquery/from_clause.rs index e3ad40b36..1584d0970 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/subquery/from_clause.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/subquery/from_clause.rs @@ -1,53 +1,52 @@ -use pg_raw_parse::{Node, deparse}; +use pg_raw_parse::Node; use crate::frontend::router::parser::rewrite::statement::{ - insert::{Subquery, SubqueryType}, - subquery::SubqueryRewrite, + insert::Subquery, subquery::SubqueryRewrite, }; impl<'a, 'mem> SubqueryRewrite<'a, 'mem> { pub(crate) fn handle_subquery_in_from_clause(&self) -> Option { - let from_clause = match self.outer_stmt { - Node::UpdateStmt(update) => update.from_clause(), - Node::SelectStmt(select) => select.from_clause(), - _ => return None, - }; - - let from_clause = from_clause.first()?; - - match from_clause { - Node::RangeSubselect(range_subselect) => { - let subquery = range_subselect.subquery(); - if range_subselect.alias().is_some() - && let Node::SelectStmt(inner_subselect) = subquery - && (!inner_subselect.from_clause().is_empty() - || !matches!(inner_subselect.where_clause(), Node::None)) - { - // There's no `is_sharded` check here, because we could have - // a case where we need to execute a subquery that's NOT sharded, - // such as generate_series, which would resolve into a - // `ShardingKeyUpdate`, requiring coordination here to - // (1) resolve the actual VALUES and (2) split apart to he - // relevant shards to ensure that `ShardingKeyUpdate` is executed - // as intended. - // - // TODO(jkaczman): The above case does not work yet. ShardingKeyUpdate - // does not work as-is for this; this will require actual changes - // to that multi-step machinery to handle those cases. However, - // subqueries correctly fetches the rows, so we have that information already. - - let subquery_sql = deparse(subquery).unwrap(); - let subquery_sql = subquery_sql.as_str(); - - return Some(Subquery { - inner_stmt: subquery_sql.to_owned(), - identifier: SubqueryType::FromClause, - }); - } + parse_from_clause(self.outer_stmt) + .is_some() + .then_some(Subquery::FromClause) + } +} + +pub fn parse_from_clause<'a>(outer_stmt: Node<'a>) -> Option> { + let from_clause = match outer_stmt { + Node::UpdateStmt(update) => update.from_clause(), + Node::SelectStmt(select) => select.from_clause(), + _ => return None, + }; + + let from_clause = from_clause.first()?; + + match from_clause { + Node::RangeSubselect(range_subselect) => { + let subquery = range_subselect.subquery(); + if range_subselect.alias().is_some() + && let Node::SelectStmt(inner_subselect) = subquery + && inner_subselect.values_lists().is_empty() + { + // There's no `is_sharded` check here, because we could have + // a case where we need to execute a subquery that's NOT sharded, + // such as generate_series, which would resolve into a + // `ShardingKeyUpdate`, requiring coordination here to + // (1) resolve the actual VALUES and (2) split apart to he + // relevant shards to ensure that `ShardingKeyUpdate` is executed + // as intended. + // + // TODO(jkaczman): The above case does not work yet. ShardingKeyUpdate + // does not work as-is for this; this will require actual changes + // to that multi-step machinery to handle those cases. However, + // subqueries correctly fetches the rows, so we have that information already. + + return Some(Node::SelectStmt(inner_subselect)); } - _ => return None, } - None + _ => return None, } + + None } diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/subquery/insert_select.rs b/pgdog/src/frontend/router/parser/rewrite/statement/subquery/insert_select.rs index b9bab13fd..ed2194ead 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/subquery/insert_select.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/subquery/insert_select.rs @@ -1,11 +1,8 @@ -use pg_raw_parse::{Node, deparse}; +use pg_raw_parse::Node; use crate::frontend::router::parser::{ StatementParser, - rewrite::statement::{ - insert::{Subquery, SubqueryType}, - subquery::SubqueryRewrite, - }, + rewrite::statement::{insert::Subquery, subquery::SubqueryRewrite}, }; impl<'a, 'mem> SubqueryRewrite<'a, 'mem> { @@ -31,13 +28,22 @@ impl<'a, 'mem> SubqueryRewrite<'a, 'mem> { if is_insert_sharded { // We're selecting data from something into a sharded table. // Even if it's a function like generate_series, we need to evaluate it. - return Some(Subquery { - inner_stmt: deparse(select).unwrap().as_str().to_string(), - identifier: SubqueryType::InsertSelect, - }); + return Some(Subquery::InsertSelect); } } None } } + +pub(crate) fn parse_insert_select<'a>(entire_ast: Node<'a>) -> Option> { + if let Node::InsertStmt(insert_stmt) = entire_ast + && let Node::SelectStmt(select) = insert_stmt.select_stmt() + // This won't be a SubLink node as the select statement is naturally apart of the insert's AST. + && !select.from_clause().is_empty() + { + return Some(insert_stmt.select_stmt()); + } + + None +} diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/subquery/mod.rs b/pgdog/src/frontend/router/parser/rewrite/statement/subquery/mod.rs index 037c6e5f0..b49e4865a 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/subquery/mod.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/subquery/mod.rs @@ -12,13 +12,14 @@ use crate::frontend::router::parser::{ pub(crate) mod from_clause; pub(crate) mod insert_select; pub(crate) mod where_clause; +pub(crate) mod with_clause; impl<'a> StatementRewrite<'a> { /// Handle planning Subqueries to later be executed and re-written by the multi step code. pub(crate) fn plan_subquery<'mem>( &mut self, outer_stmt: Node<'mem>, - ) -> Result, Error> { + ) -> Result, Error> { let mut subquery_rewrite = SubqueryRewrite::new(outer_stmt, self); subquery_rewrite.plan_subquery() @@ -50,24 +51,37 @@ impl<'a, 'mem> SubqueryRewrite<'a, 'mem> { } } - fn plan_subquery(&mut self) -> Result, Error> { + fn plan_subquery(&mut self) -> Result, Error> { + // Determine if the statement has a CTE. + // e.g., WITH c AS (..) SELECT .. WHERE IN (SELECT .. FROM c) + if let Some(list_of_subqueries) = self.handle_subquery_in_with_clause() + && !list_of_subqueries.is_empty() + { + return Ok(list_of_subqueries); + } + + let where_clause_ctes = self.handle_ctes_in_where_clause_subquery(); + if !where_clause_ctes.is_empty() { + return Ok(where_clause_ctes); + } + // Determine if there's a SubLink node in the WHERE clause of a SELECT / UPDATE / DELETE statement. if let Some(subquery) = self.handle_subquery_in_where()? { - return Ok(Some(subquery)); + return Ok(vec![subquery]); } // Determine if there's a RangeSubselect within a FROM clause of a SELECT / UPDATE statement. if let Some(subquery) = self.handle_subquery_in_from_clause() { - return Ok(Some(subquery)); + return Ok(vec![subquery]); } // Determine if there's a SELECT FROM in an INSERT. // E.g., INSERT INTO .. SELECT FROM generate_series(..) if let Some(subquery) = self.handle_subquery_in_insert_select() { - return Ok(Some(subquery)); + return Ok(vec![subquery]); } - Ok(None) + Ok(vec![]) } /// Returns the outer statement's Shard. diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/subquery/where_clause.rs b/pgdog/src/frontend/router/parser/rewrite/statement/subquery/where_clause.rs index e1cb6f73e..2835df088 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/subquery/where_clause.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/subquery/where_clause.rs @@ -1,12 +1,8 @@ -use pg_raw_parse::{Node, deparse, nodes::SubLink}; +use pg_raw_parse::{Node, nodes::SubLink}; use crate::frontend::router::parser::{ StatementParser, - rewrite::statement::{ - Error, - insert::{Subquery, SubqueryType}, - subquery::SubqueryRewrite, - }, + rewrite::statement::{Error, insert::Subquery, subquery::SubqueryRewrite}, }; impl<'a, 'mem> SubqueryRewrite<'a, 'mem> { @@ -43,19 +39,19 @@ impl<'a, 'mem> SubqueryRewrite<'a, 'mem> { return Ok(None); } - let subselect_sql = deparse(subselect).unwrap(); - let subselect_sql = subselect_sql.as_str().to_string(); - - return Ok(Some(Subquery { - inner_stmt: subselect_sql, - identifier: SubqueryType::WhereClause, - })); + return Ok(Some(Subquery::WhereClause)); } Ok(None) } } +pub(crate) fn parse_where_clause<'a>(outer_stmt: Node<'a>) -> Option> { + let (sub_link, _) = has_sub_link_in_where(outer_stmt)?; + + Some(sub_link.subselect()) +} + /// When we have a SELECT / UPDATE / DELETE, is there a SubLink present in its /// WhereClause? If there is, return it. /// diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/subquery/with_clause.rs b/pgdog/src/frontend/router/parser/rewrite/statement/subquery/with_clause.rs new file mode 100644 index 000000000..2bae915c8 --- /dev/null +++ b/pgdog/src/frontend/router/parser/rewrite/statement/subquery/with_clause.rs @@ -0,0 +1,118 @@ +//! Handle CTEs +//! + +use crate::frontend::router::parser::rewrite::statement::{ + insert::Subquery, + subquery::{SubqueryRewrite, where_clause::parse_where_clause}, +}; + +use pg_raw_parse::{Node, nodes::WithClause, walk::walk}; + +impl<'a, 'mem> SubqueryRewrite<'a, 'mem> { + pub(crate) fn handle_subquery_in_with_clause(&self) -> Option> { + let with_clause = match self.outer_stmt { + Node::SelectStmt(select) => select.with_clause(), + Node::UpdateStmt(update) => update.with_clause(), + _ => return None, + }?; + + let cte_names = cte_names(with_clause); + let ctes = with_clause.ctes(); + + let mut list_of_subqueries: Vec = Vec::new(); + for cte in ctes { + if !reads_tables(cte.ctequery(), &cte_names) { + continue; + } + + list_of_subqueries.push(Subquery::Cte(cte.ctename().unwrap().to_string())) + } + + Some(list_of_subqueries) + } +} + +impl<'a, 'mem> SubqueryRewrite<'a, 'mem> { + /// If there's a CTE within a WHERE clause, we need to evaluate the CTE separately from the WHERE, + /// so that <> + pub(crate) fn handle_ctes_in_where_clause_subquery(&self) -> Vec { + let Some(Node::SelectStmt(subselect)) = parse_where_clause(self.outer_stmt) else { + return vec![]; + }; + + let Some(with_clause) = subselect.with_clause() else { + return vec![]; + }; + + let cte_names = cte_names(with_clause); + + with_clause + .ctes() + .into_iter() + .filter(|cte| reads_tables(cte.ctequery(), &cte_names)) + .filter_map(|cte| cte.ctename()) + .map(|cte_name| Subquery::WhereClauseCte(cte_name.to_string())) + .collect() + } +} + +fn cte_names(with_clause: &WithClause) -> Vec<&str> { + with_clause + .ctes() + .into_iter() + .filter_map(|cte| cte.ctename()) + .collect() +} + +/// Determine if the CTE references actual tables (as opposed to other CTEs). +/// If not, then we don't need to re-write! +fn reads_tables(cte_query: Node<'_>, cte_names: &[&str]) -> bool { + let mut reads_tables = false; + + walk(cte_query, |node| { + if let Node::RangeVar(range_var) = node { + let is_cte_reference = range_var.schemaname().is_none() + && range_var + .relname() + .is_some_and(|relname| cte_names.contains(&relname)); + + if !is_cte_reference { + reads_tables = true; + } + } + }); + + reads_tables +} + +/// If `entire_ast` contains a CTE with alias `alias`, then return the corresponding +/// AST for the query of that CTE. +/// This is used in our multi-step code, but provided here as it's parsing logic +pub(crate) fn parse_with_clause<'a>(entire_ast: Node<'a>, alias: &String) -> Option> { + let with_clause = match entire_ast { + Node::SelectStmt(select) => select.with_clause(), + Node::UpdateStmt(update) => update.with_clause(), + _ => return None, + }?; + + let ctes = with_clause.ctes(); + + for cte in ctes { + let cte_query = cte.ctequery(); + + if let Node::SelectStmt(cte_query) = cte_query + && !cte_query.values_lists().is_empty() + { + // If it does NOT have an empty VALUES list, then it has already been re-written by us (or we don't need to eval it) + continue; + } + + if let Some(cte_name) = cte.ctename() + && cte_name == alias + { + return Some(cte.ctequery()); + } + } + + None +}