From ba7aa569b9ef0cec4faf45d9a4f3f0418f026618 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Mon, 13 Jul 2026 06:46:17 -0600 Subject: [PATCH 1/3] Don't panic on missing binds in insert_split The code on the old parser, surprisingly, treated this operation as infallible, and just sent an incorrect number of parameters to the shards (which given that we don't behave transactionally for split inserts, could lead to a partial insert). This doesn't fix the lack of transactional behavior, but at least in the new port we won't send any rows at all if we don't have enough bind parameters --- .../router/parser/rewrite/statement/insert.rs | 25 ++++++++++--------- .../router/parser/rewrite/statement/plan.rs | 2 +- 2 files changed, 14 insertions(+), 13 deletions(-) diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs b/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs index 0c1744502..71b879f72 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs @@ -55,7 +55,7 @@ impl InsertSplit { } /// Build a ClientRequest from this split and the original request. - pub fn build_request(&self, request: &ClientRequest) -> ClientRequest { + pub fn build_request(&self, request: &ClientRequest) -> Result { let mut new_request = ClientRequest::default(); let mut has_parse = false; @@ -78,7 +78,7 @@ impl InsertSplit { ProtocolMessage::Query(new_query) } ProtocolMessage::Bind(bind) => { - let new_bind = self.extract_bind_params(bind); + let new_bind = self.extract_bind_params(bind)?; ProtocolMessage::Bind(new_bind) } other => other.clone(), @@ -103,28 +103,26 @@ impl InsertSplit { new_request.last_parse = Some(split_parse); } - new_request + Ok(new_request) } /// 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 { + fn extract_bind_params(&self, bind: &Bind) -> Result { 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"); + .parameter(*param as usize - 1)? + .ok_or(Error::MissingParameter(*param))?; new.push_param(param.parameter().clone(), param.format()); } - new + Ok(new) } cfg_select! { not(feature = "new_parser") => { - fn extract_bind_params(&self, bind: &Bind) -> Bind { + fn extract_bind_params(&self, bind: &Bind) -> Result { let params: Vec = self .params .iter() @@ -152,7 +150,7 @@ impl InsertSplit { .as_deref() .unwrap_or_else(|| bind.statement()); - Bind::new_params_codes(statement_name, ¶ms, &codes) + Ok(Bind::new_params_codes(statement_name, ¶ms, &codes)) } } _ => {} @@ -160,7 +158,10 @@ impl InsertSplit { } /// Build separate ClientRequests for each insert split. -pub fn build_split_requests(splits: &[InsertSplit], request: &ClientRequest) -> Vec { +pub fn build_split_requests( + splits: &[InsertSplit], + request: &ClientRequest, +) -> Result, Error> { splits .iter() .map(|split| split.build_request(request)) diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs index 627316696..53cb51d32 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs @@ -143,7 +143,7 @@ impl RewritePlan { // separately (e.g. go/pq with Parse, Describe, Sync). We don't need to rewrite // those since insert split will return the same row(s) as multi-tuple insert. if !self.insert_split.is_empty() && request.is_executable() { - let requests = build_split_requests(&self.insert_split, request); + let requests = build_split_requests(&self.insert_split, request)?; return Ok(RewriteResult::InsertSplit(requests)); } From 24e076a7b6d18227ee8428032f118eeea07005e7 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Mon, 13 Jul 2026 07:01:21 -0600 Subject: [PATCH 2/3] fix tests --- .../router/parser/rewrite/statement/insert.rs | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs b/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs index 71b879f72..21fd1abdc 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs @@ -589,13 +589,13 @@ mod tests { ); // First split uses params 0 and 1 - let extracted = splits[0].extract_bind_params(&bind); + let extracted = splits[0].extract_bind_params(&bind).unwrap(); assert_eq!(extracted.params_raw().len(), 2); assert_eq!(extracted.params_raw()[0].data.as_ref(), b"p0"); assert_eq!(extracted.params_raw()[1].data.as_ref(), b"p1"); // Second split uses params 2 and 3 - let extracted = splits[1].extract_bind_params(&bind); + let extracted = splits[1].extract_bind_params(&bind).unwrap(); assert_eq!(extracted.params_raw().len(), 2); assert_eq!(extracted.params_raw()[0].data.as_ref(), b"p2"); assert_eq!(extracted.params_raw()[1].data.as_ref(), b"p3"); @@ -616,7 +616,7 @@ mod tests { ); // Second split uses params 2 and 3 (Text, Binary) - let extracted = splits[1].extract_bind_params(&bind); + let extracted = splits[1].extract_bind_params(&bind).unwrap(); assert_eq!(extracted.params_raw().len(), 2); assert_eq!(extracted.params_raw()[0].data.as_ref(), b"p2"); assert_eq!(extracted.params_raw()[1].data.as_ref(), b"p3"); @@ -634,7 +634,7 @@ mod tests { &[Format::Binary], // Uniform format ); - let extracted = splits[0].extract_bind_params(&bind); + let extracted = splits[0].extract_bind_params(&bind).unwrap(); assert_eq!(extracted.params_raw().len(), 1); assert_eq!(extracted.format_codes_raw().len(), 1); assert_eq!(extracted.format_codes_raw()[0], Format::Binary); @@ -655,13 +655,13 @@ mod tests { // First split: statement uses $1 with literal, bind extracts param 0 assert_eq!(splits[0].stmt(), "INSERT INTO t (a, b) VALUES ($1, 'lit1')"); - let extracted = splits[0].extract_bind_params(&bind); + let extracted = splits[0].extract_bind_params(&bind).unwrap(); assert_eq!(extracted.params_raw().len(), 1); assert_eq!(extracted.params_raw()[0].data.as_ref(), b"value_for_param1"); // Second split: statement uses $1 (renumbered from $2) with literal, bind extracts param 1 assert_eq!(splits[1].stmt(), "INSERT INTO t (a, b) VALUES ($1, 'lit2')"); - let extracted = splits[1].extract_bind_params(&bind); + let extracted = splits[1].extract_bind_params(&bind).unwrap(); assert_eq!(extracted.params_raw().len(), 1); assert_eq!(extracted.params_raw()[0].data.as_ref(), b"value_for_param2"); } @@ -683,7 +683,7 @@ 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)"); - let extracted = splits[0].extract_bind_params(&bind); + let extracted = splits[0].extract_bind_params(&bind).unwrap(); assert_eq!(extracted.params_raw().len(), 2); assert_eq!(extracted.params_raw()[0].data.as_ref(), b"p1"); assert_eq!(extracted.params_raw()[1].data.as_ref(), b"p2"); @@ -693,7 +693,7 @@ mod tests { splits[1].stmt(), "INSERT INTO t (a, b) VALUES ($1, 'literal')" ); - let extracted = splits[1].extract_bind_params(&bind); + let extracted = splits[1].extract_bind_params(&bind).unwrap(); assert_eq!(extracted.params_raw().len(), 1); assert_eq!(extracted.params_raw()[0].data.as_ref(), b"p3"); } From 0c50bf2ee9e3680b694b49649322c6032340d8f1 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Mon, 13 Jul 2026 08:59:41 -0600 Subject: [PATCH 3/3] Add unit test --- .../router/parser/rewrite/statement/insert.rs | 20 +++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs b/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs index 21fd1abdc..5bebede88 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs @@ -697,4 +697,24 @@ mod tests { assert_eq!(extracted.params_raw().len(), 1); assert_eq!(extracted.params_raw()[0].data.as_ref(), b"p3"); } + + #[test] + #[cfg(feature = "new_parser")] + fn test_extract_bind_params_incorrect_count() { + let splits = parse_and_split("INSERT INTO t (a, b) VALUES ($1, $2), ($3, $4)"); + let bind = Bind::new_params( + "test", + &[ + Parameter::new(b"p1"), + Parameter::new(b"p2"), + Parameter::new(b"p3"), + ], + ); + + std::assert_matches!(splits[0].extract_bind_params(&bind), Ok(_)); + std::assert_matches!( + splits[1].extract_bind_params(&bind), + Err(Error::MissingParameter(_)) + ); + } }