Skip to content

Don't panic on missing binds in insert_split #1181

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sgrif merged 3 commits into from
Jul 13, 2026
+42 −21
Merged

Don't panic on missing binds in insert_split #1181

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ba7aa56
Don't panic on missing binds in insert_split
sgrif Jul 13, 2026
24e076a
fix tests
sgrif Jul 13, 2026
0c50bf2
Add unit test
sgrif Jul 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
61 changes: 41 additions & 20 deletions pgdog/src/frontend/router/parser/rewrite/statement/insert.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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<ClientRequest, Error> {
let mut new_request = ClientRequest::default();
let mut has_parse = false;

Expand All @@ -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(),
Expand All @@ -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<Bind, Error> {
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))?;
Comment thread
sgrif marked this conversation as resolved.
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<Bind, Error> {
let params: Vec<Parameter> = self
.params
.iter()
Expand Down Expand Up @@ -152,15 +150,18 @@ impl InsertSplit {
.as_deref()
.unwrap_or_else(|| bind.statement());

Bind::new_params_codes(statement_name, &params, &codes)
Ok(Bind::new_params_codes(statement_name, &params, &codes))
}
}
_ => {}
}
}

/// Build separate ClientRequests for each insert split.
pub fn build_split_requests(splits: &[InsertSplit], request: &ClientRequest) -> Vec<ClientRequest> {
pub fn build_split_requests(
splits: &[InsertSplit],
request: &ClientRequest,
) -> Result<Vec<ClientRequest>, Error> {
splits
.iter()
.map(|split| split.build_request(request))
Expand Down Expand Up @@ -588,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");
Expand All @@ -615,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");
Expand All @@ -633,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);
Expand All @@ -654,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");
}
Expand All @@ -682,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");
Expand All @@ -692,8 +693,28 @@ 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");
}

#[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(_))
);
}
}
2 changes: 1 addition & 1 deletion pgdog/src/frontend/router/parser/rewrite/statement/plan.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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));
}

Expand Down
Loading