diff --git a/Cargo.lock b/Cargo.lock index 3ae92a7a1..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=a4a2a2c#a4a2a2ce04ae22056fbcf7f65194412c216a2140" +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 8f58ad884..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 = "a4a2a2c", 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] 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..e8ff52458 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) -> 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"), + } + } + _ => 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}"