From f22c774c97e57cd020442a7537fb5450f10b321e Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Fri, 3 Jul 2026 12:37:53 -0600 Subject: [PATCH 1/9] Port StatementRewrite::sharding_key_update to the new parser More specifically, this ports over the glue code as well as the last subroutine of that function, `create_stmts`. This should be the beefiest function we have to port (dear god please let this be the beefiest one). The new implementation diverges from the old in a couple of ways: - We deparse things lazily instead of eagerly - The old code was only deparsing eagerly due to the slowness of deparsing. This should no longer be slow enough to justify that. If needed we can memoize the result, but I'd rather have it be a memoized method than something upstream code cares about - We do not rewrite the bind parameters, instead making every query use all of them - The old code did a recursive mutable AST walk to look for any bind parameters in the query, rewrite/compress the numbers to be a contiguous list starting from 1, and then mapepd that to the original. I *really* don't want to have a mutable AST walk for the new parser, so instead I just shove a no-op CTE that uses all the params in the queries and copy the original bind over. - ShardingKeyUpdate holds the entire original UpdateStmt rather than smaller pieces - I actually tried having this hold onto references of the original query, but ultimately this struct ends up in a field that can't is a sibling of the original AST, so it can't hold a lifetime like that. I probably will revist that at some point, but for now copying the whole thing made more sense than copying each smaller piece into its own allocation The unit tests look a little funky now to allow the new code to pass. What they're testing is technically a little bit weaker than before (it just checks that the bind parameters are all in the right places, it gives zero shits about the numbers). However, once the port is finished, I'll likely delete all but one of those unit tests. The complex logic they were testing simply no longer exists, and we really just need one test to assert that the where clause gets stuck on all 3 subqueries, and that all 3 queries use the same bind parameters as the original update The CTE hack for the bind params really is funky, but I couldn't come up with any less hacky way to do it. I would very much prefer to stick `SELECT $1, $2, ...` somewhere instead, but that won't work. We can't rely on the client having sent us OIDs for the bind parameters in their `Parse` message, and any binds that appear in a top level select will get inferred to `text`. Since we can't assume any random blob of bytes is valid UTF-8, that would cause the query to fail. Putting the original update with `false AND` in the where clause was the simplest way I could think of to ensure the query uses every bind parameter, and keep the same type inference that we had before. Other than the caveats mentioned above, the actual logic is mostly a line for line port. The syntax is certainly very different for constructing these trees, but it's a clean mapping each way. --- Cargo.lock | 2 +- pgdog/Cargo.toml | 2 +- .../client/query_engine/multi_step/update.rs | 14 +- pgdog/src/frontend/router/parser/cache/ast.rs | 49 +- .../router/parser/rewrite/statement/update.rs | 787 +++++++++++++----- 5 files changed, 649 insertions(+), 205 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 72f8f8fcb..380a29f00 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3299,7 +3299,7 @@ dependencies = [ [[package]] name = "pg_raw_parse" version = "0.1.0" -source = "git+https://github.com/pgdogdev/pg_raw_parse.git?rev=6f5d35c#6f5d35ccda90a465f1a224976c30ea41e1b1a858" +source = "git+https://github.com/pgdogdev/pg_raw_parse.git?rev=714b553#714b553572bfebfe3cfce10bdd710833c05efa36" dependencies = [ "bindgen 0.72.1", "cc", diff --git a/pgdog/Cargo.toml b/pgdog/Cargo.toml index 66794d9a4..24e05a48d 100644 --- a/pgdog/Cargo.toml +++ b/pgdog/Cargo.toml @@ -80,7 +80,7 @@ smallvec = "1" reqwest = { version = "0.12", default-features = false, features = ["rustls-tls-webpki-roots-no-provider", "json"] } hex = "0.4" x509-parser = "0.18" -pg_raw_parse = { git = "https://github.com/pgdogdev/pg_raw_parse.git", rev = "6f5d35c", optional = true } +pg_raw_parse = { git = "https://github.com/pgdogdev/pg_raw_parse.git", rev = "714b553", optional = true } itertools = "0.15.0" [target.'cfg(unix)'.dependencies] diff --git a/pgdog/src/frontend/client/query_engine/multi_step/update.rs b/pgdog/src/frontend/client/query_engine/multi_step/update.rs index 9570d2b24..78443e4f8 100644 --- a/pgdog/src/frontend/client/query_engine/multi_step/update.rs +++ b/pgdog/src/frontend/client/query_engine/multi_step/update.rs @@ -99,7 +99,7 @@ impl<'a> UpdateMulti<'a> { context: &mut QueryEngineContext<'_>, row: Row, ) -> Result<(), Error> { - let mut request = self.rewrite.insert.build_request( + let mut request = self.rewrite.build_insert_request( context.client_request, &row.row_description, &row.data_row, @@ -140,12 +140,8 @@ impl<'a> UpdateMulti<'a> { } self.delete_row(context).await?; - self.execute_request_internal( - context, - &mut request, - self.rewrite.insert.is_returning(), - ) - .await?; + self.execute_request_internal(context, &mut request, self.rewrite.is_returning()) + .await?; self.engine .process_server_message(context, CommandComplete::new("UPDATE 1").message()?) // We only allow to update one row at a time. @@ -167,9 +163,7 @@ impl<'a> UpdateMulti<'a> { ) -> Result { let cluster = self.engine.backend.cluster()?; let schema = cluster.schema(); - let Some(table) = self.rewrite.target_table() else { - return Ok(false); - }; + let table = self.rewrite.target_table(); let Some(relation) = schema.table(table, cluster.user(), context.params.search_path()) else { diff --git a/pgdog/src/frontend/router/parser/cache/ast.rs b/pgdog/src/frontend/router/parser/cache/ast.rs index 778e5284b..289a5f0ad 100644 --- a/pgdog/src/frontend/router/parser/cache/ast.rs +++ b/pgdog/src/frontend/router/parser/cache/ast.rs @@ -1,5 +1,9 @@ +#[cfg(feature = "new_parser")] +use itertools::*; use once_cell::sync::OnceCell; use pg_query::{NodeEnum, ParseResult, parse, parse_raw}; +#[cfg(feature = "new_parser")] +use pg_raw_parse::{Owned, StmtList}; use pgdog_config::QueryParserEngine; use std::fmt::Debug; use std::ops::Deref; @@ -40,7 +44,7 @@ pub struct AstInner { pub ast: ParseResult, // FIXME(sage): Rename to ast when parser port is done #[cfg(feature = "new_parser")] - pub new_ast: pg_raw_parse::ParseResult, + pub new_ast: Owned, /// AST stats. pub stats: Mutex, /// Rewrite plan. @@ -53,10 +57,29 @@ pub struct AstInner { impl AstInner { /// Create new AST record, with no rewrite or comment routing. - pub fn new(ast: ParseResult) -> Self { + #[cfg(feature = "new_parser")] + pub(crate) fn new(ast: Owned) -> Self { + let deparse_results = ast + .iter() + .map(|s| pg_raw_parse::deparse(s)) + .collect::, _>>() + .unwrap(); + Self { + new_ast: ast, + ast: pg_query::parse(&deparse_results.iter().map(|r| r.as_str()).join("; ")).unwrap(), + stats: Mutex::new(Stats::new()), + rewrite_plan: RewritePlan::default(), + fingerprint: OnceCell::new(), + query_without_comment: "".into(), + } + } + + pub(crate) fn old(ast: ParseResult) -> Self { Self { #[cfg(feature = "new_parser")] - new_ast: pg_raw_parse::parse(&ast.deparse().unwrap()).unwrap(), + new_ast: pg_raw_parse::parse(&ast.deparse().unwrap()) + .unwrap() + .into_inner(), ast, stats: Mutex::new(Stats::new()), rewrite_plan: RewritePlan::default(), @@ -133,7 +156,9 @@ impl Ast { inner: Arc::new(AstInner { stats: Mutex::new(stats), #[cfg(feature = "new_parser")] - new_ast: pg_raw_parse::parse(&ast.deparse().unwrap()).unwrap(), + new_ast: pg_raw_parse::parse(&ast.deparse().unwrap()) + .unwrap() + .into_inner(), ast, rewrite_plan, fingerprint: OnceCell::new(), @@ -171,10 +196,22 @@ impl Ast { comment_role: None, comment_shard: None, query_parser_engine, - inner: Arc::new(AstInner::new(ast)), + inner: Arc::new(AstInner::old(ast)), }) } + /// Create new AST from a parse result. + #[cfg(feature = "new_parser")] + pub fn from_raw_stmts(stmts: Owned) -> Self { + Self { + cached: true, + comment_role: None, + comment_shard: None, + query_parser_engine: QueryParserEngine::default(), + inner: Arc::new(AstInner::new(stmts)), + } + } + /// Create new AST from a parse result. pub fn from_parse_result(parse_result: ParseResult) -> Self { Self { @@ -182,7 +219,7 @@ impl Ast { comment_role: None, comment_shard: None, query_parser_engine: QueryParserEngine::default(), - inner: Arc::new(AstInner::new(parse_result)), + inner: Arc::new(AstInner::old(parse_result)), } } diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs index 62d5ed17e..741270493 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs @@ -1,5 +1,8 @@ -use std::{collections::HashMap, ops::Deref, sync::Arc}; +#[cfg(not(feature = "new_parser"))] +use std::collections::HashMap; +use std::{ops::Deref, sync::Arc}; +#[cfg(not(feature = "new_parser"))] use pg_query::{ Node as PgNode, NodeEnum, protobuf::{ @@ -9,17 +12,24 @@ use pg_query::{ }, }; #[cfg(feature = "new_parser")] -use pg_raw_parse::make::owned; -#[cfg(feature = "new_parser")] -use pg_raw_parse::{DeparseResult, Node, deparse, nodes}; -use pgdog_config::{QueryParserEngine, RewriteMode}; +use pg_raw_parse::{ + DeparseResult, Node, Owned, deparse, + make::owned, + nodes, + walk::{self, Recurse}, +}; +#[cfg(not(feature = "new_parser"))] +use pgdog_config::QueryParserEngine; +use pgdog_config::RewriteMode; +#[cfg(not(feature = "new_parser"))] +use crate::frontend::router::parser::rewrite::statement::visitor::visit_and_mutate_nodes; use crate::{ frontend::{ BufferedQuery, ClientRequest, router::{ Ast, - parser::{Column, Table, Value, rewrite::statement::visitor::visit_and_mutate_nodes}, + parser::{Column, Table, Value}, sharding::ShardedTable, }, }, @@ -35,11 +45,21 @@ use super::*; pub(crate) struct Statement { pub(crate) ast: Ast, pub(crate) stmt: String, + #[cfg(not(feature = "new_parser"))] pub(crate) params: Vec, } impl Statement { /// Create new Bind message for the statement from original Bind. + #[cfg(feature = "new_parser")] + pub(crate) fn rewrite_bind(&self, bind: &Bind) -> Result { + let mut new = bind.clone(); + new.anonymize(); + Ok(new) + } + + /// Create new Bind message for the statement from original Bind. + #[cfg(not(feature = "new_parser"))] pub(crate) fn rewrite_bind(&self, bind: &Bind) -> Result { let mut new = Bind::new_statement(""); // We use anonymous prepared // statements for execution. @@ -102,15 +122,11 @@ impl Deref for ShardingKeyUpdate { } impl ShardingKeyUpdate { - pub(crate) fn target_table(&self) -> Option> { - self.insert.table.as_ref().map(Table::from) - } - pub(crate) fn sharded_table<'a>( &self, sharded_tables: &'a [ShardedTable], ) -> Option<&'a ShardedTable> { - let table = self.target_table()?; + let table = self.target_table(); sharded_tables.iter().find(|sharded| { if let Some(name) = sharded.name.as_ref() @@ -126,12 +142,22 @@ impl ShardingKeyUpdate { return false; } - self.insert.mapping.contains_key(&sharded.column) + #[cfg(feature = "new_parser")] + { + self.from_update + .targetList() + .iter() + .any(|rt| rt.name() == Some(&*sharded.column)) + } + #[cfg(not(feature = "new_parser"))] + { + self.insert.mapping.contains_key(&sharded.column) + } }) } } -#[derive(Debug, Clone)] +#[derive(Debug)] pub(crate) struct Inner { /// Fetch the whole old row. pub(crate) select: Statement, @@ -140,13 +166,213 @@ pub(crate) struct Inner { /// Delete old row from shard. pub(crate) delete: Statement, /// Partial insert statement. + #[cfg(not(feature = "new_parser"))] pub(crate) insert: Insert, + /// Update this is being constructed from + #[cfg(feature = "new_parser")] + // FIXME(sage): There's no reason we need to own this, but this struct is + // ultimately a child of AstInner, where the statement is borrowed from, + // so we can't add a lifetime here. We should see if we can pass the update + // in later as needed + from_update: Owned, +} + +impl Inner { + #[cfg(feature = "new_parser")] + pub(crate) fn target_table(&self) -> Table<'_> { + Table::from( + self.from_update + .relation() + .expect("UPDATE always has table"), + ) + } + + #[cfg(not(feature = "new_parser"))] + pub(crate) fn target_table(&self) -> Table<'_> { + Table::from(&self.insert.table) + } + + /// Build an INSERT statement built from an existing + /// UPDATE statement and a row returned by a SELECT statement. + #[cfg(feature = "new_parser")] + pub(crate) fn build_insert_request( + &self, + request: &ClientRequest, + row_description: &RowDescription, + data_row: &DataRow, + ) -> Result { + let params = request.parameters()?; + + let mut bind = Bind::new_statement(""); + let mut columns = Vec::new(); + let mut values = Vec::new(); + let mut columns_str = Vec::new(); + let mut values_str = Vec::new(); + + let mut bind_idx = 1; + for (row_idx, field) in row_description.iter().enumerate() { + columns_str.push(format!(r#""{}""#, field.name.replace("\"", "\"\""))); // Escape " + + if let Some(value) = self.from_update.targetList().iter().find_map(|rt| { + if rt.name() == Some(&*field.name) { + Some(rt.val()) + } else { + None + } + }) { + let Ok(value) = Value::try_from(value) else { + values_str.push( + deparse_expr([value])? + .as_str() + .trim_start_matches("SELECT ") + .to_owned(), + ); + continue; + }; + + values_str.push(format!("${bind_idx}")); + match value { + Value::Placeholder(number) => { + let param = params + .as_ref() + .expect("param") + .parameter(number as usize - 1)? + .ok_or(Error::MissingParameter(number as u16))?; + bind.push_param(param.parameter().clone(), param.format()) + } + + Value::Integer(int) => { + bind.push_param(Parameter::new(int.to_string().as_bytes()), Format::Text) + } + + Value::String(s) => bind.push_param(Parameter::new(s.as_bytes()), Format::Text), + + Value::Float(f) => { + bind.push_param(Parameter::new(f.to_string().as_bytes()), Format::Text) + } + + Value::Boolean(b) => bind.push_param( + Parameter::new(if b { "t".as_bytes() } else { "f".as_bytes() }), + Format::Text, + ), + + Value::Vector(vec) => { + bind.push_param(Parameter::new(&vec.encode(Format::Text)?), Format::Text) + } + + Value::Null => bind.push_param(Parameter::new_null(), Format::Text), + } + } else { + let value = data_row + .get_raw(row_idx) + .ok_or(Error::MissingColumn(row_idx))?; + + if value.is_null() { + bind.push_param(Parameter::new_null(), Format::Text); + } else { + bind.push_param(Parameter::new(value), Format::Text); + } + + values_str.push(format!("${bind_idx}")); + } + + columns.push(&*field.name); + values.push(bind_idx); + bind_idx += 1; + } + + let returning_list = self.from_update.returningList(); + let insert = owned(|mem| { + let mut insert = mem.make_node::(); + insert + .as_mut() + .set_relation(mem.make_unique(self.from_update.relation())); + let cols = columns + .into_iter() + .map(|c| { + mem.make_ResTarget(Some(c), mem.empty(), mem.none()) + .uncast() + }) + .collect::>(); + insert.as_mut().set_cols(mem.make_List(&cols)); + let mut select = mem.make_node::(); + let values = values + .into_iter() + .map(|v| mem.make_ParamRef(v).uncast()) + .collect::>(); + select + .as_mut() + .set_valuesLists(mem.make_List(&[mem.make_List(&values)])); + insert.as_mut().set_selectStmt(select.uncast()); + insert + .as_mut() + .set_returningList(mem.make_unique(returning_list)); + mem.make_List(&[mem.make_RawStmt(insert.uncast())]) + }); + + let stmt = format!( + "INSERT INTO {} ({}) VALUES ({}){}", + self.target_table(), + columns_str.join(", "), + values_str.join(", "), + if returning_list.is_empty() { + String::from("") + } else { + format!( + " RETURNING {}", + deparse_expr(returning_list)? + .as_str() + .trim_start_matches("SELECT ") + ) + } + ); + + //// Build the AST to be used with the router. + //// It's identical to the string-generated statement above. + let ast = Ast::from_raw_stmts(insert); + + let mut req = ClientRequest::from(vec![ + ProtocolMessage::from(Parse::new_anonymous(&stmt)), + Describe::new_statement("").into(), // So we get both T and t, + bind.into(), + Execute::new().into(), + Sync.into(), + ]); + req.ast = Some(ast); + Ok(req) + } + + #[cfg(not(feature = "new_parser"))] + /// Build an INSERT statement built from an existing + /// UPDATE statement and a row returned by a SELECT statement. + pub(crate) fn build_insert_request( + &self, + request: &ClientRequest, + row_description: &RowDescription, + data_row: &DataRow, + ) -> Result { + self.insert + .build_request(request, row_description, data_row) + } + + #[cfg(feature = "new_parser")] + /// Do we have to return the rows to the client? + pub(crate) fn is_returning(&self) -> bool { + !self.from_update.returningList().is_empty() + } + + #[cfg(not(feature = "new_parser"))] + /// Do we have to return the rows to the client? + pub(crate) fn is_returning(&self) -> bool { + !self.insert.returning_list.is_empty() && self.insert.returnin_list_deparsed.is_some() + } } /// Partially built INSERT statement. -#[derive(Debug, Clone)] +#[derive(Debug)] +#[cfg(not(feature = "new_parser"))] pub(crate) struct Insert { - pub(super) table: Option, + pub(super) table: RangeVar, /// Mapping of column name to `column name = value` from /// the original UPDATE statement. pub(super) mapping: HashMap, @@ -156,6 +382,7 @@ pub(crate) struct Insert { pub(super) returnin_list_deparsed: Option, } +#[cfg(not(feature = "new_parser"))] impl Insert { /// Build an INSERT statement built from an existing /// UPDATE statement and a row returned by a SELECT statement. @@ -253,7 +480,7 @@ impl Insert { } let insert = InsertStmt { - relation: self.table.clone(), + relation: Some(self.table.clone()), cols: columns, select_stmt: Some(Box::new(PgNode { node: Some(NodeEnum::SelectStmt(Box::new(SelectStmt { @@ -273,7 +500,7 @@ impl Insert { ..Default::default() }; - let table = self.table.as_ref().map(Table::from).unwrap(); // SAFETY: We check that UPDATE has a table. + let table = Table::from(&self.table); // This is probably one of the few places in the code where // we shouldn't use the parser. It's quicker to concatenate strings @@ -311,11 +538,6 @@ impl Insert { req.ast = Some(ast); Ok(req) } - - /// Do we have to return the rows to the client? - pub(crate) fn is_returning(&self) -> bool { - !self.returning_list.is_empty() && self.returnin_list_deparsed.is_some() - } } impl<'a> StatementRewrite<'a> { @@ -326,16 +548,14 @@ impl<'a> StatementRewrite<'a> { return Ok(()); } - let stmt_old = self + #[cfg(not(feature = "new_parser"))] + let Some(NodeEnum::UpdateStmt(stmt)) = self .stmt .stmts .first() .and_then(|stmt| stmt.stmt.as_ref().map(|stmt| stmt.node.as_ref())) - .flatten(); - - let stmt_old = if let Some(NodeEnum::UpdateStmt(stmt)) = stmt_old { - stmt - } else { + .flatten() + else { // TODO: Handle EXPLAIN ANALYZE which needs to execute. // We could return a combined plan for all 3 queries // we need to execute. @@ -350,39 +570,21 @@ impl<'a> StatementRewrite<'a> { return Ok(()); }; - if let Some(value) = self.sharding_key_update_check( - #[cfg(not(feature = "new_parser"))] - stmt_old, - #[cfg(feature = "new_parser")] - stmt, - )? { - #[cfg(feature = "new_parser")] - // Convert from pg_raw_parse::nodes::ResTarget to pg_query::protobuf::ResTarget - let parse_result = { - let sql = deparse_node(value)?; - pg_query::parse(sql.as_str())?.protobuf - }; - #[cfg(feature = "new_parser")] - let value = { - let raw_stmt = parse_result.stmts.first().unwrap(); - let Some(NodeEnum::SelectStmt(s)) = raw_stmt.stmt.as_ref().unwrap().node.as_ref() - else { - unreachable!() - }; - let Some(NodeEnum::ResTarget(r)) = s.target_list.first().unwrap().node.as_ref() - else { - unreachable!() - }; - r - }; + if let Some(value) = self.sharding_key_update_check(stmt)? { // Without a WHERE clause, this is a huge // cross-shard rewrite. - if stmt_old.where_clause.is_none() { + #[cfg(feature = "new_parser")] + if let Node::None = stmt.whereClause() { + return Err(Error::WhereClauseMissing); + } + #[cfg(not(feature = "new_parser"))] + if stmt.where_clause.is_none() { return Err(Error::WhereClauseMissing); } plan.sharding_key_update = Some(create_stmts( - stmt_old, + stmt, value, + #[cfg(not(feature = "new_parser"))] self.schema.query_parser_engine, )?); } @@ -416,7 +618,7 @@ impl<'a> StatementRewrite<'a> { Ok(Some(shard_key_assignment)) } else { let expr = shard_key_assignment.val(); - let expr = deparse_expr(expr)?; + let expr = deparse_expr([expr])?; // FIXME: // // We can technically support this. We can inject this into @@ -494,6 +696,7 @@ impl<'a> StatementRewrite<'a> { /// Visit all ParamRef nodes in a ParseResult and renumber them sequentially. /// Returns a sorted list of the original parameter numbers. +#[cfg(not(feature = "new_parser"))] fn rewrite_params(parse_result: &mut ParseResult) -> Result, Error> { let mut params = HashMap::new(); @@ -521,6 +724,7 @@ fn rewrite_params(parse_result: &mut ParseResult) -> Result, Error> { } #[derive(Debug, Clone)] +#[cfg(not(feature = "new_parser"))] pub(super) enum UpdateValue { Value(Box), Expr(String), // We deparse the expression because we can't handle it yet. @@ -541,6 +745,7 @@ pub(super) enum UpdateValue { /// /// This allows us to build a partial INSERT statement. /// +#[cfg(not(feature = "new_parser"))] fn res_targets_to_insert_res_targets( stmt: &UpdateStmt, query_parser_engine: QueryParserEngine, @@ -572,6 +777,7 @@ fn res_targets_to_insert_res_targets( /// /// Transforms `SET column = value` into `column = value` expression /// for use in shard routing validation. +#[cfg(not(feature = "new_parser"))] fn res_target_to_a_expr(res_target: &ResTarget) -> AExpr { let column_ref = ColumnRef { fields: vec![PgNode { @@ -595,6 +801,7 @@ fn res_target_to_a_expr(res_target: &ResTarget) -> AExpr { } } +#[cfg(not(feature = "new_parser"))] fn select_star() -> Vec { vec![PgNode { node: Some(NodeEnum::ResTarget(Box::new(ResTarget { @@ -612,6 +819,7 @@ fn select_star() -> Vec { }] } +#[cfg(not(feature = "new_parser"))] fn parse_result(node: NodeEnum) -> ParseResult { ParseResult { version: pg_query::PG_VERSION_NUM as i32, @@ -625,23 +833,23 @@ fn parse_result(node: NodeEnum) -> ParseResult { /// Deparse an expression node by wrapping it in a SELECT statement. #[cfg(feature = "new_parser")] -fn deparse_expr(node: Node<'_>) -> Result { - let res_target = owned(|mem| mem.make_ResTarget(None, mem.empty(), mem.make_unique(node))); - deparse_node(&res_target) -} - -#[cfg(feature = "new_parser")] -fn deparse_node(node: &nodes::ResTarget) -> Result { +fn deparse_expr<'a>(nodes: impl IntoIterator>) -> Result { let node = owned(|mem| { let mut select = mem.make_node::(); - select - .as_mut() - .set_targetList(mem.make_List(&[mem.make_unique(node)])); + let res_targets = nodes + .into_iter() + .map(|node| match node { + Node::ResTarget(r) => mem.make_unique(r), + _ => mem.make_ResTarget(None, mem.empty(), mem.make_unique(node)), + }) + .collect::>(); + select.as_mut().set_targetList(mem.make_List(&res_targets)); mem.make_RawStmt(select.uncast()) }); deparse(&*node).map_err(Into::into) } +#[cfg(not(feature = "new_parser"))] fn deparse_expr_old( node: &PgNode, query_parser_engine: QueryParserEngine, @@ -659,6 +867,7 @@ fn deparse_expr_old( } /// Deparse a list of expressions by wrapping them into a SELECT statement. +#[cfg(not(feature = "new_parser"))] fn deparse_list( list: &[PgNode], query_parser_engine: QueryParserEngine, @@ -685,6 +894,132 @@ fn deparse_list( Ok(Some(string)) } +#[cfg(feature = "new_parser")] +fn create_stmts<'a>( + stmt: &'a nodes::UpdateStmt, + new_value: &'a nodes::ResTarget, +) -> Result { + let has_bind_params = walk::walk_manual(stmt.into(), |node| match node { + Node::ParamRef(_) => std::ops::ControlFlow::Break(()), + _ => Recurse::yes(), + }) + .is_some(); + + // Stick the original UPDATE in a CTE with `AND false` at the end of the + // WHERE clause to ensure all bind parameters are used, and have their + // types inferred the same way they would have in the original query + let use_all_params = if has_bind_params { + Some(owned(|mem| { + let mut noop_query = mem.make_unique(stmt); + noop_query.as_mut().set_whereClause( + mem.make_BoolExpr( + nodes::BoolExprType::AND_EXPR, + mem.make_List(&[ + mem.make_A_Const(pg_raw_parse::ConstValue::Boolean(false)) + .uncast(), + mem.make_unique(stmt.whereClause()), + ]), + ) + .uncast(), + ); + let cte = mem.make_CommonTableExpr( + "__pgdog_use_bind_params", + mem.empty(), + noop_query.uncast(), + ); + mem.make_WithClause(mem.make_List(&[cte]), false) + })) + } else { + None + }; + + let select_star = owned(|mem| { + let mut select_stmt = mem.make_node::(); + select_stmt.as_mut().set_targetList( + mem.make_List(&[mem.make_ResTarget( + None, + mem.empty(), + mem.make_ColumnRef(mem.make_List(&[mem.make_node::().uncast()])) + .uncast(), + )]), + ); + select_stmt.as_mut().set_fromClause( + mem.make_List(&[mem + .make_unique(stmt.relation().expect("UPDATE always has a table")) + .uncast()]), + ); + select_stmt + .as_mut() + .set_withClause(mem.make_unique(use_all_params.as_deref())); + select_stmt + }); + + let select = owned(|mem| { + let mut select_stmt = mem.make_unique(&*select_star); + select_stmt + .as_mut() + .set_whereClause(mem.make_unique(stmt.whereClause())); + // let params = rewrite_params(&mut select)?; + mem.make_List(&[mem.make_RawStmt(select_stmt.uncast())]) + }); + + let select = Statement { + stmt: deparse(select.first().unwrap())?.as_str().to_owned(), + ast: Ast::from_raw_stmts(select), + }; + + let delete = owned(|mem| { + let mut delete = mem.make_node::(); + delete + .as_mut() + .set_relation(mem.make_unique(stmt.relation())); + delete + .as_mut() + .set_whereClause(mem.make_unique(stmt.whereClause())); + delete + .as_mut() + .set_withClause(mem.make_unique(use_all_params.as_deref())); + // let params = rewrite_params(&mut delete)?; + mem.make_List(&[mem.make_RawStmt(delete.uncast())]) + }); + + let delete = Statement { + stmt: deparse(delete.first().unwrap())?.as_str().to_owned(), + ast: Ast::from_raw_stmts(delete), + }; + + let check = owned(|mem| { + let mut select_stmt = mem.make_unique(&*select_star); + select_stmt.as_mut().set_whereClause( + mem.make_A_Expr( + nodes::A_Expr_Kind::AEXPR_OP, + mem.make_List(&[mem.make_String(Some("=")).uncast()]), + mem.make_ColumnRef(mem.make_List(&[mem.make_String(new_value.name()).uncast()])) + .uncast(), + mem.make_unique(new_value.val()), + ) + .uncast(), + ); + // let params = rewrite_params(&mut select_stmt)?; + mem.make_List(&[mem.make_RawStmt(select_stmt.uncast())]) + }); + + let check = Statement { + stmt: deparse(check.first().unwrap())?.as_str().to_owned(), + ast: Ast::from_raw_stmts(check), + }; + + Ok(ShardingKeyUpdate { + inner: Arc::new(Inner { + select, + delete, + check, + from_update: owned(|mem| mem.make_unique(stmt)), + }), + }) +} + +#[cfg(not(feature = "new_parser"))] fn create_stmts( stmt: &UpdateStmt, new_value: &ResTarget, @@ -768,7 +1103,7 @@ fn create_stmts( delete, check, insert: Insert { - table: stmt.relation.clone(), + table: stmt.relation.clone().expect("UPDATE always has table"), mapping: res_targets_to_insert_res_targets(stmt, query_parser_engine)?, returning_list: stmt.returning_list.clone(), returnin_list_deparsed: deparse_list(&stmt.returning_list, query_parser_engine)?, @@ -782,6 +1117,7 @@ mod test { use crate::frontend::router::sharding::ShardedTable; use pg_query::parse; use pgdog_config::Rewrite; + use regex::Regex; use crate::backend::schema::Schema; use crate::backend::{ShardedTables, replication::ShardedSchemas}; @@ -844,67 +1180,82 @@ mod test { #[test] fn test_select_basic_where_param() { - let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2") - .unwrap() - .unwrap(); + let sql = "UPDATE sharded SET id = $1 WHERE email = $2"; + let result = run_test(sql).unwrap().unwrap(); // SELECT should have WHERE clause with param renumbered to $1 - assert_eq!(result.select.stmt, "SELECT * FROM sharded WHERE email = $1"); + assert_equivalent_sql( + &result.select.stmt, + "SELECT * FROM sharded WHERE email = $1", + ); + #[cfg(feature = "new_parser")] + assert_all_bind_params_used(sql, &result.select.stmt); + #[cfg(not(feature = "new_parser"))] assert_eq!(result.select.params, vec![2]); let schema = default_schema(); let tables = schema.tables.tables(); - assert_eq!(result.target_table().unwrap().name, "sharded"); + assert_eq!(result.target_table().name, "sharded"); assert_eq!(result.sharded_table(tables).unwrap().column, "id"); } #[test] fn test_select_multiple_where_params() { - let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2 AND name = $3") - .unwrap() - .unwrap(); + let query = "UPDATE sharded SET id = $1 WHERE email = $2 AND name = $3"; + let result = run_test(query).unwrap().unwrap(); - assert_eq!( - result.select.stmt, - "SELECT * FROM sharded WHERE email = $1 AND name = $2" + assert_equivalent_sql( + &result.select.stmt, + "SELECT * FROM sharded WHERE email = $1 AND name = $2", ); + #[cfg(feature = "new_parser")] + assert_all_bind_params_used(query, &result.select.stmt); + #[cfg(not(feature = "new_parser"))] assert_eq!(result.select.params, vec![2, 3]); - assert!(!result.insert.is_returning()); + assert!(!result.is_returning()); } #[test] fn test_select_non_sequential_params() { - // Params in WHERE are $3 and $5, should be renumbered to $1 and $2 - let result = run_test( - "UPDATE sharded SET id = $1, value = $2, other = $4 WHERE email = $3 AND name = $5", - ) - .unwrap() - .unwrap(); - - assert_eq!( - result.select.stmt, - "SELECT * FROM sharded WHERE email = $1 AND name = $2" + // Resulting query must be valid with the client's bind params, even + // though the select stmt will only use $3 and $5 + let query = + "UPDATE sharded SET id = $1, value = $2, other = $4 WHERE email = $3 AND name = $5"; + let result = run_test(query).unwrap().unwrap(); + + assert_equivalent_sql( + &result.select.stmt, + "SELECT * FROM sharded WHERE email = $1 AND name = $2", ); + #[cfg(feature = "new_parser")] + assert_all_bind_params_used(query, &result.select.stmt); + #[cfg(not(feature = "new_parser"))] assert_eq!(result.select.params, vec![3, 5]); } #[test] fn test_select_single_where_param() { - let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2") - .unwrap() - .unwrap(); + let query = "UPDATE sharded SET id = $1 WHERE email = $2"; + let result = run_test(query).unwrap().unwrap(); - assert_eq!(result.select.stmt, "SELECT * FROM sharded WHERE email = $1"); + assert_equivalent_sql( + &result.select.stmt, + "SELECT * FROM sharded WHERE email = $1", + ); + #[cfg(feature = "new_parser")] + assert_all_bind_params_used(query, &result.select.stmt); + #[cfg(not(feature = "new_parser"))] assert_eq!(result.select.params, vec![2]); } #[test] fn test_delete_basic() { - let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2") - .unwrap() - .unwrap(); + let query = "UPDATE sharded SET id = $1 WHERE email = $2"; + let result = run_test(query).unwrap().unwrap(); - assert_eq!(result.delete.stmt, "DELETE FROM sharded WHERE email = $1"); + assert_equivalent_sql(&result.delete.stmt, "DELETE FROM sharded WHERE email = $1"); + #[cfg(feature = "new_parser")] + assert_all_bind_params_used(query, &result.delete.stmt); assert!(result.sharded_table(&[]).is_none()); assert!( @@ -929,78 +1280,89 @@ mod test { #[test] fn test_delete_multiple_where_params() { - let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2 AND name = $3") - .unwrap() - .unwrap(); + let query = "UPDATE sharded SET id = $1 WHERE email = $2 AND name = $3"; + let result = run_test(query).unwrap().unwrap(); - assert_eq!( - result.delete.stmt, - "DELETE FROM sharded WHERE email = $1 AND name = $2" + assert_equivalent_sql( + &result.delete.stmt, + "DELETE FROM sharded WHERE email = $1 AND name = $2", ); + #[cfg(feature = "new_parser")] + assert_all_bind_params_used(query, &result.delete.stmt); } #[test] fn test_no_params_in_where() { - let result = run_test("UPDATE sharded SET id = $1 WHERE email = 'test@example.com'") - .unwrap() - .unwrap(); + let query = "UPDATE sharded SET id = $1 WHERE email = 'test@example.com'"; + let result = run_test(query).unwrap().unwrap(); - assert_eq!( - result.select.stmt, - "SELECT * FROM sharded WHERE email = 'test@example.com'" + assert_equivalent_sql( + &result.select.stmt, + "SELECT * FROM sharded WHERE email = 'test@example.com'", ); + #[cfg(feature = "new_parser")] + assert_all_bind_params_used(query, &result.select.stmt); + #[cfg(not(feature = "new_parser"))] assert_eq!(result.select.params, Vec::::new()); } #[test] fn test_where_with_in_clause() { - let result = run_test("UPDATE sharded SET id = $1 WHERE email IN ($2, $3, $4)") - .unwrap() - .unwrap(); + let query = "UPDATE sharded SET id = $1 WHERE email IN ($2, $3, $4)"; + let result = run_test(query).unwrap().unwrap(); - assert_eq!( - result.select.stmt, - "SELECT * FROM sharded WHERE email IN ($1, $2, $3)" + assert_equivalent_sql( + &result.select.stmt, + "SELECT * FROM sharded WHERE email IN ($1, $2, $3)", ); + #[cfg(feature = "new_parser")] + assert_all_bind_params_used(query, &result.select.stmt); + #[cfg(not(feature = "new_parser"))] assert_eq!(result.select.params, vec![2, 3, 4]); } #[test] fn test_where_with_comparison_operators() { - let result = run_test("UPDATE sharded SET id = $1 WHERE count > $2 AND count < $3") - .unwrap() - .unwrap(); + let query = "UPDATE sharded SET id = $1 WHERE count > $2 AND count < $3"; + let result = run_test(query).unwrap().unwrap(); - assert_eq!( - result.select.stmt, - "SELECT * FROM sharded WHERE count > $1 AND count < $2" + assert_equivalent_sql( + &result.select.stmt, + "SELECT * FROM sharded WHERE count > $1 AND count < $2", ); + #[cfg(feature = "new_parser")] + assert_all_bind_params_used(query, &result.select.stmt); + #[cfg(not(feature = "new_parser"))] assert_eq!(result.select.params, vec![2, 3]); } #[test] fn test_where_with_or_condition() { - let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2 OR name = $3") - .unwrap() - .unwrap(); + let query = "UPDATE sharded SET id = $1 WHERE email = $2 OR name = $3"; + let result = run_test(query).unwrap().unwrap(); - assert_eq!( - result.select.stmt, - "SELECT * FROM sharded WHERE email = $1 OR name = $2" + assert_equivalent_sql( + &result.select.stmt, + "SELECT * FROM sharded WHERE email = $1 OR name = $2", ); + #[cfg(feature = "new_parser")] + assert_all_bind_params_used(query, &result.select.stmt); + #[cfg(not(feature = "new_parser"))] assert_eq!(result.select.params, vec![2, 3]); } #[test] fn test_high_param_numbers() { - let result = run_test("UPDATE sharded SET id = $10 WHERE email = $20 AND name = $30") - .unwrap() - .unwrap(); + let query = "UPDATE sharded SET id = $10 WHERE email = $20 AND name = $30"; + let result = run_test(query).unwrap().unwrap(); - assert_eq!( - result.select.stmt, - "SELECT * FROM sharded WHERE email = $1 AND name = $2" + assert_equivalent_sql( + &result.select.stmt, + "SELECT * FROM sharded WHERE email = $1 AND name = $2", ); + #[cfg(feature = "new_parser")] + assert_all_bind_params_used(query, &result.select.stmt); + #[cfg(not(feature = "new_parser"))] assert_eq!(result.select.params, vec![20, 30]); } @@ -1013,120 +1375,138 @@ mod test { #[test] fn test_where_with_like() { - let result = run_test("UPDATE sharded SET id = $1 WHERE email LIKE $2") - .unwrap() - .unwrap(); + let query = "UPDATE sharded SET id = $1 WHERE email LIKE $2"; + let result = run_test(query).unwrap().unwrap(); - assert_eq!( - result.select.stmt, - "SELECT * FROM sharded WHERE email LIKE $1" + assert_equivalent_sql( + &result.select.stmt, + "SELECT * FROM sharded WHERE email LIKE $1", ); + #[cfg(feature = "new_parser")] + assert_all_bind_params_used(query, &result.select.stmt); + #[cfg(not(feature = "new_parser"))] assert_eq!(result.select.params, vec![2]); } #[test] fn test_where_with_is_null() { - let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2 AND deleted_at IS NULL") - .unwrap() - .unwrap(); + let query = "UPDATE sharded SET id = $1 WHERE email = $2 AND deleted_at IS NULL"; + let result = run_test(query).unwrap().unwrap(); - assert_eq!( - result.select.stmt, - "SELECT * FROM sharded WHERE email = $1 AND deleted_at IS NULL" + assert_equivalent_sql( + &result.select.stmt, + "SELECT * FROM sharded WHERE email = $1 AND deleted_at IS NULL", ); + #[cfg(feature = "new_parser")] + assert_all_bind_params_used(query, &result.select.stmt); + #[cfg(not(feature = "new_parser"))] assert_eq!(result.select.params, vec![2]); } #[test] fn test_where_with_between() { - let result = run_test("UPDATE sharded SET id = $1 WHERE created_at BETWEEN $2 AND $3") - .unwrap() - .unwrap(); + let query = "UPDATE sharded SET id = $1 WHERE created_at BETWEEN $2 AND $3"; + let result = run_test(query).unwrap().unwrap(); - assert_eq!( - result.select.stmt, - "SELECT * FROM sharded WHERE created_at BETWEEN $1 AND $2" + assert_equivalent_sql( + &result.select.stmt, + "SELECT * FROM sharded WHERE created_at BETWEEN $1 AND $2", ); + #[cfg(feature = "new_parser")] + assert_all_bind_params_used(query, &result.select.stmt); + #[cfg(not(feature = "new_parser"))] assert_eq!(result.select.params, vec![2, 3]); } #[test] fn test_same_param_used_twice() { // Same parameter $2 used twice in WHERE clause - let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2 OR name = $2") - .unwrap() - .unwrap(); + let query = "UPDATE sharded SET id = $1 WHERE email = $2 OR name = $2"; + let result = run_test(query).unwrap().unwrap(); // Both occurrences should be renumbered to $1 - assert_eq!( - result.select.stmt, - "SELECT * FROM sharded WHERE email = $1 OR name = $1" + assert_equivalent_sql( + &result.select.stmt, + "SELECT * FROM sharded WHERE email = $1 OR name = $1", ); // Only one unique param in the mapping + #[cfg(feature = "new_parser")] + assert_all_bind_params_used(query, &result.select.stmt); + #[cfg(not(feature = "new_parser"))] assert_eq!(result.select.params, vec![2]); } #[test] fn test_same_param_used_multiple_times() { // $2 used three times - let result = run_test("UPDATE sharded SET id = $1 WHERE a = $2 AND b = $2 AND c = $2") - .unwrap() - .unwrap(); + let query = "UPDATE sharded SET id = $1 WHERE a = $2 AND b = $2 AND c = $2"; + let result = run_test(query).unwrap().unwrap(); - assert_eq!( - result.select.stmt, - "SELECT * FROM sharded WHERE a = $1 AND b = $1 AND c = $1" + assert_equivalent_sql( + &result.select.stmt, + "SELECT * FROM sharded WHERE a = $1 AND b = $1 AND c = $1", ); + #[cfg(feature = "new_parser")] + assert_all_bind_params_used(query, &result.select.stmt); + #[cfg(not(feature = "new_parser"))] assert_eq!(result.select.params, vec![2]); } #[test] fn test_mixed_repeated_and_unique_params() { // $2 used twice, $3 used once - let result = run_test("UPDATE sharded SET id = $1 WHERE a = $2 AND b = $3 AND c = $2") - .unwrap() - .unwrap(); + let query = "UPDATE sharded SET id = $1 WHERE a = $2 AND b = $3 AND c = $2"; + let result = run_test(query).unwrap().unwrap(); - assert_eq!( - result.select.stmt, - "SELECT * FROM sharded WHERE a = $1 AND b = $2 AND c = $1" + assert_equivalent_sql( + &result.select.stmt, + "SELECT * FROM sharded WHERE a = $1 AND b = $2 AND c = $1", ); + #[cfg(feature = "new_parser")] + assert_all_bind_params_used(query, &result.select.stmt); + #[cfg(not(feature = "new_parser"))] assert_eq!(result.select.params, vec![2, 3]); } #[test] fn test_repeated_params_in_in_clause() { // Same param repeated in IN clause (unusual but valid) - let result = run_test("UPDATE sharded SET id = $1 WHERE email IN ($2, $3, $2)") - .unwrap() - .unwrap(); + let query = "UPDATE sharded SET id = $1 WHERE email IN ($2, $3, $2)"; + let result = run_test(query).unwrap().unwrap(); - assert_eq!( - result.select.stmt, - "SELECT * FROM sharded WHERE email IN ($1, $2, $1)" + assert_equivalent_sql( + &result.select.stmt, + "SELECT * FROM sharded WHERE email IN ($1, $2, $1)", ); + #[cfg(feature = "new_parser")] + assert_all_bind_params_used(query, &result.select.stmt); + #[cfg(not(feature = "new_parser"))] assert_eq!(result.select.params, vec![2, 3]); } #[test] fn test_delete_with_repeated_params() { - let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2 OR name = $2") - .unwrap() - .unwrap(); + let query = "UPDATE sharded SET id = $1 WHERE email = $2 OR name = $2"; + let result = run_test(query).unwrap().unwrap(); - assert_eq!( - result.delete.stmt, - "DELETE FROM sharded WHERE email = $1 OR name = $1" + assert_equivalent_sql( + &result.delete.stmt, + "DELETE FROM sharded WHERE email = $1 OR name = $1", ); + #[cfg(feature = "new_parser")] + assert_all_bind_params_used(query, &result.delete.stmt); + #[cfg(not(feature = "new_parser"))] assert_eq!(result.delete.params, vec![2]); } #[test] fn test_sharding_key_not_changed() { - let result = run_test("UPDATE sharded SET id = $1 WHERE id = $1 AND email = $2") - .unwrap() - .unwrap(); - assert_eq!(result.check.stmt, "SELECT * FROM sharded WHERE id = $1"); + let query = "UPDATE sharded SET id = $1 WHERE id = $1 AND email = $2"; + let result = run_test(query).unwrap().unwrap(); + assert_equivalent_sql(&result.check.stmt, "SELECT * FROM sharded WHERE id = $1"); + #[cfg(feature = "new_parser")] + assert_all_bind_params_used(query, &result.check.stmt); + #[cfg(not(feature = "new_parser"))] assert_eq!(result.check.params, vec![1]); } @@ -1232,6 +1612,7 @@ mod test { } #[test] + #[cfg(not(feature = "new_parser"))] fn test_return_rows() { let result = run_test("UPDATE sharded SET id = $1 WHERE id = $2 RETURNING *") .unwrap() @@ -1249,6 +1630,7 @@ mod test { } #[test] + #[cfg(not(feature = "new_parser"))] fn test_res_targets_to_insert_res_targets_expr_branch() { // Test that expression assignments (non-simple values) are deparsed correctly // and stored as UpdateValue::Expr in the insert mapping. @@ -1269,6 +1651,7 @@ mod test { } #[test] + #[cfg(not(feature = "new_parser"))] fn test_res_targets_to_insert_res_targets_expr_arithmetic() { // Test arithmetic expressions are deparsed correctly let result = run_test("UPDATE sharded SET id = $1, counter = counter + 1 WHERE id = $2") @@ -1283,6 +1666,7 @@ mod test { } #[test] + #[cfg(not(feature = "new_parser"))] fn test_res_targets_to_insert_res_targets_expr_coalesce() { // Test COALESCE expressions are deparsed correctly let result = @@ -1325,8 +1709,7 @@ mod test { ))]); let insert_request = result - .insert - .build_request(&request, &row_description, &data_row) + .build_insert_request(&request, &row_description, &data_row) .unwrap(); // Get the query from the request to verify the INSERT statement @@ -1353,4 +1736,34 @@ mod test { stmt ); } + + const BIND_PARAM_PATTERN: &str = r#"(\$\d+)"#; + + fn bind_param_regex() -> Regex { + Regex::new(BIND_PARAM_PATTERN).unwrap() + } + + fn assert_equivalent_sql(actual: &str, expected: &str) { + let pattern = bind_param_regex().replace_all(expected, "BIND_PARAM"); + let pattern = regex::escape(&pattern).replace("BIND_PARAM", BIND_PARAM_PATTERN); + let regex = Regex::new(&pattern).unwrap(); + assert!( + regex.is_match(actual), + "Expected {:?} to match {:?}", + actual, + pattern + ) + } + + fn assert_all_bind_params_used(original: &str, rewritten: &str) { + for captures in bind_param_regex().captures_iter(original) { + let param = &captures[1]; + assert!( + rewritten.contains(param), + "Expected {:?} to use param {}", + rewritten, + param + ); + } + } } From 322781fe300c3778e366c86b483257b294405a8f Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Fri, 3 Jul 2026 13:13:21 -0600 Subject: [PATCH 2/9] cfg code only used in the new parser --- pgdog/src/frontend/router/parser/rewrite/statement/update.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs index 741270493..1a2476c1e 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs @@ -1755,6 +1755,7 @@ mod test { ) } + #[cfg(feature = "new_parser")] fn assert_all_bind_params_used(original: &str, rewritten: &str) { for captures in bind_param_regex().captures_iter(original) { let param = &captures[1]; From d4b7a9557f39deb2d1ae6b6c89c6323706665802 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Sun, 5 Jul 2026 13:43:35 -0600 Subject: [PATCH 3/9] Use a mutable AST walk to mimic the old parser's behavior --- Cargo.lock | 2 +- pgdog/Cargo.toml | 2 +- .../router/parser/rewrite/statement/update.rs | 438 +++++++----------- 3 files changed, 172 insertions(+), 270 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 380a29f00..22db4f3d8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3299,7 +3299,7 @@ dependencies = [ [[package]] name = "pg_raw_parse" version = "0.1.0" -source = "git+https://github.com/pgdogdev/pg_raw_parse.git?rev=714b553#714b553572bfebfe3cfce10bdd710833c05efa36" +source = "git+https://github.com/pgdogdev/pg_raw_parse.git?rev=7a7ace5#7a7ace52b1fbe279d5a30e6d7267077b5a298990" dependencies = [ "bindgen 0.72.1", "cc", diff --git a/pgdog/Cargo.toml b/pgdog/Cargo.toml index 24e05a48d..473e24f1e 100644 --- a/pgdog/Cargo.toml +++ b/pgdog/Cargo.toml @@ -80,7 +80,7 @@ smallvec = "1" reqwest = { version = "0.12", default-features = false, features = ["rustls-tls-webpki-roots-no-provider", "json"] } hex = "0.4" x509-parser = "0.18" -pg_raw_parse = { git = "https://github.com/pgdogdev/pg_raw_parse.git", rev = "714b553", optional = true } +pg_raw_parse = { git = "https://github.com/pgdogdev/pg_raw_parse.git", rev = "7a7ace5", optional = true } itertools = "0.15.0" [target.'cfg(unix)'.dependencies] diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs index 1a2476c1e..de08c9895 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs @@ -1,3 +1,5 @@ +#[cfg(feature = "new_parser")] +use indexmap::IndexSet; #[cfg(not(feature = "new_parser"))] use std::collections::HashMap; use std::{ops::Deref, sync::Arc}; @@ -12,12 +14,7 @@ use pg_query::{ }, }; #[cfg(feature = "new_parser")] -use pg_raw_parse::{ - DeparseResult, Node, Owned, deparse, - make::owned, - nodes, - walk::{self, Recurse}, -}; +use pg_raw_parse::{DeparseResult, Node, NodeMut, Owned, deparse, make::owned, nodes, walk}; #[cfg(not(feature = "new_parser"))] use pgdog_config::QueryParserEngine; use pgdog_config::RewriteMode; @@ -45,21 +42,14 @@ use super::*; pub(crate) struct Statement { pub(crate) ast: Ast, pub(crate) stmt: String, + #[cfg(feature = "new_parser")] + pub(crate) params: IndexSet, #[cfg(not(feature = "new_parser"))] pub(crate) params: Vec, } impl Statement { /// Create new Bind message for the statement from original Bind. - #[cfg(feature = "new_parser")] - pub(crate) fn rewrite_bind(&self, bind: &Bind) -> Result { - let mut new = bind.clone(); - new.anonymize(); - Ok(new) - } - - /// Create new Bind message for the statement from original Bind. - #[cfg(not(feature = "new_parser"))] pub(crate) fn rewrite_bind(&self, bind: &Bind) -> Result { let mut new = Bind::new_statement(""); // We use anonymous prepared // statements for execution. @@ -696,6 +686,19 @@ impl<'a> StatementRewrite<'a> { /// Visit all ParamRef nodes in a ParseResult and renumber them sequentially. /// Returns a sorted list of the original parameter numbers. +#[cfg(feature = "new_parser")] +fn rewrite_params(node: NodeMut<'_, '_>) -> IndexSet { + let mut params = IndexSet::new(); + walk::walk_mut(node, |node| match node { + NodeMut::ParamRef(mut param) => { + params.insert(param.number as _); + param.set_number(params.get_index_of(&(param.number as u16)).unwrap() as i32 + 1) + } + _ => (), + }); + params +} + #[cfg(not(feature = "new_parser"))] fn rewrite_params(parse_result: &mut ParseResult) -> Result, Error> { let mut params = HashMap::new(); @@ -899,40 +902,6 @@ fn create_stmts<'a>( stmt: &'a nodes::UpdateStmt, new_value: &'a nodes::ResTarget, ) -> Result { - let has_bind_params = walk::walk_manual(stmt.into(), |node| match node { - Node::ParamRef(_) => std::ops::ControlFlow::Break(()), - _ => Recurse::yes(), - }) - .is_some(); - - // Stick the original UPDATE in a CTE with `AND false` at the end of the - // WHERE clause to ensure all bind parameters are used, and have their - // types inferred the same way they would have in the original query - let use_all_params = if has_bind_params { - Some(owned(|mem| { - let mut noop_query = mem.make_unique(stmt); - noop_query.as_mut().set_whereClause( - mem.make_BoolExpr( - nodes::BoolExprType::AND_EXPR, - mem.make_List(&[ - mem.make_A_Const(pg_raw_parse::ConstValue::Boolean(false)) - .uncast(), - mem.make_unique(stmt.whereClause()), - ]), - ) - .uncast(), - ); - let cte = mem.make_CommonTableExpr( - "__pgdog_use_bind_params", - mem.empty(), - noop_query.uncast(), - ); - mem.make_WithClause(mem.make_List(&[cte]), false) - })) - } else { - None - }; - let select_star = owned(|mem| { let mut select_stmt = mem.make_node::(); select_stmt.as_mut().set_targetList( @@ -949,25 +918,25 @@ fn create_stmts<'a>( .uncast()]), ); select_stmt - .as_mut() - .set_withClause(mem.make_unique(use_all_params.as_deref())); - select_stmt }); + let mut params = IndexSet::new(); let select = owned(|mem| { let mut select_stmt = mem.make_unique(&*select_star); select_stmt .as_mut() .set_whereClause(mem.make_unique(stmt.whereClause())); - // let params = rewrite_params(&mut select)?; + params = rewrite_params(select_stmt.as_mut().into()); mem.make_List(&[mem.make_RawStmt(select_stmt.uncast())]) }); let select = Statement { stmt: deparse(select.first().unwrap())?.as_str().to_owned(), ast: Ast::from_raw_stmts(select), + params, }; + let mut params = IndexSet::new(); let delete = owned(|mem| { let mut delete = mem.make_node::(); delete @@ -976,18 +945,17 @@ fn create_stmts<'a>( delete .as_mut() .set_whereClause(mem.make_unique(stmt.whereClause())); - delete - .as_mut() - .set_withClause(mem.make_unique(use_all_params.as_deref())); - // let params = rewrite_params(&mut delete)?; + params = rewrite_params(delete.as_mut().into()); mem.make_List(&[mem.make_RawStmt(delete.uncast())]) }); let delete = Statement { stmt: deparse(delete.first().unwrap())?.as_str().to_owned(), ast: Ast::from_raw_stmts(delete), + params, }; + let mut params = IndexSet::new(); let check = owned(|mem| { let mut select_stmt = mem.make_unique(&*select_star); select_stmt.as_mut().set_whereClause( @@ -1000,13 +968,14 @@ fn create_stmts<'a>( ) .uncast(), ); - // let params = rewrite_params(&mut select_stmt)?; + params = rewrite_params(select_stmt.as_mut().into()); mem.make_List(&[mem.make_RawStmt(select_stmt.uncast())]) }); let check = Statement { stmt: deparse(check.first().unwrap())?.as_str().to_owned(), ast: Ast::from_raw_stmts(check), + params, }; Ok(ShardingKeyUpdate { @@ -1115,9 +1084,10 @@ fn create_stmts( #[cfg(test)] mod test { use crate::frontend::router::sharding::ShardedTable; + #[cfg(feature = "new_parser")] + use indexmap::indexset; use pg_query::parse; use pgdog_config::Rewrite; - use regex::Regex; use crate::backend::schema::Schema; use crate::backend::{ShardedTables, replication::ShardedSchemas}; @@ -1125,6 +1095,13 @@ mod test { use super::*; + #[cfg(not(feature = "new_parser"))] + macro_rules! indexset { + ($t:tt) => { + vec![$t] + }; + } + fn default_db_schema() -> Schema { Schema::default() } @@ -1180,18 +1157,13 @@ mod test { #[test] fn test_select_basic_where_param() { - let sql = "UPDATE sharded SET id = $1 WHERE email = $2"; - let result = run_test(sql).unwrap().unwrap(); + let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2") + .unwrap() + .unwrap(); // SELECT should have WHERE clause with param renumbered to $1 - assert_equivalent_sql( - &result.select.stmt, - "SELECT * FROM sharded WHERE email = $1", - ); - #[cfg(feature = "new_parser")] - assert_all_bind_params_used(sql, &result.select.stmt); - #[cfg(not(feature = "new_parser"))] - assert_eq!(result.select.params, vec![2]); + assert_eq!(result.select.stmt, "SELECT * FROM sharded WHERE email = $1"); + assert_eq!(result.select.params, indexset![2]); let schema = default_schema(); let tables = schema.tables.tables(); @@ -1201,61 +1173,51 @@ mod test { #[test] fn test_select_multiple_where_params() { - let query = "UPDATE sharded SET id = $1 WHERE email = $2 AND name = $3"; - let result = run_test(query).unwrap().unwrap(); + let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2 AND name = $3") + .unwrap() + .unwrap(); - assert_equivalent_sql( - &result.select.stmt, - "SELECT * FROM sharded WHERE email = $1 AND name = $2", + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE email = $1 AND name = $2" ); - #[cfg(feature = "new_parser")] - assert_all_bind_params_used(query, &result.select.stmt); - #[cfg(not(feature = "new_parser"))] - assert_eq!(result.select.params, vec![2, 3]); + assert_eq!(result.select.params, indexset![2, 3]); assert!(!result.is_returning()); } #[test] fn test_select_non_sequential_params() { - // Resulting query must be valid with the client's bind params, even - // though the select stmt will only use $3 and $5 - let query = - "UPDATE sharded SET id = $1, value = $2, other = $4 WHERE email = $3 AND name = $5"; - let result = run_test(query).unwrap().unwrap(); - - assert_equivalent_sql( - &result.select.stmt, - "SELECT * FROM sharded WHERE email = $1 AND name = $2", + // Params in WHERE are $3 and $5, should be renumbered to $1 and $2 + let result = run_test( + "UPDATE sharded SET id = $1, value = $2, other = $4 WHERE email = $3 AND name = $5", + ) + .unwrap() + .unwrap(); + + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE email = $1 AND name = $2" ); - #[cfg(feature = "new_parser")] - assert_all_bind_params_used(query, &result.select.stmt); - #[cfg(not(feature = "new_parser"))] - assert_eq!(result.select.params, vec![3, 5]); + assert_eq!(result.select.params, indexset![3, 5]); } #[test] fn test_select_single_where_param() { - let query = "UPDATE sharded SET id = $1 WHERE email = $2"; - let result = run_test(query).unwrap().unwrap(); + let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2") + .unwrap() + .unwrap(); - assert_equivalent_sql( - &result.select.stmt, - "SELECT * FROM sharded WHERE email = $1", - ); - #[cfg(feature = "new_parser")] - assert_all_bind_params_used(query, &result.select.stmt); - #[cfg(not(feature = "new_parser"))] - assert_eq!(result.select.params, vec![2]); + assert_eq!(result.select.stmt, "SELECT * FROM sharded WHERE email = $1"); + assert_eq!(result.select.params, indexset![2]); } #[test] fn test_delete_basic() { - let query = "UPDATE sharded SET id = $1 WHERE email = $2"; - let result = run_test(query).unwrap().unwrap(); + let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2") + .unwrap() + .unwrap(); - assert_equivalent_sql(&result.delete.stmt, "DELETE FROM sharded WHERE email = $1"); - #[cfg(feature = "new_parser")] - assert_all_bind_params_used(query, &result.delete.stmt); + assert_eq!(result.delete.stmt, "DELETE FROM sharded WHERE email = $1"); assert!(result.sharded_table(&[]).is_none()); assert!( @@ -1280,90 +1242,79 @@ mod test { #[test] fn test_delete_multiple_where_params() { - let query = "UPDATE sharded SET id = $1 WHERE email = $2 AND name = $3"; - let result = run_test(query).unwrap().unwrap(); + let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2 AND name = $3") + .unwrap() + .unwrap(); - assert_equivalent_sql( - &result.delete.stmt, - "DELETE FROM sharded WHERE email = $1 AND name = $2", + assert_eq!( + result.delete.stmt, + "DELETE FROM sharded WHERE email = $1 AND name = $2" ); - #[cfg(feature = "new_parser")] - assert_all_bind_params_used(query, &result.delete.stmt); } #[test] fn test_no_params_in_where() { - let query = "UPDATE sharded SET id = $1 WHERE email = 'test@example.com'"; - let result = run_test(query).unwrap().unwrap(); + let result = run_test("UPDATE sharded SET id = $1 WHERE email = 'test@example.com'") + .unwrap() + .unwrap(); - assert_equivalent_sql( - &result.select.stmt, - "SELECT * FROM sharded WHERE email = 'test@example.com'", + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE email = 'test@example.com'" ); - #[cfg(feature = "new_parser")] - assert_all_bind_params_used(query, &result.select.stmt); - #[cfg(not(feature = "new_parser"))] - assert_eq!(result.select.params, Vec::::new()); + assert!(result.select.params.is_empty()); } #[test] fn test_where_with_in_clause() { - let query = "UPDATE sharded SET id = $1 WHERE email IN ($2, $3, $4)"; - let result = run_test(query).unwrap().unwrap(); + let result = run_test("UPDATE sharded SET id = $1 WHERE email IN ($2, $3, $4)") + .unwrap() + .unwrap(); - assert_equivalent_sql( - &result.select.stmt, - "SELECT * FROM sharded WHERE email IN ($1, $2, $3)", + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE email IN ($1, $2, $3)" ); - #[cfg(feature = "new_parser")] - assert_all_bind_params_used(query, &result.select.stmt); - #[cfg(not(feature = "new_parser"))] - assert_eq!(result.select.params, vec![2, 3, 4]); + assert_eq!(result.select.params, indexset![2, 3, 4]); } #[test] fn test_where_with_comparison_operators() { - let query = "UPDATE sharded SET id = $1 WHERE count > $2 AND count < $3"; - let result = run_test(query).unwrap().unwrap(); + let result = run_test("UPDATE sharded SET id = $1 WHERE count > $2 AND count < $3") + .unwrap() + .unwrap(); - assert_equivalent_sql( - &result.select.stmt, - "SELECT * FROM sharded WHERE count > $1 AND count < $2", + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE count > $1 AND count < $2" ); - #[cfg(feature = "new_parser")] - assert_all_bind_params_used(query, &result.select.stmt); - #[cfg(not(feature = "new_parser"))] - assert_eq!(result.select.params, vec![2, 3]); + assert_eq!(result.select.params, indexset![2, 3]); } #[test] fn test_where_with_or_condition() { - let query = "UPDATE sharded SET id = $1 WHERE email = $2 OR name = $3"; - let result = run_test(query).unwrap().unwrap(); + let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2 OR name = $3") + .unwrap() + .unwrap(); - assert_equivalent_sql( - &result.select.stmt, - "SELECT * FROM sharded WHERE email = $1 OR name = $2", + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE email = $1 OR name = $2" ); - #[cfg(feature = "new_parser")] - assert_all_bind_params_used(query, &result.select.stmt); - #[cfg(not(feature = "new_parser"))] - assert_eq!(result.select.params, vec![2, 3]); + assert_eq!(result.select.params, indexset![2, 3]); } #[test] fn test_high_param_numbers() { - let query = "UPDATE sharded SET id = $10 WHERE email = $20 AND name = $30"; - let result = run_test(query).unwrap().unwrap(); + let result = run_test("UPDATE sharded SET id = $10 WHERE email = $20 AND name = $30") + .unwrap() + .unwrap(); - assert_equivalent_sql( - &result.select.stmt, - "SELECT * FROM sharded WHERE email = $1 AND name = $2", + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE email = $1 AND name = $2" ); - #[cfg(feature = "new_parser")] - assert_all_bind_params_used(query, &result.select.stmt); - #[cfg(not(feature = "new_parser"))] - assert_eq!(result.select.params, vec![20, 30]); + assert_eq!(result.select.params, indexset![20, 30]); } #[test] @@ -1375,139 +1326,121 @@ mod test { #[test] fn test_where_with_like() { - let query = "UPDATE sharded SET id = $1 WHERE email LIKE $2"; - let result = run_test(query).unwrap().unwrap(); + let result = run_test("UPDATE sharded SET id = $1 WHERE email LIKE $2") + .unwrap() + .unwrap(); - assert_equivalent_sql( - &result.select.stmt, - "SELECT * FROM sharded WHERE email LIKE $1", + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE email LIKE $1" ); - #[cfg(feature = "new_parser")] - assert_all_bind_params_used(query, &result.select.stmt); - #[cfg(not(feature = "new_parser"))] - assert_eq!(result.select.params, vec![2]); + assert_eq!(result.select.params, indexset![2]); } #[test] fn test_where_with_is_null() { - let query = "UPDATE sharded SET id = $1 WHERE email = $2 AND deleted_at IS NULL"; - let result = run_test(query).unwrap().unwrap(); + let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2 AND deleted_at IS NULL") + .unwrap() + .unwrap(); - assert_equivalent_sql( - &result.select.stmt, - "SELECT * FROM sharded WHERE email = $1 AND deleted_at IS NULL", + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE email = $1 AND deleted_at IS NULL" ); - #[cfg(feature = "new_parser")] - assert_all_bind_params_used(query, &result.select.stmt); - #[cfg(not(feature = "new_parser"))] - assert_eq!(result.select.params, vec![2]); + assert_eq!(result.select.params, indexset![2]); } #[test] fn test_where_with_between() { - let query = "UPDATE sharded SET id = $1 WHERE created_at BETWEEN $2 AND $3"; - let result = run_test(query).unwrap().unwrap(); + let result = run_test("UPDATE sharded SET id = $1 WHERE created_at BETWEEN $2 AND $3") + .unwrap() + .unwrap(); - assert_equivalent_sql( - &result.select.stmt, - "SELECT * FROM sharded WHERE created_at BETWEEN $1 AND $2", + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE created_at BETWEEN $1 AND $2" ); - #[cfg(feature = "new_parser")] - assert_all_bind_params_used(query, &result.select.stmt); - #[cfg(not(feature = "new_parser"))] - assert_eq!(result.select.params, vec![2, 3]); + assert_eq!(result.select.params, indexset![2, 3]); } #[test] fn test_same_param_used_twice() { // Same parameter $2 used twice in WHERE clause - let query = "UPDATE sharded SET id = $1 WHERE email = $2 OR name = $2"; - let result = run_test(query).unwrap().unwrap(); + let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2 OR name = $2") + .unwrap() + .unwrap(); // Both occurrences should be renumbered to $1 - assert_equivalent_sql( - &result.select.stmt, - "SELECT * FROM sharded WHERE email = $1 OR name = $1", + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE email = $1 OR name = $1" ); // Only one unique param in the mapping - #[cfg(feature = "new_parser")] - assert_all_bind_params_used(query, &result.select.stmt); - #[cfg(not(feature = "new_parser"))] - assert_eq!(result.select.params, vec![2]); + assert_eq!(result.select.params, indexset![2]); } #[test] fn test_same_param_used_multiple_times() { // $2 used three times - let query = "UPDATE sharded SET id = $1 WHERE a = $2 AND b = $2 AND c = $2"; - let result = run_test(query).unwrap().unwrap(); + let result = run_test("UPDATE sharded SET id = $1 WHERE a = $2 AND b = $2 AND c = $2") + .unwrap() + .unwrap(); - assert_equivalent_sql( - &result.select.stmt, - "SELECT * FROM sharded WHERE a = $1 AND b = $1 AND c = $1", + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE a = $1 AND b = $1 AND c = $1" ); - #[cfg(feature = "new_parser")] - assert_all_bind_params_used(query, &result.select.stmt); - #[cfg(not(feature = "new_parser"))] - assert_eq!(result.select.params, vec![2]); + assert_eq!(result.select.params, indexset![2]); } #[test] fn test_mixed_repeated_and_unique_params() { // $2 used twice, $3 used once - let query = "UPDATE sharded SET id = $1 WHERE a = $2 AND b = $3 AND c = $2"; - let result = run_test(query).unwrap().unwrap(); + let result = run_test("UPDATE sharded SET id = $1 WHERE a = $2 AND b = $3 AND c = $2") + .unwrap() + .unwrap(); - assert_equivalent_sql( - &result.select.stmt, - "SELECT * FROM sharded WHERE a = $1 AND b = $2 AND c = $1", + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE a = $1 AND b = $2 AND c = $1" ); - #[cfg(feature = "new_parser")] - assert_all_bind_params_used(query, &result.select.stmt); - #[cfg(not(feature = "new_parser"))] - assert_eq!(result.select.params, vec![2, 3]); + assert_eq!(result.select.params, indexset![2, 3]); } #[test] fn test_repeated_params_in_in_clause() { // Same param repeated in IN clause (unusual but valid) - let query = "UPDATE sharded SET id = $1 WHERE email IN ($2, $3, $2)"; - let result = run_test(query).unwrap().unwrap(); + let result = run_test("UPDATE sharded SET id = $1 WHERE email IN ($2, $3, $2)") + .unwrap() + .unwrap(); - assert_equivalent_sql( - &result.select.stmt, - "SELECT * FROM sharded WHERE email IN ($1, $2, $1)", + assert_eq!( + result.select.stmt, + "SELECT * FROM sharded WHERE email IN ($1, $2, $1)" ); - #[cfg(feature = "new_parser")] - assert_all_bind_params_used(query, &result.select.stmt); - #[cfg(not(feature = "new_parser"))] - assert_eq!(result.select.params, vec![2, 3]); + assert_eq!(result.select.params, indexset![2, 3]); } #[test] fn test_delete_with_repeated_params() { - let query = "UPDATE sharded SET id = $1 WHERE email = $2 OR name = $2"; - let result = run_test(query).unwrap().unwrap(); + let result = run_test("UPDATE sharded SET id = $1 WHERE email = $2 OR name = $2") + .unwrap() + .unwrap(); - assert_equivalent_sql( - &result.delete.stmt, - "DELETE FROM sharded WHERE email = $1 OR name = $1", + assert_eq!( + result.delete.stmt, + "DELETE FROM sharded WHERE email = $1 OR name = $1" ); - #[cfg(feature = "new_parser")] - assert_all_bind_params_used(query, &result.delete.stmt); - #[cfg(not(feature = "new_parser"))] - assert_eq!(result.delete.params, vec![2]); + assert_eq!(result.delete.params, indexset![2]); } #[test] fn test_sharding_key_not_changed() { - let query = "UPDATE sharded SET id = $1 WHERE id = $1 AND email = $2"; - let result = run_test(query).unwrap().unwrap(); - assert_equivalent_sql(&result.check.stmt, "SELECT * FROM sharded WHERE id = $1"); - #[cfg(feature = "new_parser")] - assert_all_bind_params_used(query, &result.check.stmt); - #[cfg(not(feature = "new_parser"))] - assert_eq!(result.check.params, vec![1]); + let result = run_test("UPDATE sharded SET id = $1 WHERE id = $1 AND email = $2") + .unwrap() + .unwrap(); + assert_eq!(result.check.stmt, "SELECT * FROM sharded WHERE id = $1"); + assert_eq!(result.check.params, indexset![1]); } #[test] @@ -1736,35 +1669,4 @@ mod test { stmt ); } - - const BIND_PARAM_PATTERN: &str = r#"(\$\d+)"#; - - fn bind_param_regex() -> Regex { - Regex::new(BIND_PARAM_PATTERN).unwrap() - } - - fn assert_equivalent_sql(actual: &str, expected: &str) { - let pattern = bind_param_regex().replace_all(expected, "BIND_PARAM"); - let pattern = regex::escape(&pattern).replace("BIND_PARAM", BIND_PARAM_PATTERN); - let regex = Regex::new(&pattern).unwrap(); - assert!( - regex.is_match(actual), - "Expected {:?} to match {:?}", - actual, - pattern - ) - } - - #[cfg(feature = "new_parser")] - fn assert_all_bind_params_used(original: &str, rewritten: &str) { - for captures in bind_param_regex().captures_iter(original) { - let param = &captures[1]; - assert!( - rewritten.contains(param), - "Expected {:?} to use param {}", - rewritten, - param - ); - } - } } From 574b297b9f44374137bb1934ee70f608bd16d823 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Sun, 5 Jul 2026 14:37:37 -0600 Subject: [PATCH 4/9] Use parser instead of string construction --- Cargo.lock | 2 +- pgdog/Cargo.toml | 2 +- .../router/parser/rewrite/statement/update.rs | 159 +++++++----------- 3 files changed, 65 insertions(+), 98 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 22db4f3d8..cf52ea477 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3299,7 +3299,7 @@ dependencies = [ [[package]] name = "pg_raw_parse" version = "0.1.0" -source = "git+https://github.com/pgdogdev/pg_raw_parse.git?rev=7a7ace5#7a7ace52b1fbe279d5a30e6d7267077b5a298990" +source = "git+https://github.com/pgdogdev/pg_raw_parse.git?rev=34e5987#34e598782a612d2dc3d894264602240232870fb5" dependencies = [ "bindgen 0.72.1", "cc", diff --git a/pgdog/Cargo.toml b/pgdog/Cargo.toml index 473e24f1e..195619ed8 100644 --- a/pgdog/Cargo.toml +++ b/pgdog/Cargo.toml @@ -80,7 +80,7 @@ smallvec = "1" reqwest = { version = "0.12", default-features = false, features = ["rustls-tls-webpki-roots-no-provider", "json"] } hex = "0.4" x509-parser = "0.18" -pg_raw_parse = { git = "https://github.com/pgdogdev/pg_raw_parse.git", rev = "7a7ace5", optional = true } +pg_raw_parse = { git = "https://github.com/pgdogdev/pg_raw_parse.git", rev = "34e5987", optional = true } itertools = "0.15.0" [target.'cfg(unix)'.dependencies] diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs index de08c9895..9b3bc6e69 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs @@ -14,7 +14,9 @@ use pg_query::{ }, }; #[cfg(feature = "new_parser")] -use pg_raw_parse::{DeparseResult, Node, NodeMut, Owned, deparse, make::owned, nodes, walk}; +use pg_raw_parse::make::{owned, try_owned}; +#[cfg(feature = "new_parser")] +use pg_raw_parse::{DeparseResult, Node, NodeMut, Owned, deparse, nodes, walk}; #[cfg(not(feature = "new_parser"))] use pgdog_config::QueryParserEngine; use pgdog_config::RewriteMode; @@ -194,85 +196,68 @@ impl Inner { let params = request.parameters()?; let mut bind = Bind::new_statement(""); - let mut columns = Vec::new(); - let mut values = Vec::new(); - let mut columns_str = Vec::new(); - let mut values_str = Vec::new(); - - let mut bind_idx = 1; - for (row_idx, field) in row_description.iter().enumerate() { - columns_str.push(format!(r#""{}""#, field.name.replace("\"", "\"\""))); // Escape " - - if let Some(value) = self.from_update.targetList().iter().find_map(|rt| { - if rt.name() == Some(&*field.name) { - Some(rt.val()) - } else { - None - } - }) { - let Ok(value) = Value::try_from(value) else { - values_str.push( - deparse_expr([value])? - .as_str() - .trim_start_matches("SELECT ") - .to_owned(), - ); - continue; - }; - - values_str.push(format!("${bind_idx}")); - match value { - Value::Placeholder(number) => { - let param = params - .as_ref() - .expect("param") - .parameter(number as usize - 1)? - .ok_or(Error::MissingParameter(number as u16))?; - bind.push_param(param.parameter().clone(), param.format()) - } + let returning_list = self.from_update.returningList(); - Value::Integer(int) => { - bind.push_param(Parameter::new(int.to_string().as_bytes()), Format::Text) + let insert = try_owned(|mem| -> Result<_, Error> { + let mut columns = Vec::new(); + let mut values = Vec::new(); + for (idx, field) in row_description.iter().enumerate() { + columns.push(&*field.name); + if let Some(value) = self.from_update.targetList().iter().find_map(|rt| { + if rt.name() == Some(&*field.name) { + Some(rt.val()) + } else { + None } - - Value::String(s) => bind.push_param(Parameter::new(s.as_bytes()), Format::Text), - - Value::Float(f) => { - bind.push_param(Parameter::new(f.to_string().as_bytes()), Format::Text) + }) { + let Ok(value) = Value::try_from(value) else { + values.push(mem.make_unique(value)); + continue; + }; + + match value { + Value::Placeholder(number) => { + let param = params + .as_ref() + .expect("param") + .parameter(number as usize - 1)? + .ok_or(Error::MissingParameter(number as u16))?; + bind.push_param(param.parameter().clone(), param.format()) + } + + Value::Integer(int) => bind + .push_param(Parameter::new(int.to_string().as_bytes()), Format::Text), + + Value::String(s) => { + bind.push_param(Parameter::new(s.as_bytes()), Format::Text) + } + + Value::Float(f) => { + bind.push_param(Parameter::new(f.to_string().as_bytes()), Format::Text) + } + + Value::Boolean(b) => bind.push_param( + Parameter::new(if b { "t".as_bytes() } else { "f".as_bytes() }), + Format::Text, + ), + + Value::Vector(vec) => bind + .push_param(Parameter::new(&vec.encode(Format::Text)?), Format::Text), + + Value::Null => bind.push_param(Parameter::new_null(), Format::Text), } + } else { + let value = data_row.get_raw(idx).ok_or(Error::MissingColumn(idx))?; - Value::Boolean(b) => bind.push_param( - Parameter::new(if b { "t".as_bytes() } else { "f".as_bytes() }), - Format::Text, - ), - - Value::Vector(vec) => { - bind.push_param(Parameter::new(&vec.encode(Format::Text)?), Format::Text) + if value.is_null() { + bind.push_param(Parameter::new_null(), Format::Text); + } else { + bind.push_param(Parameter::new(value), Format::Text); } - - Value::Null => bind.push_param(Parameter::new_null(), Format::Text), - } - } else { - let value = data_row - .get_raw(row_idx) - .ok_or(Error::MissingColumn(row_idx))?; - - if value.is_null() { - bind.push_param(Parameter::new_null(), Format::Text); - } else { - bind.push_param(Parameter::new(value), Format::Text); } - - values_str.push(format!("${bind_idx}")); + values.push(mem.make_ParamRef(bind.params_raw().len() as _).uncast()); } - columns.push(&*field.name); - values.push(bind_idx); - bind_idx += 1; - } - - let returning_list = self.from_update.returningList(); - let insert = owned(|mem| { let mut insert = mem.make_node::(); insert .as_mut() @@ -286,10 +271,6 @@ impl Inner { .collect::>(); insert.as_mut().set_cols(mem.make_List(&cols)); let mut select = mem.make_node::(); - let values = values - .into_iter() - .map(|v| mem.make_ParamRef(v).uncast()) - .collect::>(); select .as_mut() .set_valuesLists(mem.make_List(&[mem.make_List(&values)])); @@ -297,32 +278,18 @@ impl Inner { insert .as_mut() .set_returningList(mem.make_unique(returning_list)); - mem.make_List(&[mem.make_RawStmt(insert.uncast())]) - }); - - let stmt = format!( - "INSERT INTO {} ({}) VALUES ({}){}", - self.target_table(), - columns_str.join(", "), - values_str.join(", "), - if returning_list.is_empty() { - String::from("") - } else { - format!( - " RETURNING {}", - deparse_expr(returning_list)? - .as_str() - .trim_start_matches("SELECT ") - ) - } - ); + Ok(mem.make_List(&[mem.make_RawStmt(insert.uncast())])) + })?; + dbg!(&insert); + let stmt = deparse(insert.first().unwrap())?; + dbg!(stmt.as_str()); //// Build the AST to be used with the router. //// It's identical to the string-generated statement above. let ast = Ast::from_raw_stmts(insert); let mut req = ClientRequest::from(vec![ - ProtocolMessage::from(Parse::new_anonymous(&stmt)), + ProtocolMessage::from(Parse::new_anonymous(stmt.as_str())), Describe::new_statement("").into(), // So we get both T and t, bind.into(), Execute::new().into(), From 84b62a6073a10c30438a5dd26dba4b21721b4b97 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Sun, 5 Jul 2026 14:55:22 -0600 Subject: [PATCH 5/9] Fix tests on old parser --- pgdog/src/frontend/router/parser/rewrite/statement/update.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs index 9b3bc6e69..0afb0792c 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs @@ -1064,8 +1064,8 @@ mod test { #[cfg(not(feature = "new_parser"))] macro_rules! indexset { - ($t:tt) => { - vec![$t] + ($($t:tt)*) => { + vec![$($t)*] }; } From 62258934fa771e42ade1318bc9aee3258ee567aa Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Sun, 5 Jul 2026 15:11:13 -0600 Subject: [PATCH 6/9] Don't attempt to convert literals in updates to binds We're not caching the prepared statement, we're always using text format anyway, there's just no reason to do anything other than put the exact nodes we got in for the output --- .../router/parser/rewrite/statement/update.rs | 52 +++++-------------- 1 file changed, 13 insertions(+), 39 deletions(-) diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs index 0afb0792c..1e0fbcc90 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs @@ -33,8 +33,8 @@ use crate::{ }, }, net::{ - Bind, DataRow, Describe, Execute, Flush, Format, FromDataType, Parse, ProtocolMessage, - Query, RowDescription, Sync, bind::Parameter, + Bind, DataRow, Describe, Execute, Flush, Format, Parse, ProtocolMessage, Query, + RowDescription, Sync, bind::Parameter, }, }; @@ -203,6 +203,7 @@ impl Inner { let mut values = Vec::new(); for (idx, field) in row_description.iter().enumerate() { columns.push(&*field.name); + if let Some(value) = self.from_update.targetList().iter().find_map(|rt| { if rt.name() == Some(&*field.name) { Some(rt.val()) @@ -210,41 +211,14 @@ impl Inner { None } }) { - let Ok(value) = Value::try_from(value) else { + if let Ok(Value::Placeholder(number)) = Value::try_from(value) { + let param = params + .and_then(|p| p.parameter(number as usize - 1).transpose()) + .ok_or(Error::MissingParameter(number as u16))??; + bind.push_param(param.parameter().clone(), param.format()); + values.push(mem.make_ParamRef(bind.params_raw().len() as _).uncast()); + } else { values.push(mem.make_unique(value)); - continue; - }; - - match value { - Value::Placeholder(number) => { - let param = params - .as_ref() - .expect("param") - .parameter(number as usize - 1)? - .ok_or(Error::MissingParameter(number as u16))?; - bind.push_param(param.parameter().clone(), param.format()) - } - - Value::Integer(int) => bind - .push_param(Parameter::new(int.to_string().as_bytes()), Format::Text), - - Value::String(s) => { - bind.push_param(Parameter::new(s.as_bytes()), Format::Text) - } - - Value::Float(f) => { - bind.push_param(Parameter::new(f.to_string().as_bytes()), Format::Text) - } - - Value::Boolean(b) => bind.push_param( - Parameter::new(if b { "t".as_bytes() } else { "f".as_bytes() }), - Format::Text, - ), - - Value::Vector(vec) => bind - .push_param(Parameter::new(&vec.encode(Format::Text)?), Format::Text), - - Value::Null => bind.push_param(Parameter::new_null(), Format::Text), } } else { let value = data_row.get_raw(idx).ok_or(Error::MissingColumn(idx))?; @@ -254,8 +228,8 @@ impl Inner { } else { bind.push_param(Parameter::new(value), Format::Text); } + values.push(mem.make_ParamRef(bind.params_raw().len() as _).uncast()); } - values.push(mem.make_ParamRef(bind.params_raw().len() as _).uncast()); } let mut insert = mem.make_node::(); @@ -280,9 +254,7 @@ impl Inner { .set_returningList(mem.make_unique(returning_list)); Ok(mem.make_List(&[mem.make_RawStmt(insert.uncast())])) })?; - dbg!(&insert); let stmt = deparse(insert.first().unwrap())?; - dbg!(stmt.as_str()); //// Build the AST to be used with the router. //// It's identical to the string-generated statement above. @@ -1595,6 +1567,7 @@ mod test { Field::bigint("id"), Field::text("email"), Field::text("other_col"), + Field::text("other_other_col"), ]); // Create a mock data row with values for columns not in the UPDATE SET clause @@ -1602,6 +1575,7 @@ mod test { data_row.add("1"); // id - will be overwritten by mapping data_row.add("old@example.com"); // email - will be overwritten by mapping data_row.add("other_value"); // other_col - from existing row + data_row.add("other_other_value"); // other_other_col - from existing row // Create a simple query request (not prepared statement) let request = ClientRequest::from(vec![ProtocolMessage::from(Query::new( From 86cc54b5eca4c01fb250bdbe52e8ff2b414c560e Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Sun, 5 Jul 2026 15:17:39 -0600 Subject: [PATCH 7/9] Minor cleanup --- .../router/parser/rewrite/statement/update.rs | 19 +++++++------------ 1 file changed, 7 insertions(+), 12 deletions(-) diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs index 1e0fbcc90..4893fe244 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs @@ -194,15 +194,16 @@ impl Inner { data_row: &DataRow, ) -> Result { let params = request.parameters()?; - let mut bind = Bind::new_statement(""); - let returning_list = self.from_update.returningList(); let insert = try_owned(|mem| -> Result<_, Error> { let mut columns = Vec::new(); let mut values = Vec::new(); for (idx, field) in row_description.iter().enumerate() { - columns.push(&*field.name); + columns.push( + mem.make_ResTarget(Some(&*field.name), mem.empty(), mem.none()) + .uncast(), + ); if let Some(value) = self.from_update.targetList().iter().find_map(|rt| { if rt.name() == Some(&*field.name) { @@ -221,6 +222,7 @@ impl Inner { values.push(mem.make_unique(value)); } } else { + // This column wasn't changed, get the value from the select let value = data_row.get_raw(idx).ok_or(Error::MissingColumn(idx))?; if value.is_null() { @@ -236,14 +238,7 @@ impl Inner { insert .as_mut() .set_relation(mem.make_unique(self.from_update.relation())); - let cols = columns - .into_iter() - .map(|c| { - mem.make_ResTarget(Some(c), mem.empty(), mem.none()) - .uncast() - }) - .collect::>(); - insert.as_mut().set_cols(mem.make_List(&cols)); + insert.as_mut().set_cols(mem.make_List(&columns)); let mut select = mem.make_node::(); select .as_mut() @@ -251,7 +246,7 @@ impl Inner { insert.as_mut().set_selectStmt(select.uncast()); insert .as_mut() - .set_returningList(mem.make_unique(returning_list)); + .set_returningList(mem.make_unique(self.from_update.returningList())); Ok(mem.make_List(&[mem.make_RawStmt(insert.uncast())])) })?; let stmt = deparse(insert.first().unwrap())?; From f82004a488dc22fad5d0f0de0dfc40bf80553293 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Sun, 5 Jul 2026 15:27:55 -0600 Subject: [PATCH 8/9] Fix cfg'd uses --- pgdog/src/frontend/router/parser/rewrite/statement/update.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs index 4893fe244..b9cc5bf7c 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs @@ -23,6 +23,8 @@ use pgdog_config::RewriteMode; #[cfg(not(feature = "new_parser"))] use crate::frontend::router::parser::rewrite::statement::visitor::visit_and_mutate_nodes; +#[cfg(not(feature = "new_parser"))] +use crate::net::FromDataType; use crate::{ frontend::{ BufferedQuery, ClientRequest, From 8bcdef1a3d32e2482c57c87afdc1f3b2d95f8f1c Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Sun, 5 Jul 2026 15:30:55 -0600 Subject: [PATCH 9/9] Fix test --- pgdog/src/frontend/router/parser/rewrite/statement/update.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs index b9cc5bf7c..c2e0a320d 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs @@ -1564,6 +1564,7 @@ mod test { Field::bigint("id"), Field::text("email"), Field::text("other_col"), + #[cfg(feature = "new_parser")] Field::text("other_other_col"), ]); @@ -1572,6 +1573,7 @@ mod test { data_row.add("1"); // id - will be overwritten by mapping data_row.add("old@example.com"); // email - will be overwritten by mapping data_row.add("other_value"); // other_col - from existing row + #[cfg(feature = "new_parser")] data_row.add("other_other_value"); // other_other_col - from existing row // Create a simple query request (not prepared statement)