diff --git a/Cargo.lock b/Cargo.lock index 22520deab..1d982e3a8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3313,7 +3313,7 @@ dependencies = [ [[package]] name = "pg_raw_parse" version = "0.1.0" -source = "git+https://github.com/pgdogdev/pg_raw_parse.git?rev=b2a15ec#b2a15ec8a4a3a0a9a8d702901e6e336c63aec72b" +source = "git+https://github.com/pgdogdev/pg_raw_parse.git?rev=6421337#6421337b411e91b4b0ca553b8146108c2a273667" dependencies = [ "bindgen 0.72.1", "cc", diff --git a/pgdog/Cargo.toml b/pgdog/Cargo.toml index 03d111e7f..7f0f4ab26 100644 --- a/pgdog/Cargo.toml +++ b/pgdog/Cargo.toml @@ -80,7 +80,7 @@ smallvec = "1" reqwest.workspace = true hex = "0.4" x509-parser = "0.18" -pg_raw_parse = { git = "https://github.com/pgdogdev/pg_raw_parse.git", rev = "b2a15ec", optional = true } +pg_raw_parse = { git = "https://github.com/pgdogdev/pg_raw_parse.git", rev = "6421337", optional = true } itertools = "0.15.0" [target.'cfg(unix)'.dependencies] diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs b/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs index 7bcc63ea3..0c1744502 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs @@ -1,9 +1,17 @@ +#[cfg(feature = "new_parser")] +use indexmap::IndexSet; +#[cfg(not(feature = "new_parser"))] use pg_query::{Node, NodeEnum}; -use pgdog_config::{QueryParserEngine, RewriteMode}; +#[cfg(feature = "new_parser")] +use pg_raw_parse::{Node, NodeMut, deparse, make, nodes, walk}; +#[cfg(not(feature = "new_parser"))] +use pgdog_config::QueryParserEngine; +use pgdog_config::RewriteMode; use crate::frontend::router::Ast; use crate::frontend::router::parser::Cache; use crate::frontend::{BufferedQuery, ClientRequest}; +#[cfg(any(test, not(feature = "new_parser")))] use crate::net::messages::bind::{Format, Parameter}; use crate::net::{Bind, Parse, ProtocolMessage, Query}; @@ -14,6 +22,9 @@ pub struct InsertSplit { /// Parameter positions in the original Bind message /// that should be used to build the Bind message specific to this /// insert statement. + #[cfg(feature = "new_parser")] + params: IndexSet, + #[cfg(not(feature = "new_parser"))] params: Vec, /// The split up INSERT statement with parameters and/or values. @@ -28,11 +39,6 @@ pub struct InsertSplit { } impl InsertSplit { - /// Get the parameter positions from the original Bind message. - pub fn params(&self) -> &[u16] { - &self.params - } - /// Get the SQL statement. pub fn stmt(&self) -> &str { &self.stmt @@ -101,35 +107,55 @@ impl InsertSplit { } /// Extract specific parameters from a Bind message based on this split's param indices. + #[cfg(feature = "new_parser")] fn extract_bind_params(&self, bind: &Bind) -> Bind { - let params: Vec = self - .params - .iter() - .filter_map(|&idx| bind.params_raw().get(idx as usize).cloned()) - .collect(); - - let codes: Vec = if bind.format_codes_raw().len() == 1 { - // Uniform format: keep it - bind.format_codes_raw().clone() - } else if bind.format_codes_raw().len() == bind.params_raw().len() { - // One-to-one mapping: extract corresponding codes - self.params - .iter() - .filter_map(|&idx| bind.format_codes_raw().get(idx as usize).copied()) - .collect() - } else { - // No codes (all text) - Vec::new() - }; + let mut new = Bind::new_statement(self.statement_name().unwrap_or_default()); + for param in &self.params { + let param = bind + .parameter(*param as usize - 1) + .expect("Remove when port done") + .ok_or(Error::MissingParameter(*param)) + .expect("Remove when port done"); + new.push_param(param.parameter().clone(), param.format()); + } - // Use the split's registered statement name if available, - // otherwise fall back to the original bind's statement name. - let statement_name = self - .statement_name - .as_deref() - .unwrap_or_else(|| bind.statement()); + new + } - Bind::new_params_codes(statement_name, ¶ms, &codes) + cfg_select! { + not(feature = "new_parser") => { + fn extract_bind_params(&self, bind: &Bind) -> Bind { + let params: Vec = self + .params + .iter() + .filter_map(|&idx| bind.params_raw().get(idx as usize).cloned()) + .collect(); + + let codes: Vec = if bind.format_codes_raw().len() == 1 { + // Uniform format: keep it + bind.format_codes_raw().clone() + } else if bind.format_codes_raw().len() == bind.params_raw().len() { + // One-to-one mapping: extract corresponding codes + self.params + .iter() + .filter_map(|&idx| bind.format_codes_raw().get(idx as usize).copied()) + .collect() + } else { + // No codes (all text) + Vec::new() + }; + + // Use the split's registered statement name if available, + // otherwise fall back to the original bind's statement name. + let statement_name = self + .statement_name + .as_deref() + .unwrap_or_else(|| bind.statement()); + + Bind::new_params_codes(statement_name, ¶ms, &codes) + } + } + _ => {} } } @@ -158,23 +184,40 @@ impl StatementRewrite<'_> { /// INSERT INTO my_table (id, value) VALUES ($1, $2) -- These are copied from params $3 and $4 /// ``` /// - pub(super) fn split_insert(&mut self, plan: &mut RewritePlan) -> Result<(), Error> { + #[cfg(feature = "new_parser")] + pub(super) fn split_insert( + &mut self, + insert: &nodes::InsertStmt, + plan: &mut RewritePlan, + ) -> Result<(), Error> { // Don't rewrite INSERTs in unsharded databases. if self.schema.shards == 1 || self.schema.rewrite.split_inserts != RewriteMode::Rewrite { return Ok(()); } - let splits: Vec<(Vec, String)> = { - let values_lists = match self.get_insert_values_lists() { - Some(lists) if lists.len() > 1 => lists, - _ => return Ok(()), - }; + let mut splits = Vec::new(); + make::try_owned(|mem| { + let mut copy = mem.make_unique(insert); - values_lists - .iter() - .map(|values_list| self.build_single_tuple_insert(values_list)) - .collect::, _>>()? - }; + if let Node::SelectStmt(select) = insert.select_stmt() { + for list in select.values_lists() { + let (params, select) = self.build_single_tuple_select(mem, list); + copy.as_mut().set_select_stmt(select.uncast()); + splits.push((params, deparse(&*copy)?.as_str().to_string())); + } + } + + Ok::<_, Error>(copy) + })?; + + if splits.len() <= 1 { + return Ok(()); + } + + // FIXME(sage): This is extremely duplicated with the work we do for + // multi-step updates. (#1178 for inserts) Both pieces of code have + // distinct sets of bugs. We should unify those two parts of the code + // base and make this behave consistently. // Now create Ast for each split (needs mutable borrow of prepared_statements) let cache = Cache::get(); @@ -211,7 +254,66 @@ impl StatementRewrite<'_> { Ok(()) } + cfg_select! { + not(feature = "new_parser") => { + pub(super) fn split_insert(&mut self, plan: &mut RewritePlan) -> Result<(), Error> { + // Don't rewrite INSERTs in unsharded databases. + if self.schema.shards == 1 || self.schema.rewrite.split_inserts != RewriteMode::Rewrite { + return Ok(()); + } + + let splits: Vec<(Vec, String)> = { + let values_lists = match self.get_insert_values_lists() { + Some(lists) if lists.len() > 1 => lists, + _ => return Ok(()), + }; + + values_lists + .iter() + .map(|values_list| self.build_single_tuple_insert(values_list)) + .collect::, _>>()? + }; + + // Now create Ast for each split (needs mutable borrow of prepared_statements) + let cache = Cache::get(); + let ctx = self.ast_context(); + for (params, stmt) in splits { + let query = if self.extended { + BufferedQuery::Prepared(Parse::named("", &stmt)) + } else { + BufferedQuery::Query(Query::new(&stmt)) + }; + let ast = cache + .query(&query, &ctx, self.prepared_statements) + .map_err(|e| Error::Cache(e.to_string()))?; + + // If this is a named prepared statement, register the split in the global cache + // and store the assigned name for use in Bind messages. + let statement_name = if self.prepared { + // Name will be assigned by `insert`. + let mut parse = Parse::named("", &stmt); + self.prepared_statements.insert(&mut parse); + Some(parse.name().to_owned()) + } else { + None + }; + + plan.insert_split.push(InsertSplit { + params, + stmt, + ast, + statement_name, + }); + } + + Ok(()) + } + } + _ => {} + } + /// Get the values_lists from an INSERT statement, if present. + #[cfg(not(feature = "new_parser"))] fn get_insert_values_lists(&self) -> Option<&[Node]> { let stmt = self.stmt.stmts.first()?; let node = stmt.stmt.as_ref()?; @@ -229,36 +331,64 @@ impl StatementRewrite<'_> { /// Build a single-tuple INSERT from the original statement with just one values_list. /// Returns the parameter positions (0-indexed) and the SQL string. - fn build_single_tuple_insert(&self, values_list: &Node) -> Result<(Vec, String), Error> { - let mut ast = self.stmt.clone(); - let mut params = Vec::new(); - - // Collect parameter references from this values_list - Self::collect_params(values_list, &mut params); - - // Renumber parameters to start from $1 - let mut new_values_list = values_list.clone(); - Self::renumber_params(&mut new_values_list, ¶ms); - - // Replace the values_lists with just this one tuple - if let Some(stmt) = ast.stmts.first_mut() - && let Some(node) = stmt.stmt.as_mut() - && let Some(NodeEnum::InsertStmt(insert)) = node.node.as_mut() - && let Some(select) = insert.select_stmt.as_mut() - && let Some(NodeEnum::SelectStmt(select_stmt)) = select.node.as_mut() - { - select_stmt.values_lists = vec![new_values_list]; - } + #[cfg(feature = "new_parser")] + fn build_single_tuple_select<'mem>( + &self, + mem: make::MemoryToken<'mem>, + values_list: Node<'_>, + ) -> (IndexSet, make::Unique<'mem, &'mem nodes::SelectStmt>) { + let mut tuple = mem.make_unique(values_list); + + let mut params = IndexSet::new(); + walk::walk_mut(tuple.as_mut(), |node| match node { + NodeMut::ParamRef(param) => { + params.insert(param.number as _); + param.set_number(params.get_index_of(&(param.number as u16)).unwrap() as i32 + 1) + } + _ => (), + }); + + let mut select = mem.make_node::(); + select.as_mut().set_values_lists(mem.make_list(&[tuple])); + (params, select) + } + + cfg_select! { + not(feature = "new_parser") => { + fn build_single_tuple_insert(&self, values_list: &Node) -> Result<(Vec, String), Error> { + let mut ast = self.stmt.clone(); + let mut params = Vec::new(); + + // Collect parameter references from this values_list + Self::collect_params(values_list, &mut params); + + // Renumber parameters to start from $1 + let mut new_values_list = values_list.clone(); + Self::renumber_params(&mut new_values_list, ¶ms); + + // Replace the values_lists with just this one tuple + if let Some(stmt) = ast.stmts.first_mut() + && let Some(node) = stmt.stmt.as_mut() + && let Some(NodeEnum::InsertStmt(insert)) = node.node.as_mut() + && let Some(select) = insert.select_stmt.as_mut() + && let Some(NodeEnum::SelectStmt(select_stmt)) = select.node.as_mut() + { + select_stmt.values_lists = vec![new_values_list]; + } - let stmt = match self.schema.query_parser_engine { - QueryParserEngine::PgQueryProtobuf => ast.deparse(), - QueryParserEngine::PgQueryRaw => ast.deparse_raw(), - }?; + let stmt = match self.schema.query_parser_engine { + QueryParserEngine::PgQueryProtobuf => ast.deparse(), + QueryParserEngine::PgQueryRaw => ast.deparse_raw(), + }?; - Ok((params, stmt)) + Ok((params, stmt)) + } + } + _ => {} } /// Collect all parameter references from a node tree. + #[cfg(not(feature = "new_parser"))] fn collect_params(node: &Node, params: &mut Vec) { if let Some(node_enum) = &node.node { match node_enum { @@ -281,6 +411,7 @@ impl StatementRewrite<'_> { } /// Renumber parameters in a node tree based on their position in the params list. + #[cfg(not(feature = "new_parser"))] fn renumber_params(node: &mut Node, params: &[u16]) { if let Some(node_enum) = &mut node.node { match node_enum { @@ -334,6 +465,13 @@ mod tests { fn parse_and_split(sql: &str) -> Vec { let mut ast = pg_query::parse(sql).unwrap().protobuf; + #[cfg(feature = "new_parser")] + let root = pg_raw_parse::parse(sql).unwrap(); + #[cfg(feature = "new_parser")] + let insert = match root.stmts().next() { + Some(Node::InsertStmt(insert)) => insert, + _ => unreachable!(), + }; let mut prepared = PreparedStatements::default(); let schema = default_schema(); let db_schema = default_db_schema(); @@ -348,6 +486,9 @@ mod tests { search_path: None, }); let mut plan = RewritePlan::default(); + #[cfg(feature = "new_parser")] + rewriter.split_insert(&insert, &mut plan).unwrap(); + #[cfg(not(feature = "new_parser"))] rewriter.split_insert(&mut plan).unwrap(); plan.insert_split } @@ -359,14 +500,20 @@ mod tests { assert_eq!(splits.len(), 2); // First tuple uses params 0 and 1 (original $1, $2) - assert_eq!(splits[0].params(), &[0, 1]); + #[cfg(feature = "new_parser")] + assert_eq!(splits[0].params.as_slice(), &[1, 2]); + #[cfg(not(feature = "new_parser"))] + assert_eq!(splits[0].params, &[0, 1]); assert_eq!( splits[0].stmt(), "INSERT INTO my_table (id, value) VALUES ($1, $2)" ); // Second tuple uses params 2 and 3 (original $3, $4), renumbered to $1, $2 - assert_eq!(splits[1].params(), &[2, 3]); + #[cfg(feature = "new_parser")] + assert_eq!(splits[1].params.as_slice(), &[3, 4]); + #[cfg(not(feature = "new_parser"))] + assert_eq!(splits[1].params, &[2, 3]); assert_eq!( splits[1].stmt(), "INSERT INTO my_table (id, value) VALUES ($1, $2)" @@ -388,13 +535,13 @@ mod tests { assert_eq!(splits.len(), 2); // No params for literal values - assert!(splits[0].params().is_empty()); + assert!(splits[0].params.is_empty()); assert_eq!( splits[0].stmt(), "INSERT INTO my_table (id, value) VALUES (1, 'a')" ); - assert!(splits[1].params().is_empty()); + assert!(splits[1].params.is_empty()); assert_eq!( splits[1].stmt(), "INSERT INTO my_table (id, value) VALUES (2, 'b')" @@ -408,13 +555,19 @@ mod tests { assert_eq!(splits.len(), 2); - assert_eq!(splits[0].params(), &[0]); + #[cfg(feature = "new_parser")] + assert_eq!(splits[0].params.as_slice(), &[1]); + #[cfg(not(feature = "new_parser"))] + assert_eq!(splits[0].params, &[0]); assert_eq!( splits[0].stmt(), "INSERT INTO my_table (id, value) VALUES ($1, 'a')" ); - assert_eq!(splits[1].params(), &[1]); + #[cfg(feature = "new_parser")] + assert_eq!(splits[1].params.as_slice(), &[2]); + #[cfg(not(feature = "new_parser"))] + assert_eq!(splits[1].params, &[1]); assert_eq!( splits[1].stmt(), "INSERT INTO my_table (id, value) VALUES ($1, 'b')" @@ -529,7 +682,6 @@ mod tests { // First split: uses params 0 and 1 (original $1, $2) assert_eq!(splits[0].stmt(), "INSERT INTO t (a, b) VALUES ($1, $2)"); - assert_eq!(splits[0].params(), &[0, 1]); let extracted = splits[0].extract_bind_params(&bind); assert_eq!(extracted.params_raw().len(), 2); assert_eq!(extracted.params_raw()[0].data.as_ref(), b"p1"); @@ -540,7 +692,6 @@ mod tests { splits[1].stmt(), "INSERT INTO t (a, b) VALUES ($1, 'literal')" ); - assert_eq!(splits[1].params(), &[2]); let extracted = splits[1].extract_bind_params(&bind); assert_eq!(extracted.params_raw().len(), 1); assert_eq!(extracted.params_raw()[0].data.as_ref(), b"p3"); diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs index 830d99f74..e8e8fcf59 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs @@ -113,18 +113,20 @@ impl<'a> StatementRewrite<'a> { pub fn maybe_rewrite(&mut self, node: Node<'a>) -> Result { let mut plan = RewritePlan::default(); - make::try_owned(|mem| { - match node { - Node::InsertStmt(_) - | Node::SelectStmt(_) - | Node::UpdateStmt(_) - | Node::DeleteStmt(_) => walk::walk(node, |node| match node { - Node::ParamRef(param) => plan.params = plan.params.max(param.number as u16), - _ => (), - }), - _ => (), - } + match node { + Node::InsertStmt(_) + | Node::SelectStmt(_) + | Node::UpdateStmt(_) + | Node::DeleteStmt(_) => walk::walk(node, |node| match node { + Node::ParamRef(param) => plan.params = plan.params.max(param.number as u16), + _ => {} + }), + Node::PrepareStmt(_) | Node::ExecuteStmt(_) | Node::ExplainStmt(_) => {} + // We can't do anything with DDL statements + _ => return Ok(plan), + } + make::try_owned(|mem| { let mut node = mem.make_unique(node); // Handle top-level PREPARE/EXECUTE statements. @@ -170,7 +172,10 @@ impl<'a> StatementRewrite<'a> { plan.stmt = Some(pg_raw_parse::deparse(node.as_ref())?.as_str().to_owned()); } - self.split_insert(&mut plan)?; + if let Node::InsertStmt(insert) = node.as_ref() { + self.split_insert(insert, &mut plan)?; + } + if let Node::UpdateStmt(stmt) = node.as_ref() { self.sharding_key_update(stmt, &mut plan)?; }