From 5f66ce4d03f42d83218da7048f41e5b9a76754bf Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Fri, 10 Jul 2026 07:46:23 -0600 Subject: [PATCH 1/3] Port StatementRewrite::rewrite_unique_id This is mostly a direct port, other than lifting where we get the UniqueId generator up a level so `rewrite_unique_id` can be infallible. (NB: The only reason getting the generator can fail is if the config is invalid. We should be checking that config way earlier, and fetching the generator should be infallible). This marks the final function in `StatementRewrite` that needed to be ported. We can now begin the structural change of removing the old AST from that struct entirely. However, that will be a large enough change on its own to warrant a separate PR, so as of this commit we still do the deparse/reparse/deparse/reparse dance. --- Cargo.lock | 2 +- pgdog/Cargo.toml | 2 +- .../rewrite/statement/aggregate/engine.rs | 4 +- .../parser/rewrite/statement/aggregate/mod.rs | 2 +- .../router/parser/rewrite/statement/mod.rs | 63 ++++--- .../parser/rewrite/statement/unique_id.rs | 173 ++++++++++++++---- 6 files changed, 173 insertions(+), 73 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 3ae92a7a1..1abf1017b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3313,7 +3313,7 @@ dependencies = [ [[package]] name = "pg_raw_parse" version = "0.1.0" -source = "git+https://github.com/pgdogdev/pg_raw_parse.git?rev=a4a2a2c#a4a2a2ce04ae22056fbcf7f65194412c216a2140" +source = "git+https://github.com/pgdogdev/pg_raw_parse.git?rev=fb9dc29#fb9dc29f494833ab39f1f51f2bdba3250801274d" dependencies = [ "bindgen 0.72.1", "cc", diff --git a/pgdog/Cargo.toml b/pgdog/Cargo.toml index 8f58ad884..1f0747aa1 100644 --- a/pgdog/Cargo.toml +++ b/pgdog/Cargo.toml @@ -80,7 +80,7 @@ smallvec = "1" reqwest.workspace = true hex = "0.4" x509-parser = "0.18" -pg_raw_parse = { git = "https://github.com/pgdogdev/pg_raw_parse.git", rev = "a4a2a2c", optional = true } +pg_raw_parse = { git = "https://github.com/pgdogdev/pg_raw_parse.git", rev = "fb9dc29", optional = true } itertools = "0.15.0" [target.'cfg(unix)'.dependencies] diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/engine.rs b/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/engine.rs index 4b2eb7d59..a4a1f6bc6 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/engine.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/engine.rs @@ -26,7 +26,7 @@ impl AggregatesRewrite { /// Rewrite a SELECT query in-place, adding helper aggregates when necessary. #[cfg(feature = "new_parser")] pub(crate) fn rewrite_select<'a>( - mut select: nodes::SelectStmtMut<'a, '_>, + select: &mut nodes::SelectStmtMut<'a, '_>, mem: make::MemoryToken<'a>, aggregate: &Aggregate, ) -> RewriteOutput { @@ -455,7 +455,7 @@ mod tests { let aggregate = Aggregate::parse(&*stmt, &Default::default()); output = Some(AggregatesRewrite::rewrite_select( - stmt.as_mut(), + &mut stmt.as_mut(), mem, &aggregate, )); diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/mod.rs b/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/mod.rs index 412b83a10..ca74a61ef 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/mod.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/mod.rs @@ -17,7 +17,7 @@ impl StatementRewrite<'_> { #[cfg(feature = "new_parser")] pub(super) fn rewrite_aggregates<'a>( &mut self, - select: SelectStmtMut<'a, '_>, + select: &mut SelectStmtMut<'a, '_>, mem: MemoryToken<'a>, plan: &mut RewritePlan, schema: &Schema, diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs index a0caf845f..830d99f74 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs @@ -1,9 +1,10 @@ //! Statement rewriter. +#[cfg(not(feature = "new_parser"))] use pg_query::Node as PgNode; use pg_query::protobuf::ParseResult; #[cfg(feature = "new_parser")] -use pg_raw_parse::{Node, NodeMut, make, walk}; +use pg_raw_parse::{Node, NodeMut, make, transform, walk}; #[cfg(not(feature = "new_parser"))] use pgdog_config::QueryParserEngine; @@ -142,47 +143,45 @@ impl<'a> StatementRewrite<'a> { // Track the next parameter number to use let mut next_param = plan.params as i32 + 1; - - self.stmt.stmts[0] = pg_query::parse(pg_raw_parse::deparse(node.as_ref())?.as_str())? - .protobuf - .stmts - .pop() - .unwrap(); - let extended = self.extended; - visitor::visit_and_mutate_nodes(self.stmt, |node| -> Result, Error> { - match Self::rewrite_unique_id(node, extended, &mut next_param)? { - Some(replacement) => { - plan.unique_ids += 1; - self.rewritten = true; - Ok(Some(replacement)) - } - None => Ok(None), + let generator = crate::unique_id::UniqueId::generator()?; + transform::transform(&mut node, |node| { + if let Some(replacement) = Self::rewrite_unique_id( + node.as_ref(), + mem, + generator, + self.extended, + &mut next_param, + ) { + plan.unique_ids += 1; + self.rewritten = true; + node.replace(replacement); + None + } else { + Some(node) } - })?; - - let reparsed = pg_raw_parse::parse(&pg_query::deparse(self.stmt)?)?; - let mut node = reparsed.stmts().next().unwrap(); - let mut select_stmt; - if let Node::SelectStmt(select) = node { - select_stmt = mem.make_unique(select); - self.rewrite_aggregates(select_stmt.as_mut(), mem, &mut plan, self.db_schema)?; - self.limit_offset(&select_stmt, &mut plan); - node = (&*select_stmt).into(); - self.stmt.stmts = pg_query::parse(pg_raw_parse::deparse(node)?.as_str())? - .protobuf - .stmts; + }); + + if let NodeMut::SelectStmt(mut select) = node.as_mut() { + self.rewrite_aggregates(&mut select, mem, &mut plan, self.db_schema)?; + self.limit_offset(&select, &mut plan); } if self.rewritten { - plan.stmt = Some(pg_raw_parse::deparse(node)?.as_str().to_owned()); + plan.stmt = Some(pg_raw_parse::deparse(node.as_ref())?.as_str().to_owned()); } self.split_insert(&mut plan)?; - if let Node::UpdateStmt(stmt) = node { + if let Node::UpdateStmt(stmt) = node.as_ref() { self.sharding_key_update(stmt, &mut plan)?; } - Ok::<_, Error>(mem.make_raw_stmt(mem.make_unique(node))) + self.stmt.stmts[0] = pg_query::parse(pg_raw_parse::deparse(node.as_ref())?.as_str())? + .protobuf + .stmts + .pop() + .unwrap(); + + Ok::<_, Error>(mem.make_raw_stmt(node)) })?; Ok(plan) diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/unique_id.rs b/pgdog/src/frontend/router/parser/rewrite/statement/unique_id.rs index 9638bd2db..b064a79c5 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/unique_id.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/unique_id.rs @@ -1,35 +1,83 @@ +use super::StatementRewrite; +#[cfg(feature = "new_parser")] +use crate::unique_id::UniqueId; +#[cfg(not(feature = "new_parser"))] use pg_query::protobuf::{AConst, ParamRef, String as PgString, TypeCast, TypeName, a_const::Val}; +#[cfg(not(feature = "new_parser"))] use pg_query::{Node, NodeEnum}; - -use super::StatementRewrite; +#[cfg(feature = "new_parser")] +use pg_raw_parse::{Node, make}; impl StatementRewrite<'_> { /// Attempt to rewrite a pgdog.unique_id() call. /// /// Returns `Ok(Some(replacement_node))` if the node is a unique_id call, /// `Ok(None)` otherwise. Increments `next_param` when in extended mode. - pub(super) fn rewrite_unique_id( - node: &Node, + #[cfg(feature = "new_parser")] + pub(super) fn rewrite_unique_id<'mem>( + node: Node<'_>, + mem: make::MemoryToken<'mem>, + generator: &UniqueId, extended: bool, next_param: &mut i32, - ) -> Result, super::Error> { + ) -> Option>> { if !Self::is_unique_id(node) { - return Ok(None); + return None; } let replacement = if extended { - let param_num = *next_param; + let param_ref = mem.make_param_ref(*next_param); *next_param += 1; - Self::param_bigint(param_num) + param_ref.uncast() } else { - let unique_id = crate::unique_id::UniqueId::generator()?.next_id(); - Self::literal_bigint(unique_id) + use pg_raw_parse::ConstValue; + + let unique_id = generator.next_id(); + mem.make_a_const(ConstValue::Float(&unique_id.to_string())) + .uncast() }; - Ok(Some(replacement)) + Some( + mem.make_type_cast( + replacement, + mem.make_list(&[ + mem.make_string(Some("pg_catalog")), + mem.make_string(Some("int8")), + ]), + ) + .uncast(), + ) + } + + cfg_select! { + not(feature = "new_parser") => { + pub(super) fn rewrite_unique_id( + node: &Node, + extended: bool, + next_param: &mut i32, + ) -> Result, super::Error> { + if !Self::is_unique_id(node) { + return Ok(None); + } + + let replacement = if extended { + let param_num = *next_param; + *next_param += 1; + Self::param_bigint(param_num) + } else { + let unique_id = crate::unique_id::UniqueId::generator()?.next_id(); + Self::literal_bigint(unique_id) + }; + + Ok(Some(replacement)) + } + + } + _ => {} } /// Create a parameter reference cast to bigint: $N::bigint + #[cfg(not(feature = "new_parser"))] fn param_bigint(number: i32) -> Node { let param_ref = Node { node: Some(NodeEnum::ParamRef(ParamRef { @@ -48,6 +96,7 @@ impl StatementRewrite<'_> { } /// Create a literal value cast to bigint: ::bigint + #[cfg(not(feature = "new_parser"))] fn literal_bigint(value: i64) -> Node { let literal = Node { node: Some(NodeEnum::AConst(AConst { @@ -68,6 +117,7 @@ impl StatementRewrite<'_> { } /// Create a TypeName for bigint (int8). + #[cfg(not(feature = "new_parser"))] fn bigint_type() -> TypeName { TypeName { names: vec![ @@ -87,27 +137,44 @@ impl StatementRewrite<'_> { } /// Check if a node is a function call to pgdog.unique_id(). - fn is_unique_id(node: &Node) -> bool { - let Some(NodeEnum::FuncCall(func)) = &node.node else { + #[cfg(feature = "new_parser")] + fn is_unique_id(node: Node<'_>) -> bool { + let Node::FuncCall(func) = node else { return false; }; - // Must have exactly 2 parts: schema "pgdog" and function "unique_id" - if func.funcname.len() != 2 { - return false; - } + func.funcname() + .iter() + .filter_map(Node::as_str) + .eq(["pgdog", "unique_id"]) + } - let schema = func.funcname.first().and_then(|n| match &n.node { - Some(NodeEnum::String(s)) => Some(s.sval.as_str()), - _ => None, - }); + cfg_select! { + not(feature = "new_parser") => { + fn is_unique_id(node: &Node) -> bool { + let Some(NodeEnum::FuncCall(func)) = &node.node else { + return false; + }; - let name = func.funcname.get(1).and_then(|n| match &n.node { - Some(NodeEnum::String(s)) => Some(s.sval.as_str()), - _ => None, - }); + // Must have exactly 2 parts: schema "pgdog" and function "unique_id" + if func.funcname.len() != 2 { + return false; + } + + let schema = func.funcname.first().and_then(|n| match &n.node { + Some(NodeEnum::String(s)) => Some(s.sval.as_str()), + _ => None, + }); - matches!((schema, name), (Some("pgdog"), Some("unique_id"))) + let name = func.funcname.get(1).and_then(|n| match &n.node { + Some(NodeEnum::String(s)) => Some(s.sval.as_str()), + _ => None, + }); + + matches!((schema, name), (Some("pgdog"), Some("unique_id"))) + } + } + _ => {} } } @@ -123,6 +190,8 @@ mod tests { use crate::frontend::router::parser::rewrite::statement::RewritePlan; use crate::test_utils::set_env_var; use pg_query::protobuf::ParseResult; + #[cfg(feature = "new_parser")] + use pg_raw_parse::{Owned, nodes}; fn default_schema() -> ShardingSchema { ShardingSchema { @@ -139,48 +208,79 @@ mod tests { Schema::default() } - fn parse_first_target(sql: &str) -> Node { - let ast = pg_query::parse(sql).unwrap(); - let stmt = ast.protobuf.stmts.first().unwrap().stmt.as_ref().unwrap(); - match &stmt.node { - Some(NodeEnum::SelectStmt(select)) => { - let res_target = select.target_list.first().unwrap(); - match &res_target.node { - Some(NodeEnum::ResTarget(res)) => *res.val.as_ref().unwrap().clone(), - _ => panic!("expected ResTarget"), - } + #[cfg(feature = "new_parser")] + fn parse_first_target(sql: &str) -> Owned { + let ast = pg_raw_parse::parse(sql).unwrap(); + match ast.stmts().next().unwrap() { + Node::SelectStmt(select) => { + make::owned(|mem| mem.make_unique(select.target_list().first().unwrap())) } _ => panic!("expected SelectStmt"), } } + cfg_select! { + not(feature = "new_parser") => { + fn parse_first_target(sql: &str) -> Owned { + let ast = pg_query::parse(sql).unwrap(); + let stmt = ast.protobuf.stmts.first().unwrap().stmt.as_ref().unwrap(); + match &stmt.node { + Some(NodeEnum::SelectStmt(select)) => { + let res_target = select.target_list.first().unwrap(); + match &res_target.node { + Some(NodeEnum::ResTarget(res)) => *res.val.as_ref().unwrap().clone(), + _ => panic!("expected ResTarget"), + } + } + _ => panic!("expected SelectStmt"), + } + } + } + _ => {} + } + #[test] fn test_is_unique_id_qualified() { let node = parse_first_target("SELECT pgdog.unique_id()"); + #[cfg(feature = "new_parser")] + assert!(StatementRewrite::is_unique_id(node.val())); + #[cfg(not(feature = "new_parser"))] assert!(StatementRewrite::is_unique_id(&node)); } #[test] fn test_is_unique_id_unqualified() { let node = parse_first_target("SELECT unique_id()"); + #[cfg(feature = "new_parser")] + assert!(!StatementRewrite::is_unique_id(node.val())); + #[cfg(not(feature = "new_parser"))] assert!(!StatementRewrite::is_unique_id(&node)); } #[test] fn test_is_unique_id_wrong_schema() { let node = parse_first_target("SELECT other.unique_id()"); + #[cfg(feature = "new_parser")] + assert!(!StatementRewrite::is_unique_id(node.val())); + #[cfg(not(feature = "new_parser"))] assert!(!StatementRewrite::is_unique_id(&node)); } #[test] fn test_is_unique_id_wrong_function() { let node = parse_first_target("SELECT pgdog.other_func()"); + #[cfg(feature = "new_parser")] + assert!(!StatementRewrite::is_unique_id(node.val())); + #[cfg(not(feature = "new_parser"))] assert!(!StatementRewrite::is_unique_id(&node)); } #[test] fn test_is_unique_id_not_function() { let node = parse_first_target("SELECT 1"); + #[cfg(feature = "new_parser")] + assert!(!StatementRewrite::is_unique_id(node.val())); + #[cfg(not(feature = "new_parser"))] assert!(!StatementRewrite::is_unique_id(&node)); } @@ -221,6 +321,7 @@ mod tests { let sql = ast.deparse().unwrap(); // Value should be a bigint literal cast + #[cfg(not(feature = "new_parser"))] assert!( sql.contains("::bigint"), "Expected ::bigint cast, got: {sql}" From 97779b83d2760276cbe1c7415e011c439270f1c1 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Fri, 10 Jul 2026 07:58:31 -0600 Subject: [PATCH 2/3] Fix incorrectly changed line in old parser code --- pgdog/src/frontend/router/parser/rewrite/statement/unique_id.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/unique_id.rs b/pgdog/src/frontend/router/parser/rewrite/statement/unique_id.rs index b064a79c5..e8ff52458 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/unique_id.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/unique_id.rs @@ -221,7 +221,7 @@ mod tests { cfg_select! { not(feature = "new_parser") => { - fn parse_first_target(sql: &str) -> Owned { + fn parse_first_target(sql: &str) -> Node { let ast = pg_query::parse(sql).unwrap(); let stmt = ast.protobuf.stmts.first().unwrap().stmt.as_ref().unwrap(); match &stmt.node { From b57cf969bb044efcacebec5967b413f142e812d0 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Fri, 10 Jul 2026 09:18:31 -0600 Subject: [PATCH 3/3] bump parser version --- Cargo.lock | 2 +- pgdog/Cargo.toml | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1abf1017b..22520deab 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3313,7 +3313,7 @@ dependencies = [ [[package]] name = "pg_raw_parse" version = "0.1.0" -source = "git+https://github.com/pgdogdev/pg_raw_parse.git?rev=fb9dc29#fb9dc29f494833ab39f1f51f2bdba3250801274d" +source = "git+https://github.com/pgdogdev/pg_raw_parse.git?rev=b2a15ec#b2a15ec8a4a3a0a9a8d702901e6e336c63aec72b" dependencies = [ "bindgen 0.72.1", "cc", diff --git a/pgdog/Cargo.toml b/pgdog/Cargo.toml index 1f0747aa1..03d111e7f 100644 --- a/pgdog/Cargo.toml +++ b/pgdog/Cargo.toml @@ -80,7 +80,7 @@ smallvec = "1" reqwest.workspace = true hex = "0.4" x509-parser = "0.18" -pg_raw_parse = { git = "https://github.com/pgdogdev/pg_raw_parse.git", rev = "fb9dc29", optional = true } +pg_raw_parse = { git = "https://github.com/pgdogdev/pg_raw_parse.git", rev = "b2a15ec", optional = true } itertools = "0.15.0" [target.'cfg(unix)'.dependencies]