diff --git a/integration/rust/tests/integration/stddev.rs b/integration/rust/tests/integration/stddev.rs index 805d932b1..f5865e409 100644 --- a/integration/rust/tests/integration/stddev.rs +++ b/integration/rust/tests/integration/stddev.rs @@ -1,6 +1,25 @@ use crate::setup::{admin_sqlx, connections_sqlx}; use rust_decimal::prelude::*; -use sqlx::{Executor, Pool, Postgres, Row}; +use sqlx::{Column, Executor, Pool, Postgres, Row}; + +#[tokio::test] +async fn test_direct_stddev_hides_cross_shard_helpers() { + let conns = connections_sqlx().await; + + setup_schema(&conns, "test_direct_stddev", "int8").await; + setup_data(&conns, "test_direct_stddev", TEST_DATA).await; + + let row = conns[1] + .fetch_one( + "SELECT stddev_pop(value) AS deviation \ + FROM test_direct_stddev WHERE customer_id = 1", + ) + .await + .unwrap(); + + assert_eq!(row.columns().len(), 1); + assert_eq!(row.columns()[0].name(), "deviation"); +} #[tokio::test] async fn test_variance_numeric() { diff --git a/pgdog/src/frontend/client/query_engine/test/mod.rs b/pgdog/src/frontend/client/query_engine/test/mod.rs index 791312a6a..cf2382be9 100644 --- a/pgdog/src/frontend/client/query_engine/test/mod.rs +++ b/pgdog/src/frontend/client/query_engine/test/mod.rs @@ -28,6 +28,7 @@ mod pipeline_execution; pub(crate) mod prelude; mod prepared_syntax_error; mod replicas; +mod rewrite_aggregate; mod rewrite_extended; mod rewrite_insert_split; mod rewrite_offset; diff --git a/pgdog/src/frontend/client/query_engine/test/rewrite_aggregate.rs b/pgdog/src/frontend/client/query_engine/test/rewrite_aggregate.rs new file mode 100644 index 000000000..d3dfd5442 --- /dev/null +++ b/pgdog/src/frontend/client/query_engine/test/rewrite_aggregate.rs @@ -0,0 +1,89 @@ +use crate::frontend::router::parser::Shard; + +use super::prelude::*; +use super::test_sharded_client; + +async fn route_and_rewrite_request(messages: Vec) -> (Shard, String) { + let mut client = test_sharded_client(); + client.client_request = ClientRequest::from(messages); + + let mut engine = QueryEngine::from_client(&client).unwrap(); + let mut context = QueryEngineContext::new(&mut client); + + let rewrite = engine.parse_and_rewrite(&mut context).unwrap(); + engine + .route_query(&mut context, rewrite.as_ref()) + .await + .unwrap(); + + let route = context.client_request.route().shard().clone(); + let sql = context + .client_request + .iter() + .find_map(|message| match message { + ProtocolMessage::Query(query) => Some(query.query().to_owned()), + ProtocolMessage::Parse(parse) => Some(parse.query().to_owned()), + _ => None, + }) + .unwrap(); + + (route, sql) +} + +async fn route_and_rewrite(sql: &str) -> (Shard, String) { + route_and_rewrite_request(vec![Query::new(sql).into()]).await +} + +#[tokio::test] +async fn test_direct_aggregates_do_not_include_cross_shard_helpers() { + for function in ["avg", "stddev", "variance"] { + let (route, sql) = route_and_rewrite(&format!( + "SELECT {function}(region_id) FROM sharded WHERE id = 1" + )) + .await; + + assert!(matches!(route, Shard::Direct(_))); + assert!( + !sql.contains("__pgdog_"), + "direct {function} query included helper columns: {sql}" + ); + } +} + +#[tokio::test] +async fn test_cross_shard_aggregate_keeps_helpers() { + let (route, sql) = route_and_rewrite("SELECT stddev(region_id) FROM sharded").await; + + assert!(route.is_all()); + assert!(sql.contains("__pgdog_count_col0")); + assert!(sql.contains("__pgdog_sum_col0")); + assert!(sql.contains("__pgdog_sumsq_col0")); +} + +#[tokio::test] +async fn test_direct_extended_aggregate_does_not_include_helpers() { + let (route, sql) = route_and_rewrite_request(vec![ + Parse::new_anonymous("SELECT avg(region_id) FROM sharded WHERE id = $1").into(), + Bind::new_params("", &[Parameter::new(b"1")]).into(), + Execute::new().into(), + Sync.into(), + ]) + .await; + + assert!(matches!(route, Shard::Direct(_))); + assert!(!sql.contains("__pgdog_")); +} + +#[tokio::test] +async fn test_cross_shard_extended_aggregate_keeps_helpers() { + let (route, sql) = route_and_rewrite_request(vec![ + Parse::new_anonymous("SELECT avg(region_id) FROM sharded").into(), + Bind::new_params("", &[]).into(), + Execute::new().into(), + Sync.into(), + ]) + .await; + + assert!(route.is_all()); + assert!(sql.contains("__pgdog_count_col0")); +} diff --git a/pgdog/src/frontend/client/query_engine/test/rewrite_offset.rs b/pgdog/src/frontend/client/query_engine/test/rewrite_offset.rs index 5eccc828a..142470a2e 100644 --- a/pgdog/src/frontend/client/query_engine/test/rewrite_offset.rs +++ b/pgdog/src/frontend/client/query_engine/test/rewrite_offset.rs @@ -18,7 +18,7 @@ async fn run_test(messages: Vec) -> Option { let rewrite_result = engine.parse_and_rewrite(&mut context).unwrap(); match rewrite_result { - Some(RewriteResult::InPlace { offset }) => offset, + Some(RewriteResult::InPlace { offset, .. }) => offset, other => panic!("expected InPlace, got {:?}", other), } } 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 2a2ef2fb8..635644d6f 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/engine.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/aggregate/engine.rs @@ -64,6 +64,27 @@ impl AggregatesRewrite { } } + /// Remove helper aggregate columns previously appended by [`Self::rewrite_select`]. + pub(crate) fn rollback_select<'a>( + select: &mut nodes::SelectStmtMut<'a, '_>, + mem: make::MemoryToken<'a>, + plan: &AggregateRewritePlan, + ) { + let original_len = select + .target_list() + .len() + .checked_sub(plan.helpers().len()) + .expect("aggregate helper plan cannot exceed the SELECT target list"); + let targets = select + .target_list() + .iter() + .take(original_len) + .map(|target| mem.make_unique(target)) + .collect::>(); + + select.target_list_mut().replace(mem.make_list(&targets)); + } + fn build_sum_of_squares_func<'a>( original: &nodes::FuncCall, mem: make::MemoryToken<'a>, @@ -288,4 +309,27 @@ mod tests { // Expect original STDDEV plus three helpers. assert_eq!(ast.target_list().len(), 4); } + + #[test] + fn rewrite_engine_rolls_back_helpers() { + let sql = "SELECT AVG(price), STDDEV(discount) FROM menu"; + let ast = pg_raw_parse::parse(sql).unwrap(); + + let select = make::owned(|mem| { + let Node::SelectStmt(stmt) = ast.stmts().next().unwrap() else { + unreachable!("not a select"); + }; + let mut stmt = mem.make_unique(stmt); + let aggregate = Aggregate::parse(&stmt, &Default::default()); + let output = AggregatesRewrite::rewrite_select(&mut stmt.as_mut(), mem, &aggregate); + + assert_eq!(stmt.target_list().len(), 6); + AggregatesRewrite::rollback_select(&mut stmt.as_mut(), mem, &output.plan); + assert_eq!(stmt.target_list().len(), 2); + stmt + }); + + let statement = pg_raw_parse::deparse(&*select).unwrap(); + assert!(!statement.as_str().contains("__pgdog_")); + } } diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs index 22a8a47f5..70eeb0aaf 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs @@ -166,8 +166,27 @@ impl<'a> StatementRewrite<'a> { return Err(err); } - if let NodeMut::SelectStmt(mut select) = stmt.stmt_mut() { + if matches!(stmt.stmt(), Node::SelectStmt(_)) { + let NodeMut::SelectStmt(mut select) = stmt.stmt_mut() else { + unreachable!("statement was checked to be SELECT"); + }; self.rewrite_aggregates(&mut select, mem, &mut plan, self.db_schema)?; + if !plan.aggregates.is_noop() { + aggregate::AggregatesRewrite::rollback_select(&mut select, mem, &plan.aggregates); + } + + if !plan.aggregates.is_noop() { + plan.direct_stmt = Some(pg_raw_parse::deparse(&*stmt)?.as_str().to_owned()); + + let NodeMut::SelectStmt(mut select) = stmt.stmt_mut() else { + unreachable!("statement was checked to be SELECT"); + }; + self.rewrite_aggregates(&mut select, mem, &mut plan, self.db_schema)?; + } + + let NodeMut::SelectStmt(select) = stmt.stmt_mut() else { + unreachable!("statement was checked to be SELECT"); + }; self.limit_offset(&select, &mut plan); } diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs index 82f3c723f..c516f2706 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/plan.rs @@ -1,6 +1,6 @@ use crate::frontend::{ClientRequest, PreparedStatements}; use crate::net::messages::bind::{Format, Parameter}; -use crate::net::{Bind, Parse, ProtocolMessage, Query}; +use crate::net::{Bind, ProtocolMessage}; use crate::unique_id::UniqueId; use super::insert::build_split_requests; @@ -9,6 +9,21 @@ use super::{ Error, InsertSplit, PrepareExecute, ShardingKeyUpdate, aggregate::AggregateRewritePlan, }; +fn apply_statement(request: &mut ClientRequest, stmt: &str) { + for message in request.messages.iter_mut() { + match message { + ProtocolMessage::Query(query) => query.set_query(stmt), + ProtocolMessage::Parse(parse) => { + parse.set_query(stmt); + if !parse.anonymous() { + PreparedStatements::global().write().rewrite(parse); + } + } + _ => {} + } + } +} + /// Statement rewrite plan. /// /// Executed in order of fields in this struct. @@ -30,6 +45,11 @@ pub(crate) struct RewritePlan { /// Rewritten SQL statement. pub(crate) stmt: Option, + /// Rewritten SQL statement without cross-shard aggregate helper columns. + /// + /// Restored after routing when a query only needs one shard. + pub(crate) direct_stmt: Option, + /// Prepared statements to prepend to the client request. /// Each tuple contains (name, statement) for ProtocolMessage::Prepare. pub(crate) prepare_rewrites: Vec, @@ -52,7 +72,10 @@ pub(crate) struct RewritePlan { #[derive(Debug, Clone)] pub(crate) enum RewriteResult { - InPlace { offset: Option }, + InPlace { + offset: Option, + direct_stmt: Option, + }, InsertSplit(Vec), ShardingKeyUpdate(ShardingKeyUpdate), } @@ -61,8 +84,25 @@ impl RewriteResult { pub(crate) fn apply_after_parser(&self, request: &mut ClientRequest) -> Result<(), Error> { match self { Self::InPlace { - offset: Some(offset), - } => offset.apply_after_parser(request), + offset, + direct_stmt, + } => { + if request.is_executable() + && request + .route + .as_ref() + .is_some_and(|route| !route.is_cross_shard()) + && let Some(stmt) = direct_stmt + { + apply_statement(request, stmt); + } + + if let Some(offset) = offset { + offset.apply_after_parser(request)?; + } + + Ok(()) + } _ => Ok(()), } } @@ -100,23 +140,6 @@ impl RewritePlan { Ok(()) } - /// Apply the rewrite plan to a Parse message by updating the SQL. - pub(crate) fn apply_parse(&self, parse: &mut Parse) { - if let Some(ref stmt) = self.stmt { - parse.set_query(stmt); - if !parse.anonymous() { - PreparedStatements::global().write().rewrite(parse); - } - } - } - - /// Apply the rewrite plan to a Query message by updating the SQL. - pub(crate) fn apply_query(&self, query: &mut Query) { - if let Some(ref stmt) = self.stmt { - query.set_query(stmt); - } - } - /// Apply the rewrite plan to a ClientRequest. pub(crate) fn apply(&self, request: &mut ClientRequest) -> Result { // Prepend any required Prepare messages for EXECUTE statements. @@ -136,12 +159,13 @@ impl RewritePlan { }); } + if let Some(stmt) = &self.stmt { + apply_statement(request, stmt); + } + for message in request.messages.iter_mut() { - match message { - ProtocolMessage::Parse(parse) => self.apply_parse(parse), - ProtocolMessage::Query(query) => self.apply_query(query), - ProtocolMessage::Bind(bind) => self.apply_bind(bind)?, - _ => {} + if let ProtocolMessage::Bind(bind) = message { + self.apply_bind(bind)?; } } @@ -163,6 +187,7 @@ impl RewritePlan { Ok(RewriteResult::InPlace { offset: self.offset.clone(), + direct_stmt: self.direct_stmt.clone(), }) } }