From b9479c6724d8abd158767dafd12e6802b7b51920 Mon Sep 17 00:00:00 2001 From: Sage Griffin Date: Mon, 13 Jul 2026 09:22:43 -0600 Subject: [PATCH] Remove the old parser's AST from `StatementRewrite` It is no longer used, so we can get rid of it. I opted to change the signature of `maybe_rewrite` to take a mutable reference to the statement rather than copy it and return a new one. This resulted in a significant indentation change, but no actual code changes (beyond replacing a few `as_ref` and `as_mut`s with `stmt` and `stmt_mut`. --- Cargo.lock | 2 +- pgdog/Cargo.toml | 2 +- pgdog/src/frontend/router/parser/cache/ast.rs | 34 +++++-- .../parser/rewrite/statement/auto_id.rs | 82 ++++++++-------- .../router/parser/rewrite/statement/insert.rs | 2 + .../router/parser/rewrite/statement/mod.rs | 91 +++++++++--------- .../router/parser/rewrite/statement/offset.rs | 2 + .../rewrite/statement/simple_prepared.rs | 56 +++++++---- .../parser/rewrite/statement/unique_id.rs | 94 ++++++++++--------- .../router/parser/rewrite/statement/update.rs | 2 + 10 files changed, 209 insertions(+), 158 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 1d982e3a8..1362318be 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=6421337#6421337b411e91b4b0ca553b8146108c2a273667" +source = "git+https://github.com/pgdogdev/pg_raw_parse.git?rev=6b424c4#6b424c4510ef4e9c482ef8a7fd7246b66f15de0a" dependencies = [ "bindgen 0.72.1", "cc", diff --git a/pgdog/Cargo.toml b/pgdog/Cargo.toml index 7f0f4ab26..fc724639f 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 = "6421337", optional = true } +pg_raw_parse = { git = "https://github.com/pgdogdev/pg_raw_parse.git", rev = "6b424c4", optional = true } itertools = "0.15.0" [target.'cfg(unix)'.dependencies] diff --git a/pgdog/src/frontend/router/parser/cache/ast.rs b/pgdog/src/frontend/router/parser/cache/ast.rs index 0351c3dcf..f8d9899b3 100644 --- a/pgdog/src/frontend/router/parser/cache/ast.rs +++ b/pgdog/src/frontend/router/parser/cache/ast.rs @@ -3,7 +3,7 @@ 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 pg_raw_parse::{Owned, StmtList, make}; use pgdog_config::QueryParserEngine; use std::fmt::Debug; use std::ops::Deref; @@ -108,19 +108,21 @@ impl Ast { search_path: Option<&ParameterValue>, ) -> Result { let now = Instant::now(); + #[cfg(not(feature = "new_parser"))] let mut ast = match schema.query_parser_engine { QueryParserEngine::PgQueryProtobuf => parse(query.query_without_comment), QueryParserEngine::PgQueryRaw => parse_raw(query.query_without_comment), } .map_err(Error::PgQuery)?; #[cfg(feature = "new_parser")] - let new_ast = pg_raw_parse::parse(query.query_without_comment).map_err(Error::Parse)?; + let ast = pg_raw_parse::parse(query.query_without_comment).map_err(Error::Parse)?; // Run the rewrite unconditionally. Even when a shard comment will // route the query to a specific shard, we need to know whether the // same query body (without the comment) would require a rewrite, so // `Cache::query` can decide whether this entry is safe to cache. let mut rewriter = StatementRewrite::new(StatementRewriteContext { + #[cfg(not(feature = "new_parser"))] stmt: &mut ast.protobuf, extended: query.original_query.extended(), prepared: query.original_query.prepared(), @@ -131,11 +133,18 @@ impl Ast { search_path, }); #[cfg(feature = "new_parser")] - let rewrite_plan = if let Some(node) = new_ast.stmts().next() { - rewriter.maybe_rewrite(node)? - } else { - Default::default() - }; + let mut rewrite_plan = Default::default(); + #[cfg(feature = "new_parser")] + let ast = make::try_owned(|mem| { + // FIXME(sage): We should have a parse function on mem so we don't + // need to parse and then copy the parsed tree just to throw the + // original away + let mut copy = mem.make_unique(&*ast.into_inner()); + if let Some(stmt) = copy.as_mut().into_iter().next() { + rewrite_plan = rewriter.maybe_rewrite(stmt, mem)?; + } + Ok::<_, Error>(copy) + })?; #[cfg(not(feature = "new_parser"))] let rewrite_plan = rewriter.maybe_rewrite()?; @@ -154,6 +163,10 @@ impl Ast { ); } + #[cfg(feature = "new_parser")] + let old_ast = + pg_query::parse(&pg_raw_parse::deparse_stmts(&*ast)?).map_err(Error::PgQuery)?; + Ok(Self { cached: true, comment_shard: None, @@ -162,9 +175,10 @@ impl Ast { inner: Arc::new(AstInner { stats: Mutex::new(stats), #[cfg(feature = "new_parser")] - new_ast: pg_raw_parse::parse(&ast.deparse().unwrap()) - .unwrap() - .into_inner(), + new_ast: ast, + #[cfg(feature = "new_parser")] + ast: old_ast, + #[cfg(not(feature = "new_parser"))] ast, rewrite_plan, fingerprint: OnceCell::new(), diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/auto_id.rs b/pgdog/src/frontend/router/parser/rewrite/statement/auto_id.rs index 1b32e147a..c342e4613 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/auto_id.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/auto_id.rs @@ -557,32 +557,8 @@ mod tests { db_schema: &Schema, mode: RewriteMode, ) -> Result<(String, RewritePlan), Error> { - let _guard = set_env_var("NODE_ID", "pgdog-1"); - let mut ast = pg_query::parse(sql).unwrap().protobuf; - #[cfg(feature = "new_parser")] - let stmt = pg_raw_parse::parse(sql).unwrap(); - let mut prepared = PreparedStatements::default(); let schema = sharding_schema_with_mode(mode); - let mut rewriter = StatementRewrite::new(StatementRewriteContext { - stmt: &mut ast, - extended: false, - prepared: false, - prepared_statements: &mut prepared, - schema: &schema, - db_schema, - user: "", - search_path: None, - }); - let plan = rewriter.maybe_rewrite( - #[cfg(feature = "new_parser")] - stmt.stmts().next().unwrap(), - )?; - let result = if plan.stmt.is_some() { - plan.stmt.clone().unwrap() - } else { - ast.deparse().unwrap() - }; - Ok((result, plan)) + rewrite_sql_with_sharding_schema(sql, db_schema, &schema) } #[test] @@ -768,18 +744,16 @@ mod tests { } } + #[cfg(feature = "new_parser")] fn rewrite_sql_with_sharding_schema( sql: &str, db_schema: &Schema, schema: &ShardingSchema, ) -> Result<(String, RewritePlan), Error> { let _guard = set_env_var("NODE_ID", "pgdog-1"); - let mut ast = pg_query::parse(sql).unwrap().protobuf; - #[cfg(feature = "new_parser")] - let new_ast = pg_raw_parse::parse(sql).unwrap(); + let ast = pg_raw_parse::parse(sql).unwrap(); let mut prepared = PreparedStatements::default(); let mut rewriter = StatementRewrite::new(StatementRewriteContext { - stmt: &mut ast, extended: false, prepared: false, prepared_statements: &mut prepared, @@ -788,16 +762,46 @@ mod tests { user: "", search_path: None, }); - let plan = rewriter.maybe_rewrite( - #[cfg(feature = "new_parser")] - new_ast.stmts().next().unwrap(), - )?; - let result = if plan.stmt.is_some() { - plan.stmt.clone().unwrap() - } else { - ast.deparse().unwrap() - }; - Ok((result, plan)) + let mut plan = Default::default(); + let ast = make::try_owned(|mem| { + let mut copy = mem.make_unique(&*ast.into_inner()); + plan = rewriter.maybe_rewrite(copy.as_mut().into_iter().next().unwrap(), mem)?; + Ok::<_, Error>(copy) + })?; + let sql = pg_raw_parse::deparse_stmts(&*ast)?; + Ok((sql, plan)) + } + + cfg_select! { + not(feature = "new_parser") => { + fn rewrite_sql_with_sharding_schema( + sql: &str, + db_schema: &Schema, + schema: &ShardingSchema, + ) -> Result<(String, RewritePlan), Error> { + let _guard = set_env_var("NODE_ID", "pgdog-1"); + let mut ast = pg_query::parse(sql).unwrap().protobuf; + let mut prepared = PreparedStatements::default(); + let mut rewriter = StatementRewrite::new(StatementRewriteContext { + stmt: &mut ast, + extended: false, + prepared: false, + prepared_statements: &mut prepared, + schema, + db_schema, + user: "", + search_path: None, + }); + let plan = rewriter.maybe_rewrite()?; + let result = if plan.stmt.is_some() { + plan.stmt.clone().unwrap() + } else { + ast.deparse().unwrap() + }; + Ok((result, plan)) + } + } + _ => {} } #[test] diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs b/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs index 0c1744502..f0db9c4e8 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/insert.rs @@ -464,6 +464,7 @@ mod tests { } fn parse_and_split(sql: &str) -> Vec { + #[cfg(not(feature = "new_parser"))] let mut ast = pg_query::parse(sql).unwrap().protobuf; #[cfg(feature = "new_parser")] let root = pg_raw_parse::parse(sql).unwrap(); @@ -476,6 +477,7 @@ mod tests { let schema = default_schema(); let db_schema = default_db_schema(); let mut rewriter = StatementRewrite::new(StatementRewriteContext { + #[cfg(not(feature = "new_parser"))] stmt: &mut ast, extended: false, prepared: false, diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs index e8e8fcf59..8d8421a87 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs @@ -2,9 +2,10 @@ #[cfg(not(feature = "new_parser"))] use pg_query::Node as PgNode; +#[cfg(not(feature = "new_parser"))] use pg_query::protobuf::ParseResult; #[cfg(feature = "new_parser")] -use pg_raw_parse::{Node, NodeMut, make, transform, walk}; +use pg_raw_parse::{Node, NodeMut, make, nodes, transform, walk}; #[cfg(not(feature = "new_parser"))] use pgdog_config::QueryParserEngine; @@ -35,6 +36,7 @@ pub(crate) use update::*; #[derive(Debug)] pub struct StatementRewriteContext<'a> { /// The AST of the statement we are rewriting. + #[cfg(not(feature = "new_parser"))] pub stmt: &'a mut ParseResult, /// The statement is using the extended protocol with placeholders. pub extended: bool, @@ -56,6 +58,7 @@ pub struct StatementRewriteContext<'a> { #[derive(Debug)] pub struct StatementRewrite<'a> { /// SQL statement. + #[cfg(not(feature = "new_parser"))] stmt: &'a mut ParseResult, /// The statement was rewritten. rewritten: bool, @@ -85,6 +88,7 @@ impl<'a> StatementRewrite<'a> { /// pub fn new(ctx: StatementRewriteContext<'a>) -> Self { Self { + #[cfg(not(feature = "new_parser"))] stmt: ctx.stmt, rewritten: false, extended: ctx.extended, @@ -110,14 +114,18 @@ impl<'a> StatementRewrite<'a> { /// Maybe rewrite the statement and produce a rewrite plan /// we can apply to Bind messages. #[cfg(feature = "new_parser")] - pub fn maybe_rewrite(&mut self, node: Node<'a>) -> Result { + pub fn maybe_rewrite<'mem>( + &mut self, + mut stmt: nodes::RawStmtMut<'mem, '_>, + mem: make::MemoryToken<'mem>, + ) -> Result { let mut plan = RewritePlan::default(); - match node { + match stmt.stmt() { Node::InsertStmt(_) | Node::SelectStmt(_) | Node::UpdateStmt(_) - | Node::DeleteStmt(_) => walk::walk(node, |node| match node { + | Node::DeleteStmt(_) => walk::walk(stmt.stmt(), |node| match node { Node::ParamRef(param) => plan.params = plan.params.max(param.number as u16), _ => {} }), @@ -126,27 +134,26 @@ impl<'a> StatementRewrite<'a> { _ => return Ok(plan), } - make::try_owned(|mem| { - let mut node = mem.make_unique(node); - - // Handle top-level PREPARE/EXECUTE statements. - let prepared_result = self.rewrite_simple_prepared(node.as_mut(), mem)?; - if prepared_result.rewritten { - self.rewritten = true; - plan.prepares = prepared_result.prepares; - } + // Handle top-level PREPARE/EXECUTE statements. + let prepared_result = self.rewrite_simple_prepared(stmt.stmt_mut(), mem)?; + if prepared_result.rewritten { + self.rewritten = true; + plan.prepares = prepared_result.prepares; + } - // Inject pgdog.unique_id() for missing BIGINT primary keys. - // This must run BEFORE the unique_id rewriter so the injected - // function calls get processed. - if let NodeMut::InsertStmt(insert) = node.as_mut() { - self.inject_auto_id(insert, mem, &mut plan)?; - } + // Inject pgdog.unique_id() for missing BIGINT primary keys. + // This must run BEFORE the unique_id rewriter so the injected + // function calls get processed. + if let NodeMut::InsertStmt(insert) = stmt.stmt_mut() { + self.inject_auto_id(insert, mem, &mut plan)?; + } - // Track the next parameter number to use - let mut next_param = plan.params as i32 + 1; - let generator = crate::unique_id::UniqueId::generator()?; - transform::transform(&mut node, |node| { + // Track the next parameter number to use + let mut next_param = plan.params as i32 + 1; + let generator = crate::unique_id::UniqueId::generator()?; + transform::transform_node( + stmt.stmt_mut(), + &mut transform::TransformClosure::new(|node| { if let Some(replacement) = Self::rewrite_unique_id( node.as_ref(), mem, @@ -161,33 +168,25 @@ impl<'a> StatementRewrite<'a> { } else { Some(node) } - }); - - 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_ref())?.as_str().to_owned()); - } + }), + ); - if let Node::InsertStmt(insert) = node.as_ref() { - self.split_insert(insert, &mut plan)?; - } + if let NodeMut::SelectStmt(mut select) = stmt.stmt_mut() { + self.rewrite_aggregates(&mut select, mem, &mut plan, self.db_schema)?; + self.limit_offset(&select, &mut plan); + } - if let Node::UpdateStmt(stmt) = node.as_ref() { - self.sharding_key_update(stmt, &mut plan)?; - } + if self.rewritten { + plan.stmt = Some(pg_raw_parse::deparse(&*stmt)?.as_str().to_owned()); + } - self.stmt.stmts[0] = pg_query::parse(pg_raw_parse::deparse(node.as_ref())?.as_str())? - .protobuf - .stmts - .pop() - .unwrap(); + if let Node::InsertStmt(insert) = stmt.stmt() { + self.split_insert(insert, &mut plan)?; + } - Ok::<_, Error>(mem.make_raw_stmt(node)) - })?; + if let Node::UpdateStmt(stmt) = stmt.stmt() { + self.sharding_key_update(stmt, &mut plan)?; + } Ok(plan) } diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/offset.rs b/pgdog/src/frontend/router/parser/rewrite/statement/offset.rs index 776b1fb37..79094c378 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/offset.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/offset.rs @@ -322,6 +322,7 @@ mod tests { } fn run_limit_offset(sql: &str, schema: &ShardingSchema) -> RewritePlan { + #[cfg(not(feature = "new_parser"))] let mut ast = pg_query::parse(sql).unwrap(); #[cfg(feature = "new_parser")] let stmt = pg_raw_parse::parse(sql).unwrap(); @@ -329,6 +330,7 @@ mod tests { let mut ps = PreparedStatements::default(); #[cfg_attr(feature = "new_parser", allow(unused_mut))] let mut rewrite = StatementRewrite::new(StatementRewriteContext { + #[cfg(not(feature = "new_parser"))] stmt: &mut ast.protobuf, extended: false, prepared: false, diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs b/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs index 8ce37e861..2b156c0f6 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/simple_prepared.rs @@ -207,8 +207,8 @@ mod tests { use crate::backend::ShardingSchema; use crate::backend::schema::Schema; use crate::config::PreparedStatements as PreparedStatementsLevel; + #[cfg(not(feature = "new_parser"))] use pg_query::parse; - use pg_query::protobuf::ParseResult; use pgdog_config::Rewrite; struct TestContext { @@ -235,12 +235,10 @@ mod tests { } } - fn rewrite(&mut self, sql: &str) -> Result<(ParseResult, RewritePlan), Error> { - let mut ast = parse(sql).unwrap().protobuf; - #[cfg(feature = "new_parser")] - let stmt = pg_raw_parse::parse(sql).unwrap(); + #[cfg(feature = "new_parser")] + fn rewrite(&mut self, sql: &str) -> Result<(String, RewritePlan), Error> { + let stmt = pg_raw_parse::parse(sql)?; let mut rewrite = StatementRewrite::new(StatementRewriteContext { - stmt: &mut ast, extended: false, prepared: false, prepared_statements: &mut self.ps, @@ -249,20 +247,43 @@ mod tests { user: "", search_path: None, }); - let plan = rewrite.maybe_rewrite( - #[cfg(feature = "new_parser")] - stmt.stmts().next().unwrap(), - )?; - Ok((ast, plan)) + let mut plan = Default::default(); + let ast = pg_raw_parse::make::try_owned(|mem| { + let mut copy = mem.make_unique(&*stmt.into_inner()); + plan = rewrite.maybe_rewrite(copy.as_mut().into_iter().next().unwrap(), mem)?; + Ok::<_, Error>(copy) + })?; + let sql = pg_raw_parse::deparse_stmts(&*ast)?; + Ok((sql, plan)) + } + + cfg_select! { + not(feature = "new_parser") => { + fn rewrite(&mut self, sql: &str) -> Result<(String, RewritePlan), Error> { + let mut ast = parse(sql).unwrap().protobuf; + let mut rewrite = StatementRewrite::new(StatementRewriteContext { + stmt: &mut ast, + extended: false, + prepared: false, + prepared_statements: &mut self.ps, + schema: &self.schema, + db_schema: &self.db_schema, + user: "", + search_path: None, + }); + let plan = rewrite.maybe_rewrite()?; + Ok((ast.deparse().unwrap(), plan)) + } + } + _ => {} } } #[test] fn test_rewrite_prepare() { let mut ctx = TestContext::new(); - let (ast, plan) = ctx.rewrite("PREPARE test_stmt AS SELECT $1, $2").unwrap(); + let (sql, plan) = ctx.rewrite("PREPARE test_stmt AS SELECT $1, $2").unwrap(); - let sql = ast.deparse().unwrap(); assert!( sql.contains("__pgdog_"), "PREPARE should be renamed to __pgdog_N, got: {sql}" @@ -279,9 +300,8 @@ mod tests { fn test_rewrite_execute() { let mut ctx = TestContext::new(); ctx.rewrite("PREPARE test_stmt AS SELECT 1").unwrap(); - let (ast, plan) = ctx.rewrite("EXECUTE test_stmt").unwrap(); + let (sql, plan) = ctx.rewrite("EXECUTE test_stmt").unwrap(); - let sql = ast.deparse().unwrap(); assert!( sql.contains("__pgdog_"), "EXECUTE should use global name, got: {sql}" @@ -297,9 +317,8 @@ mod tests { fn test_rewrite_execute_with_params() { let mut ctx = TestContext::new(); ctx.rewrite("PREPARE test_stmt AS SELECT $1, $2").unwrap(); - let (ast, plan) = ctx.rewrite("EXECUTE test_stmt(1, 'hello')").unwrap(); + let (sql, plan) = ctx.rewrite("EXECUTE test_stmt(1, 'hello')").unwrap(); - let sql = ast.deparse().unwrap(); assert!( sql.contains("__pgdog_"), "EXECUTE should use global name, got: {sql}" @@ -321,9 +340,8 @@ mod tests { #[test] fn test_no_rewrite_for_regular_select() { let mut ctx = TestContext::new(); - let (ast, plan) = ctx.rewrite("SELECT 1, 2, 3").unwrap(); + let (sql, plan) = ctx.rewrite("SELECT 1, 2, 3").unwrap(); - let sql = ast.deparse().unwrap(); assert_eq!(sql, "SELECT 1, 2, 3"); assert!(plan.prepares.is_empty()); assert!(plan.stmt.is_none()); 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 e8ff52458..46e98d5e1 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/unique_id.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/unique_id.rs @@ -189,7 +189,6 @@ mod tests { use crate::frontend::router::parser::StatementRewriteContext; 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}; @@ -286,9 +285,8 @@ mod tests { #[test] fn test_rewrite_select_extended_single() { - let (ast, plan) = run_test("SELECT pgdog.unique_id()", true); + let (sql, plan) = run_test("SELECT pgdog.unique_id()", true); - let sql = ast.deparse().unwrap(); assert_eq!(sql, "SELECT $1::bigint"); assert_eq!(plan.params, 0); assert_eq!(plan.unique_ids, 1); @@ -296,9 +294,8 @@ mod tests { #[test] fn test_rewrite_select_extended_with_existing_params() { - let (ast, plan) = run_test("SELECT pgdog.unique_id(), $1, $2", true); + let (sql, plan) = run_test("SELECT pgdog.unique_id(), $1, $2", true); - let sql = ast.deparse().unwrap(); assert_eq!(sql, "SELECT $3::bigint, $1, $2"); assert_eq!(plan.params, 2); assert_eq!(plan.unique_ids, 1); @@ -306,9 +303,8 @@ mod tests { #[test] fn test_rewrite_select_extended_multiple_unique_ids() { - let (ast, plan) = run_test("SELECT pgdog.unique_id(), pgdog.unique_id()", true); + let (sql, plan) = run_test("SELECT pgdog.unique_id(), pgdog.unique_id()", true); - let sql = ast.deparse().unwrap(); assert_eq!(sql, "SELECT $1::bigint, $2::bigint"); assert_eq!(plan.params, 0); assert_eq!(plan.unique_ids, 2); @@ -317,9 +313,8 @@ mod tests { #[test] fn test_rewrite_select_simple() { let _guard = set_env_var("NODE_ID", "pgdog-1"); - let (ast, plan) = run_test("SELECT pgdog.unique_id()", false); + let (sql, plan) = run_test("SELECT pgdog.unique_id()", false); - let sql = ast.deparse().unwrap(); // Value should be a bigint literal cast #[cfg(not(feature = "new_parser"))] assert!( @@ -337,9 +332,8 @@ mod tests { #[test] fn test_rewrite_select_simple_multiple_unique_ids() { let _guard = set_env_var("NODE_ID", "pgdog-1"); - let (ast, plan) = run_test("SELECT pgdog.unique_id(), pgdog.unique_id()", false); + let (sql, plan) = run_test("SELECT pgdog.unique_id(), pgdog.unique_id()", false); - let sql = ast.deparse().unwrap(); // Each unique_id call should get a different value assert!( !sql.contains("pgdog.unique_id"), @@ -350,87 +344,79 @@ mod tests { #[test] fn test_rewrite_no_unique_id() { - let (ast, plan) = run_test("SELECT 1, 2, 3", true); + let (sql, plan) = run_test("SELECT 1, 2, 3", true); - let sql = ast.deparse().unwrap(); assert_eq!(sql, "SELECT 1, 2, 3"); assert_eq!(plan.unique_ids, 0); } #[test] fn test_rewrite_insert_values() { - let (ast, plan) = run_test( + let (sql, plan) = run_test( "INSERT INTO t (id, name) VALUES (pgdog.unique_id(), 'test')", true, ); - let sql = ast.deparse().unwrap(); assert_eq!(sql, "INSERT INTO t (id, name) VALUES ($1::bigint, 'test')"); assert_eq!(plan.unique_ids, 1); } #[test] fn test_rewrite_insert_multiple_rows() { - let (ast, plan) = run_test( + let (sql, plan) = run_test( "INSERT INTO t (id) VALUES (pgdog.unique_id()), (pgdog.unique_id())", true, ); - let sql = ast.deparse().unwrap(); assert_eq!(sql, "INSERT INTO t (id) VALUES ($1::bigint), ($2::bigint)"); assert_eq!(plan.unique_ids, 2); } #[test] fn test_rewrite_insert_select() { - let (ast, plan) = run_test("INSERT INTO t (id) SELECT pgdog.unique_id() FROM s", true); + let (sql, plan) = run_test("INSERT INTO t (id) SELECT pgdog.unique_id() FROM s", true); - let sql = ast.deparse().unwrap(); assert_eq!(sql, "INSERT INTO t (id) SELECT $1::bigint FROM s"); assert_eq!(plan.unique_ids, 1); } #[test] fn test_rewrite_update_set() { - let (ast, plan) = run_test( + let (sql, plan) = run_test( "UPDATE t SET id = pgdog.unique_id() WHERE name = 'test'", true, ); - let sql = ast.deparse().unwrap(); assert_eq!(sql, "UPDATE t SET id = $1::bigint WHERE name = 'test'"); assert_eq!(plan.unique_ids, 1); } #[test] fn test_rewrite_update_where() { - let (ast, plan) = run_test( + let (sql, plan) = run_test( "UPDATE t SET name = 'new' WHERE id = pgdog.unique_id()", true, ); - let sql = ast.deparse().unwrap(); assert_eq!(sql, "UPDATE t SET name = 'new' WHERE id = $1::bigint"); assert_eq!(plan.unique_ids, 1); } #[test] fn test_rewrite_delete_where() { - let (ast, plan) = run_test("DELETE FROM t WHERE id = pgdog.unique_id()", true); + let (sql, plan) = run_test("DELETE FROM t WHERE id = pgdog.unique_id()", true); - let sql = ast.deparse().unwrap(); assert_eq!(sql, "DELETE FROM t WHERE id = $1::bigint"); assert_eq!(plan.unique_ids, 1); } #[test] fn test_rewrite_insert_returning() { - let (ast, plan) = run_test( + let (sql, plan) = run_test( "INSERT INTO t (id) VALUES (pgdog.unique_id()) RETURNING pgdog.unique_id()", true, ); - let sql = ast.deparse().unwrap(); assert_eq!( sql, "INSERT INTO t (id) VALUES ($1::bigint) RETURNING $2::bigint" @@ -440,34 +426,30 @@ mod tests { #[test] fn test_rewrite_explain_insert_select() { - let (ast, plan) = run_test( + let (sql, plan) = run_test( "EXPLAIN INSERT INTO t (id) SELECT pgdog.unique_id() FROM s", true, ); - let sql = ast.deparse().unwrap(); assert_eq!(sql, "EXPLAIN INSERT INTO t (id) SELECT $1::bigint FROM s"); assert_eq!(plan.unique_ids, 1); } #[test] fn test_rewrite_explain_select() { - let (ast, plan) = run_test("EXPLAIN SELECT pgdog.unique_id()", true); + let (sql, plan) = run_test("EXPLAIN SELECT pgdog.unique_id()", true); - let sql = ast.deparse().unwrap(); assert_eq!(sql, "EXPLAIN SELECT $1::bigint"); assert_eq!(plan.unique_ids, 1); } - fn run_test(sql: &str, extended: bool) -> (ParseResult, RewritePlan) { - let mut ast = pg_query::parse(sql).unwrap().protobuf; - #[cfg(feature = "new_parser")] + #[cfg(feature = "new_parser")] + fn run_test(sql: &str, extended: bool) -> (String, RewritePlan) { let stmt = pg_raw_parse::parse(sql).unwrap(); let mut ps = PreparedStatements::default(); let schema = default_schema(); let db_schema = default_db_schema(); let mut rewrite = StatementRewrite::new(StatementRewriteContext { - stmt: &mut ast, extended, prepared: false, prepared_statements: &mut ps, @@ -476,12 +458,40 @@ mod tests { user: "", search_path: None, }); - let plan = rewrite - .maybe_rewrite( - #[cfg(feature = "new_parser")] - stmt.stmts().next().unwrap(), - ) - .unwrap(); - (ast, plan) + let mut plan = Default::default(); + let ast = make::owned(|mem| { + let mut copy = mem.make_unique(&*stmt.into_inner()); + let stmt = copy.as_mut().into_iter().next().unwrap(); + plan = rewrite.maybe_rewrite(stmt, mem).unwrap(); + copy + }); + let sql = pg_raw_parse::deparse_stmts(&*ast).unwrap(); + (sql, plan) + } + + cfg_select! { + not(feature = "new_parser") => { + fn run_test(sql: &str, extended: bool) -> (String, RewritePlan) { + let mut ast = pg_query::parse(sql).unwrap().protobuf; + let mut ps = PreparedStatements::default(); + let schema = default_schema(); + let db_schema = default_db_schema(); + let mut rewrite = StatementRewrite::new(StatementRewriteContext { + stmt: &mut ast, + extended, + prepared: false, + prepared_statements: &mut ps, + schema: &schema, + db_schema: &db_schema, + user: "", + search_path: None, + }); + let plan = rewrite + .maybe_rewrite() + .unwrap(); + (ast.deparse().unwrap(), plan) + } + } + _ => {} } } diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs index 8707e7f0c..8d179f9dc 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/update.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/update.rs @@ -1063,6 +1063,7 @@ mod test { } fn run_test(query: &str) -> Result, Error> { + #[cfg(not(feature = "new_parser"))] let mut stmt_old = parse(query)?; #[cfg(feature = "new_parser")] let stmt = pg_raw_parse::parse(query)?; @@ -1071,6 +1072,7 @@ mod test { let mut stmts = PreparedStatements::new(); let ctx = StatementRewriteContext { + #[cfg(not(feature = "new_parser"))] stmt: &mut stmt_old.protobuf, schema: &schema, db_schema: &db_schema,