From 4e17c0b4ed004186af79421eb3ee891c7ed1a574 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Fri, 10 Jul 2026 15:44:16 -0600 Subject: [PATCH] Port StatementRewrite::split_insert_stmt to the new parser This is mostly a straight port, with some minor control flow changes to make it read better. The only notable behavior change is that the params it stores are one indexed instead of zero indexed, and that it doesn't duplicate parameters that are used multiple times. This makes it behave nearly identically to the statement struct we use for multi-step updates, which I intend to unify this with in the long term. --- Cargo.lock | 2 +- pgdog/Cargo.toml | 2 +- .../router/parser/rewrite/statement/insert.rs | 303 +++++++++++++----- .../router/parser/rewrite/statement/mod.rs | 29 +- 4 files changed, 246 insertions(+), 90 deletions(-) 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)?; }