diff --git a/integration/rust/tests/integration/having.rs b/integration/rust/tests/integration/having.rs new file mode 100644 index 000000000..ad9683e27 --- /dev/null +++ b/integration/rust/tests/integration/having.rs @@ -0,0 +1,171 @@ +use crate::setup::{admin_sqlx, connections_sqlx, connections_tokio}; +use sqlx::{Executor, Row, postgres::PgPool}; + +#[tokio::test] +async fn having_filters_merged_groups() -> Result<(), Box> { + let sharded = connections_sqlx().await.get(1).cloned().unwrap(); + reset_table(&sharded).await?; + + seed( + &sharded, + 0, + "INSERT INTO having_merge_test VALUES (1, 10), (1, 11), (2, 12)", + ) + .await?; + seed( + &sharded, + 1, + "INSERT INTO having_merge_test VALUES (1, 20), (2, 21)", + ) + .await?; + + let rows = sharded + .fetch_all( + "SELECT customer_id, COUNT(*) AS order_count \ + FROM having_merge_test \ + GROUP BY customer_id \ + HAVING COUNT(*) > 2 \ + ORDER BY customer_id", + ) + .await?; + + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].columns().len(), 2); + assert_eq!(rows[0].get::("customer_id"), 1); + assert_eq!(rows[0].get::("order_count"), 3); + + // An untyped literal takes the type of the aggregate it is compared with. + let rows = sharded + .fetch_all( + "SELECT customer_id FROM having_merge_test \ + GROUP BY customer_id HAVING MAX(order_id) > '20'", + ) + .await?; + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].get::("customer_id"), 2); + + cleanup(&sharded).await; + Ok(()) +} + +#[tokio::test] +async fn prepared_having_parameter_typed_and_untyped() -> Result<(), Box> { + let sharded = connections_sqlx().await.get(1).cloned().unwrap(); + reset_table(&sharded).await?; + + seed( + &sharded, + 0, + "INSERT INTO having_merge_test VALUES (1, 10), (1, 11)", + ) + .await?; + seed( + &sharded, + 1, + "INSERT INTO having_merge_test VALUES (1, 20), (2, 21)", + ) + .await?; + + let rows = sqlx::query( + "SELECT customer_id, COUNT(*) AS order_count \ + FROM having_merge_test \ + GROUP BY customer_id \ + HAVING COUNT(*) > $1", + ) + .bind(2_i64) + .fetch_all(&sharded) + .await?; + + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].columns().len(), 2); + assert_eq!(rows[0].get::("customer_id"), 1); + assert_eq!(rows[0].get::("order_count"), 3); + + // Parse without parameter types, so PostgreSQL has to infer $1. + let client = connections_tokio().await.remove(1); + let statement = client + .prepare( + "SELECT customer_id FROM having_merge_test \ + GROUP BY customer_id HAVING COUNT(*) > $1", + ) + .await?; + let rows = client.query(&statement, &[&2_i64]).await?; + assert_eq!(rows.len(), 1); + assert_eq!(rows[0].get::<_, i64>(0), 1); + + cleanup(&sharded).await; + Ok(()) +} + +#[tokio::test] +async fn having_rejects_shard_local_steps() -> Result<(), Box> { + let sharded = connections_sqlx().await.get(1).cloned().unwrap(); + reset_table(&sharded).await?; + let client = connections_tokio().await.remove(1); + + // Each of these would run on shards before PgDog merges the groups. + for sql in [ + "SELECT * FROM having_merge_test GROUP BY customer_id, order_id HAVING COUNT(*) > 0", + "SELECT customer_id FROM having_merge_test GROUP BY customer_id \ + HAVING COUNT(*) > 0 ORDER BY customer_id LIMIT 1", + "SELECT customer_id FROM having_merge_test GROUP BY customer_id \ + HAVING COUNT(*) > 0 OFFSET 1", + "SELECT DISTINCT ON (customer_id) customer_id FROM having_merge_test \ + GROUP BY customer_id HAVING COUNT(*) > 0", + "SELECT customer_id AS c FROM having_merge_test GROUP BY c HAVING COUNT(*) > 0", + "SELECT customer_id, COUNT(*) + 1 FROM having_merge_test \ + GROUP BY customer_id HAVING COUNT(*) > 0", + "EXECUTE having_prepared(0)", + ] { + if sql.starts_with("EXECUTE") { + client + .simple_query( + "PREPARE having_prepared AS SELECT customer_id FROM having_merge_test \ + GROUP BY customer_id HAVING COUNT(*) > $1", + ) + .await?; + } + let error = client.simple_query(sql).await.unwrap_err(); + let message = error.as_db_error().map(|error| error.message()); + assert!( + message.is_some_and(|message| message.contains("cross-shard HAVING is unsupported")), + "{sql}: {error}" + ); + } + client.simple_query("SELECT 1").await?; + + cleanup(&sharded).await; + Ok(()) +} + +async fn reset_table(pool: &PgPool) -> Result<(), sqlx::Error> { + cleanup(pool).await; + for shard in [0, 1] { + pool.execute( + format!( + "/* pgdog_shard: {shard} */ \ + CREATE TABLE having_merge_test(customer_id BIGINT, order_id BIGINT)" + ) + .as_str(), + ) + .await?; + } + admin_sqlx().await.execute("RELOAD").await?; + Ok(()) +} + +async fn seed(pool: &PgPool, shard: usize, statement: &str) -> Result<(), sqlx::Error> { + pool.execute(format!("/* pgdog_shard: {shard} */ {statement}").as_str()) + .await?; + Ok(()) +} + +async fn cleanup(pool: &PgPool) { + for shard in [0, 1] { + pool.execute( + format!("/* pgdog_shard: {shard} */ DROP TABLE IF EXISTS having_merge_test").as_str(), + ) + .await + .ok(); + } +} diff --git a/integration/rust/tests/integration/mod.rs b/integration/rust/tests/integration/mod.rs index 029213dca..a1ae52d22 100644 --- a/integration/rust/tests/integration/mod.rs +++ b/integration/rust/tests/integration/mod.rs @@ -17,6 +17,7 @@ mod cross_shard_omni_check; pub mod distinct; pub mod explain; pub mod fake_transactions; +pub mod having; pub mod idle_in_transaction; pub mod limit; pub mod maintenance_mode; diff --git a/pgdog/src/backend/pool/connection/buffer/having.rs b/pgdog/src/backend/pool/connection/buffer/having.rs index f6d8801df..c0ebb1077 100644 --- a/pgdog/src/backend/pool/connection/buffer/having.rs +++ b/pgdog/src/backend/pool/connection/buffer/having.rs @@ -89,6 +89,21 @@ fn compare(left: &Datum, operator: CompareOperator, right: &Datum) -> Result are supported for text".into(), + )); + } let ordering = left .compare_numeric(right) @@ -195,10 +210,10 @@ mod tests { true, ), ( - CompareOperator::Greater, + CompareOperator::NotEqual, Datum::Text("a".into()), Datum::Text("z".into()), - false, + true, ), ] { assert_eq!( @@ -211,13 +226,27 @@ mod tests { ); } - assert!(matches!( - compare( - &Datum::Text("2".into()), + for (left, operator, right) in [ + ( + Datum::Text("2".into()), + CompareOperator::Equal, + Datum::Integer(2), + ), + ( + Datum::Text("a".into()), + CompareOperator::Greater, + Datum::Text("B".into()), + ), + ( + Datum::Unknown("2026-01-02".into()), CompareOperator::Equal, - &Datum::Integer(2) + Datum::Unknown("2026-01-02".into()), ), - Err(Error::UnsupportedHaving(_)) - )); + ] { + assert!(matches!( + compare(&left, operator, &right), + Err(Error::UnsupportedHaving(_)) + )); + } } } diff --git a/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs b/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs index 6732ddec4..20ed6bd36 100644 --- a/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs +++ b/pgdog/src/frontend/client/query_engine/test/rewrite_projection.rs @@ -97,6 +97,79 @@ async fn cross_shard_aggregate_adds_and_tracks_helpers() { ); } +#[tokio::test] +async fn cross_shard_having_rewrites_and_replaces_route_aggregate() { + let sql = "SELECT customer_id, COUNT(*) AS orders \ + FROM purchases GROUP BY customer_id HAVING COUNT(*) > 2"; + let mut client = test_sharded_client(); + client.client_request = ClientRequest::from(vec![ProtocolMessage::Query(Query::new(sql))]); + + let engine = QueryEngine::from_client(&client).unwrap(); + let (mut context, client_request) = QueryEngineContext::new(&mut client); + let result = engine + .parse_and_rewrite(&mut context, client_request) + .await + .unwrap(); + client_request.route = Some(route(Shard::All)); + + projection::finalize_after_route( + client_request, + &Schema::default(), + result.as_ref().and_then(RewriteResult::offset_plan), + ) + .unwrap(); + + let query = match &client_request.messages[0] { + ProtocolMessage::Query(query) => query, + _ => panic!("expected Query"), + }; + assert!(query.query().contains("__pgdog_having_0")); + assert_eq!( + client_request.route().aggregate().targets().len(), + 2, + "both the selected and hidden COUNT must be merged" + ); + assert!( + client_request + .route() + .projection_rewrite_plan + .as_deref() + .expect("projection rewrite plan") + .having + .is_some() + ); +} + +#[tokio::test] +async fn direct_having_keeps_native_postgres_semantics() { + let sql = "SELECT customer_id, COUNT(*) AS orders \ + FROM purchases GROUP BY customer_id HAVING COUNT(*) > 2"; + let mut client = test_sharded_client(); + client.client_request = ClientRequest::from(vec![ProtocolMessage::Query(Query::new(sql))]); + + let engine = QueryEngine::from_client(&client).unwrap(); + let (mut context, client_request) = QueryEngineContext::new(&mut client); + let result = engine + .parse_and_rewrite(&mut context, client_request) + .await + .unwrap(); + client_request.route = Some(route(Shard::Direct(0))); + + projection::finalize_after_route( + client_request, + &Schema::default(), + result.as_ref().and_then(RewriteResult::offset_plan), + ) + .unwrap(); + + let query = match &client_request.messages[0] { + ProtocolMessage::Query(query) => query, + _ => panic!("expected Query"), + }; + assert_eq!(query.query(), sql); + assert!(client_request.route().projection_rewrite_plan.is_none()); +} + #[tokio::test] async fn named_prepared_aggregate_uses_cross_shard_variant() { let mut client = test_sharded_client(); diff --git a/pgdog/src/frontend/client/query_engine/test/unsupported_aggregation.rs b/pgdog/src/frontend/client/query_engine/test/unsupported_aggregation.rs index 536c4121c..1c1c40e31 100644 --- a/pgdog/src/frontend/client/query_engine/test/unsupported_aggregation.rs +++ b/pgdog/src/frontend/client/query_engine/test/unsupported_aggregation.rs @@ -20,7 +20,7 @@ async fn test_unsupported_aggregation_keeps_connection_usable() { ), ( "SELECT count(*) FROM sharded HAVING count(*) + 1 > 0", - "cross-shard HAVING is unsupported: aggregates would be filtered on each shard before merging", + "cross-shard HAVING is unsupported: operands must be aggregate calls, grouped columns, constants or parameters", ), ( "SELECT value FROM sharded GROUP BY value ORDER BY count(*)", diff --git a/pgdog/src/frontend/router/parser/aggregate.rs b/pgdog/src/frontend/router/parser/aggregate.rs index 7eff61390..853044b20 100644 --- a/pgdog/src/frontend/router/parser/aggregate.rs +++ b/pgdog/src/frontend/router/parser/aggregate.rs @@ -39,6 +39,23 @@ pub(crate) enum AggregateFunction { Unrecognized(String), } +impl AggregateFunction { + pub(crate) fn from_builtin_name(name: &str) -> Option { + Some(match name { + "count" => Self::Count, + "max" => Self::Max, + "min" => Self::Min, + "sum" => Self::Sum, + "avg" => Self::Avg, + "stddev" | "stddev_samp" => Self::StddevSamp, + "stddev_pop" => Self::StddevPop, + "variance" | "var_samp" => Self::VarSamp, + "var_pop" => Self::VarPop, + _ => return None, + }) + } +} + impl fmt::Display for AggregateFunction { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> Result<(), fmt::Error> { match self { @@ -94,20 +111,11 @@ fn columns_match(group_by_names: &[&str], select_names: &[&str]) -> bool { impl AggregateFunction { pub(crate) fn from_name(name: &str, schema: &Schema) -> Option { - Some(match name { - "count" => AggregateFunction::Count, - "max" => AggregateFunction::Max, - "min" => AggregateFunction::Min, - "sum" => AggregateFunction::Sum, - "avg" => AggregateFunction::Avg, - "stddev" | "stddev_samp" => AggregateFunction::StddevSamp, - "stddev_pop" => AggregateFunction::StddevPop, - "variance" | "var_samp" => AggregateFunction::VarSamp, - "var_pop" => AggregateFunction::VarPop, - name if schema.aggregate_functions.contains(name) => { - AggregateFunction::Unrecognized(name.to_owned()) - } - _ => return None, + Self::from_builtin_name(name).or_else(|| { + schema + .aggregate_functions + .contains(name) + .then(|| Self::Unrecognized(name.to_owned())) }) } } diff --git a/pgdog/src/frontend/router/parser/having.rs b/pgdog/src/frontend/router/parser/having.rs index b007bc341..57ba43a62 100644 --- a/pgdog/src/frontend/router/parser/having.rs +++ b/pgdog/src/frontend/router/parser/having.rs @@ -1,6 +1,4 @@ //! Cross-shard `HAVING` execution plan. -// Only tests build plans until the HAVING rewriter is wired in. -#![cfg_attr(not(test), allow(dead_code))] /// Comparison operators supported by the cross-shard `HAVING` evaluator. #[derive(Debug, Clone, Copy, PartialEq, Eq)] diff --git a/pgdog/src/frontend/router/parser/query/prepare.rs b/pgdog/src/frontend/router/parser/query/prepare.rs index 315dd5172..8297d3876 100644 --- a/pgdog/src/frontend/router/parser/query/prepare.rs +++ b/pgdog/src/frontend/router/parser/query/prepare.rs @@ -1,7 +1,10 @@ use pg_raw_parse::nodes::{ExecuteStmt, PrepareStmt}; use crate::{ - frontend::{BufferedQuery, PreparedStatements}, + frontend::{ + BufferedQuery, PreparedStatements, + router::parser::rewrite::statement::Error as RewriteError, + }, net::{PREPARE_TEMPLATE_NAME, Query}, }; @@ -68,7 +71,18 @@ impl QueryParser { }; match stmt { - Some(Node::SelectStmt(stmt)) => self.select(stmt, &mut context), + Some(Node::SelectStmt(stmt)) => { + let command = self.select(stmt, &mut context)?; + // The cross-shard HAVING rewrite never sees EXECUTE, so each + // shard would filter its partial groups before the merge. + if !matches!(stmt.having_clause(), Node::None) && command.route().is_cross_shard() { + return Err(RewriteError::UnsupportedHaving( + "PREPARE and EXECUTE statements are unsupported", + ) + .into()); + } + Ok(command) + } Some(Node::InsertStmt(stmt)) => self.insert(stmt.into(), &mut context), Some(Node::UpdateStmt(stmt)) => self.update(stmt.into(), &mut context), Some(Node::DeleteStmt(stmt)) => self.delete(stmt.into(), &mut context), diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/having.rs b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs new file mode 100644 index 000000000..1f710a44e --- /dev/null +++ b/pgdog/src/frontend/router/parser/rewrite/statement/having.rs @@ -0,0 +1,324 @@ +//! Route-dependent rewrite for cross-shard `HAVING`. + +use std::collections::HashSet; + +use itertools::Itertools; +use pg_raw_parse::{ConstValue, Node, make, nodes}; + +use crate::frontend::router::parser::{ + AggregateFunction, CompareOperator, Function, HavingExpr, HavingRewritePlan, HavingValue, +}; + +use super::{ + Error, + projection::{column_name, target_output_name}, +}; + +const ALIAS_PREFIX: &str = "__pgdog_having_"; + +/// Project every supported predicate operand and disable shard-local filtering. +/// +/// The original expression remains under `OR TRUE` so PostgreSQL still performs +/// parse analysis and parameter type inference. +pub(super) fn rewrite_select<'a>( + select: &mut nodes::SelectStmtMut<'a, '_>, + mem: make::MemoryToken<'a>, +) -> Result { + let having = select.having_clause(); + debug_assert!( + !matches!(having, Node::None), + "HAVING rewrite requires a HAVING clause" + ); + let having_source = mem.make_unique(having); + let having = having_source.as_ref(); + + let aliases = select + .target_list() + .iter() + .filter_map(|target| target_output_name(target).map(str::to_owned)) + .collect(); + let mut parser = Parser { + aliases, + values: Vec::new(), + next_alias: 0, + }; + let expression = parser.parse_expression(having)?; + let group_sources = select + .group_clause() + .iter() + .map(|group| mem.make_unique(group)) + .collect::>(); + for group_source in &group_sources { + let group = group_source.as_ref(); + match group { + // Shards can't resolve a SELECT alias inside the hidden projection. + Node::ColumnRef(column) if groups_by_output_alias(select, column) => { + return Err(Error::UnsupportedHaving( + "GROUP BY on a SELECT alias is unsupported", + )); + } + Node::ColumnRef(_) => { + let alias = parser.next_alias(); + parser.values.push(Projected { + alias, + node: group, + typed_like: None, + }); + } + Node::A_Const(_) => {} + _ => { + return Err(Error::UnsupportedHaving( + "GROUP BY expressions other than columns and ordinals are unsupported", + )); + } + } + } + + let hidden_aliases = parser + .values + .iter() + .map(|value| value.alias.clone()) + .collect::>(); + let targets = parser + .values + .iter() + .map(|value| { + let projected = match value.typed_like { + Some(like) => typed_like(mem, value.node, like), + None => mem.make_unique(value.node), + }; + mem.make_res_target(Some(&value.alias), mem.empty(), projected) + }) + .collect::>(); + select + .target_list_mut() + .extend(mem, mem.make_list(&targets)); + + let original = mem.make_unique(having); + let true_value = mem.make_a_const(ConstValue::Boolean(true)).uncast(); + let mut neutralized = mem.make_node::(); + neutralized + .as_mut() + .set_boolop(nodes::BoolExprType::OR_EXPR); + neutralized + .as_mut() + .set_args(mem.make_list(&[original, true_value])); + select.set_having_clause(neutralized.uncast()); + + Ok(HavingRewritePlan::new(expression, hidden_aliases)) +} + +fn groups_by_output_alias( + select: &nodes::SelectStmtMut<'_, '_>, + column: &nodes::ColumnRef, +) -> bool { + column.fields().len() == 1 + && select.target_list().iter().any(|target| { + target + .name() + .is_some_and(|name| Some(name) == column_name(column)) + }) +} + +/// `CASE WHEN false THEN ELSE END` evaluates to `value` with the +/// type PostgreSQL infers for it in a comparison with `like`. A bare untyped +/// literal or parameter in the target list would be resolved as `text` instead. +fn typed_like<'a>( + mem: make::MemoryToken<'a>, + value: Node<'_>, + like: Node<'_>, +) -> make::Unique<'a, Node<'a>> { + let mut when = mem.make_node::(); + when.as_mut() + .set_expr(mem.make_a_const(ConstValue::Boolean(false)).uncast()); + when.as_mut().set_result(mem.make_unique(like)); + let mut case = mem.make_node::(); + case.as_mut().set_args(mem.make_list(&[when.uncast()])); + case.as_mut().set_defresult(mem.make_unique(value)); + case.uncast() +} + +struct Projected<'a> { + alias: String, + node: Node<'a>, + /// Operand whose type an untyped literal or parameter should take. + typed_like: Option>, +} + +struct Parser<'a> { + aliases: HashSet, + values: Vec>, + next_alias: usize, +} + +impl<'a> Parser<'a> { + fn parse_expression(&mut self, node: Node<'a>) -> Result { + match node { + Node::A_Expr(expression) if expression.kind == nodes::A_Expr_Kind::AEXPR_OP => { + let operator = expression + .name() + .iter() + .exactly_one() + .ok() + .and_then(Node::as_str) + .and_then(Self::comparison_operator) + .ok_or_else(|| { + Error::UnsupportedHaving( + "only =, <>, !=, <, <=, > and >= comparisons are supported", + ) + })?; + let (left, right) = (expression.lexpr(), expression.rexpr()); + Ok(HavingExpr::Compare { + left: self.parse_value(left, right)?, + operator, + right: self.parse_value(right, left)?, + }) + } + _ => Err(Error::UnsupportedHaving( + "only a single comparison is supported", + )), + } + } + + fn parse_value(&mut self, node: Node<'a>, other: Node<'a>) -> Result { + if !Self::supported_value(node) { + return Err(Error::UnsupportedHaving( + "operands must be aggregate calls, grouped columns, constants or parameters", + )); + } + + let untyped = |node| matches!(node, Node::A_Const(_) | Node::ParamRef(_)); + let alias = self.next_alias(); + self.values.push(Projected { + alias: alias.clone(), + node, + typed_like: (untyped(node) && !untyped(other)).then_some(other), + }); + Ok(HavingValue::new(alias)) + } + + fn supported_value(node: Node<'_>) -> bool { + match node { + Node::A_Const(_) | Node::ParamRef(_) | Node::ColumnRef(_) => true, + Node::FuncCall(_) => Self::supported_aggregate(node), + _ => false, + } + } + + fn supported_aggregate(node: Node<'_>) -> bool { + Function::try_from(node) + .is_ok_and(|function| AggregateFunction::from_builtin_name(function.name).is_some()) + } + + fn comparison_operator(operator: &str) -> Option { + Some(match operator { + "=" => CompareOperator::Equal, + "<>" | "!=" => CompareOperator::NotEqual, + "<" => CompareOperator::Less, + "<=" => CompareOperator::LessOrEqual, + ">" => CompareOperator::Greater, + ">=" => CompareOperator::GreaterOrEqual, + _ => return None, + }) + } + + fn next_alias(&mut self) -> String { + loop { + let alias = format!("{ALIAS_PREFIX}{}", self.next_alias); + self.next_alias += 1; + if self.aliases.insert(alias.clone()) { + return alias; + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + fn rewrite(sql: &str) -> (String, HavingRewritePlan) { + let ast = pg_raw_parse::parse(sql).expect("parse HAVING query"); + let Node::SelectStmt(select) = ast.stmts().next().expect("one statement") else { + panic!("expected SELECT") + }; + let mut plan = None; + let rewritten = make::owned(|mem| { + let mut select = mem.make_unique(select); + plan = + Some(rewrite_select(&mut select.as_mut(), mem).expect("rewrite supported HAVING")); + select + }); + ( + pg_raw_parse::deparse(&*rewritten) + .expect("deparse rewritten query") + .as_str() + .to_owned(), + plan.expect("rewrite plan"), + ) + } + + #[test] + fn rewrites_single_comparison_with_parameter() { + let (sql, plan) = rewrite( + "SELECT customer_id, COUNT(*) FROM purchases \ + GROUP BY customer_id HAVING COUNT(*) > $1", + ); + + assert!(matches!( + plan.expression(), + HavingExpr::Compare { + operator: CompareOperator::Greater, + .. + } + )); + assert_eq!( + sql, + "SELECT customer_id, count(*), count(*) AS __pgdog_having_0, \ + CASE WHEN false THEN count(*) ELSE $1 END AS __pgdog_having_1, \ + customer_id AS __pgdog_having_2 \ + FROM purchases GROUP BY customer_id HAVING count(*) > $1 OR true" + ); + assert_eq!(plan.hidden_aliases().len(), 3); + } + + #[test] + fn rejects_boolean_null_test_and_cast_predicates() { + for sql in [ + "SELECT customer_id FROM purchases GROUP BY customer_id \ + HAVING COUNT(*) > 1 AND COUNT(*) < 5", + "SELECT customer_id FROM purchases GROUP BY customer_id \ + HAVING MAX(price) IS NULL", + "SELECT customer_id FROM purchases GROUP BY customer_id \ + HAVING COUNT(*)::numeric > 3", + "SELECT customer_id FROM purchases GROUP BY customer_id \ + HAVING MAX(name) ~ 'a'", + ] { + let ast = pg_raw_parse::parse(sql).unwrap(); + let Node::SelectStmt(select) = ast.stmts().next().unwrap() else { + panic!("expected SELECT") + }; + let result = make::try_owned(|mem| { + let mut select = mem.make_unique(select); + rewrite_select(&mut select.as_mut(), mem).map(|_| select) + }); + assert!(matches!(result, Err(Error::UnsupportedHaving(_))), "{sql}"); + } + } + + #[test] + fn parses_every_comparison_operator() { + for (sql, expected) in [ + ("=", CompareOperator::Equal), + ("<>", CompareOperator::NotEqual), + ("!=", CompareOperator::NotEqual), + ("<", CompareOperator::Less), + ("<=", CompareOperator::LessOrEqual), + (">", CompareOperator::Greater), + (">=", CompareOperator::GreaterOrEqual), + ] { + assert_eq!(Parser::comparison_operator(sql), Some(expected)); + } + assert_eq!(Parser::comparison_operator("~"), None); + } +} diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs index b2e600713..8e70331c0 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/mod.rs @@ -11,6 +11,7 @@ use pg_raw_parse::{Node, NodeMut, make, nodes, transform, walk}; pub(crate) mod aggregate; pub(crate) mod auto_id; pub(crate) mod error; +mod having; pub(crate) mod insert; pub(crate) mod insert_split_plan; pub(crate) mod nextval; diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/order_by.rs b/pgdog/src/frontend/router/parser/rewrite/statement/order_by.rs index d1c4bf534..60aea9925 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/order_by.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/order_by.rs @@ -4,7 +4,9 @@ use crate::backend::schema::Schema; use crate::frontend::router::parser::aggregate::AggregateFunction; use crate::frontend::router::parser::{Column, Function, OrderBy}; -use super::projection::{OrderByHelper, OrderBySource, ProjectionRewritePlan}; +use super::projection::{ + OrderByHelper, OrderBySource, ProjectionRewritePlan, column_name, is_star, target_output_name, +}; impl OrderBySource { fn name(&self) -> &str { @@ -14,21 +16,6 @@ impl OrderBySource { } } -fn column_name(column: &nodes::ColumnRef) -> Option<&str> { - column - .fields() - .into_iter() - .next_back() - .and_then(Node::as_str) -} - -fn is_star(column: &nodes::ColumnRef) -> bool { - matches!( - column.fields().into_iter().next_back(), - Some(Node::A_Star(_)) - ) -} - /// Compare complete column references so `a.price` and `b.price` remain distinct. fn same_column(left: &nodes::ColumnRef, right: &nodes::ColumnRef) -> bool { left.fields() @@ -59,16 +46,6 @@ fn star_covers(star: &nodes::ColumnRef, column: &nodes::ColumnRef) -> bool { .map(Node::as_str)))) } -fn target_output_name(target: &nodes::ResTarget) -> Option<&str> { - if let Some(alias) = target.name() { - return Some(alias); - } - match target.val() { - Node::ColumnRef(column) if !is_star(column) => column_name(column), - _ => None, - } -} - fn single_relation(select: &nodes::SelectStmtMut<'_, '_>) -> bool { let from = select.from_clause(); from.len() == 1 && matches!(from.first(), Some(Node::RangeVar(_))) @@ -404,6 +381,14 @@ mod tests { assert!(plan.order_by_helpers[0].injected); } + #[test] + fn injects_helper_when_cast_target_shares_output_name() { + let (sql, plan) = rewrite("SELECT price::int, price FROM purchases ORDER BY price"); + + assert!(sql.contains("price AS __pgdog_order_col0"), "{sql}"); + assert!(plan.order_by_helpers[0].injected); + } + #[test] fn injects_helper_for_unqualified_star_join() { let (sql, plan) = rewrite("SELECT * FROM a JOIN b ON a.id = b.a_id ORDER BY b.price"); diff --git a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs index 55c75157c..d1c6c0339 100644 --- a/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs +++ b/pgdog/src/frontend/router/parser/rewrite/statement/projection.rs @@ -6,13 +6,16 @@ use super::Error; use super::aggregate::{AggregatesRewrite, HelperKind}; +use super::having; use super::offset::{self, OffsetPlan}; use super::order_by; use crate::backend::schema::Schema; -use crate::frontend::router::parser::{Aggregate, HavingExpr, HavingRewritePlan, OrderBy}; +use crate::frontend::router::parser::{ + Aggregate, AggregateFunction, Function, HavingExpr, HavingRewritePlan, OrderBy, +}; use crate::frontend::{ClientRequest, PreparedStatements}; use crate::net::{ProtocolMessage, RowDescription}; -use pg_raw_parse::{Node, StmtList, make}; +use pg_raw_parse::{Node, StmtList, make, nodes}; use std::collections::BTreeSet; use std::sync::Arc; @@ -68,7 +71,6 @@ pub(crate) struct HavingRewriteResult { } impl HavingRewriteResult { - #[cfg_attr(not(test), allow(dead_code))] pub(crate) fn new(rewrite: HavingRewritePlan, aggregate: Aggregate) -> Self { Self { rewrite, aggregate } } @@ -81,12 +83,42 @@ impl HavingRewriteResult { self.rewrite.hidden_aliases() } - #[allow(dead_code)] pub(crate) fn aggregate(&self) -> &Aggregate { &self.aggregate } } +pub(super) fn column_name(column: &nodes::ColumnRef) -> Option<&str> { + column + .fields() + .into_iter() + .next_back() + .and_then(Node::as_str) +} + +pub(super) fn is_star(column: &nodes::ColumnRef) -> bool { + matches!( + column.fields().into_iter().next_back(), + Some(Node::A_Star(_)) + ) +} + +/// Column name PostgreSQL assigns to a SELECT target in the RowDescription. +pub(super) fn target_output_name(target: &nodes::ResTarget) -> Option<&str> { + target + .name() + .or_else(|| expression_output_name(target.val())) +} + +fn expression_output_name(node: Node<'_>) -> Option<&str> { + match node { + Node::ColumnRef(column) if !is_star(column) => column_name(column), + Node::FuncCall(function) => function.funcname().iter().next_back()?.as_str(), + Node::TypeCast(cast) => expression_output_name(cast.arg()), + _ => None, + } +} + impl ProjectionRewritePlan { pub(crate) fn is_noop(&self) -> bool { self.aggregate_helpers.is_empty() @@ -207,6 +239,42 @@ pub(crate) fn finalize_after_route( Ok(()) } +/// Query shapes that would return wrong rows once HAVING moves after the merge. +fn unsupported_with_having(select: &nodes::SelectStmt, schema: &Schema) -> Option<&'static str> { + // Aggregates are merged by SELECT list position, which star expansion shifts. + let expands_star = select + .target_list() + .iter() + .any(|target| match target.val() { + Node::ColumnRef(column) => is_star(column), + Node::A_Indirection(indirection) => matches!( + indirection.indirection().into_iter().next_back(), + Some(Node::A_Star(_)) + ), + _ => false, + }); + if expands_star { + return Some("SELECT * is unsupported"); + } + // Only whole aggregate calls are merged; an expression over them keeps one shard's value. + let wraps_aggregate = select.target_list().iter().any(|target| { + order_by::contains_aggregate(target.val(), schema) + && !Function::try_from(target.val()) + .is_ok_and(|function| AggregateFunction::from_name(function.name, schema).is_some()) + }); + if wraps_aggregate { + return Some("expressions over aggregates in the SELECT list are unsupported"); + } + // Shards would drop groups before PgDog merges and filters them. + if !matches!(select.limit_count(), Node::None) || !matches!(select.limit_offset(), Node::None) { + return Some("LIMIT and OFFSET are unsupported"); + } + if !matches!(select.distinct_clause().first(), None | Some(Node::None)) { + return Some("DISTINCT ON is unsupported"); + } + None +} + fn build( ast: &StmtList, schema: &Schema, @@ -215,16 +283,19 @@ fn build( let Some(Node::SelectStmt(select)) = ast.stmts().next() else { return Ok(None); }; - if !matches!(select.having_clause(), Node::None) { - return Err(Error::UnsupportedHaving( - "aggregates would be filtered on each shard before merging", - )); - } - - let aggregate = Aggregate::parse(select, schema); - if aggregate.is_empty() && select.sort_clause().is_empty() && !rewrite_offset { - return Ok(None); + let has_having = !matches!(select.having_clause(), Node::None); + if has_having && let Some(reason) = unsupported_with_having(select, schema) { + return Err(Error::UnsupportedHaving(reason)); } + let aggregate = if has_having { + None + } else { + let aggregate = Aggregate::parse(select, schema); + if aggregate.is_empty() && select.sort_clause().is_empty() && !rewrite_offset { + return Ok(None); + } + Some(aggregate) + }; // We can't sort by an expression with an aggregate in it, e.g., `ORDER BY count(*)`. // Each shard runs the aggregate ONLY on the rows stored directly on it, e.g., Shard A counts 2 'free' rows and Shard B counts 3. @@ -241,17 +312,45 @@ fn build( } let mut plan = ProjectionRewritePlan::default(); - let rewritten = make::owned(|mem| { + let rewritten = make::try_owned(|mem| { let mut select = mem.make_unique(select); + if has_having { + let having = having::rewrite_select(&mut select.as_mut(), mem)?; + let rewritten_aggregate = Aggregate::parse(&select, schema); + if rewritten_aggregate.is_empty() { + return Err(Error::UnsupportedHaving( + "at least one aggregate function is required", + )); + } + // Aggregate::parse skips GROUP BY items it cannot resolve. Require + // every item here so unrelated groups are never merged together. + if rewritten_aggregate.group_by().len() != select.group_clause().len() { + return Err(Error::UnsupportedHaving( + "every GROUP BY expression must be a selected column or ordinal", + )); + } + plan.having = Some(Box::new(HavingRewriteResult::new( + having, + rewritten_aggregate, + ))); + } + let aggregate = plan + .having + .as_deref() + .map(HavingRewriteResult::aggregate) + .or(aggregate.as_ref()) + .expect("non-HAVING aggregate parsed before projection rewrite"); if !aggregate.is_empty() { - plan = AggregatesRewrite::rewrite_select(&mut select.as_mut(), mem, &aggregate); + let aggregate_plan = + AggregatesRewrite::rewrite_select(&mut select.as_mut(), mem, aggregate); + plan.aggregate_helpers = aggregate_plan.aggregate_helpers; } order_by::rewrite_select(&mut select.as_mut(), mem, &mut plan); if rewrite_offset { offset::rewrite_select(&mut select.as_mut(), mem); } - select - }); + Ok::<_, Error>(select) + })?; if plan.is_noop() && !rewrite_offset { return Ok(None); } @@ -321,4 +420,127 @@ mod tests { ]); assert!(plan.drop_columns(&row_description).is_empty()); } + + #[test] + fn having_internal_aliases_do_not_collide_with_client_aliases() { + let ast = pg_raw_parse::parse( + "SELECT __pgdog_having_0, COUNT(*) \ + FROM purchases GROUP BY __pgdog_having_0 HAVING COUNT(*) > 2", + ) + .unwrap(); + let rewrite = build(&ast, &Schema::default(), false) + .unwrap() + .expect("HAVING requires a rewrite"); + + assert_eq!( + rewrite + .plan + .as_deref() + .expect("projection plan") + .having + .as_ref() + .expect("HAVING plan") + .hidden_aliases(), + &["__pgdog_having_1", "__pgdog_having_2", "__pgdog_having_3"] + ); + } + + #[test] + fn having_projects_unselected_grouping_column() { + let ast = pg_raw_parse::parse( + "SELECT COUNT(*) FROM purchases \ + GROUP BY customer_id HAVING COUNT(*) > 2", + ) + .unwrap(); + let rewrite = build(&ast, &Schema::default(), false) + .unwrap() + .expect("HAVING requires a rewrite"); + let plan = rewrite.plan.expect("projection plan"); + let aggregate = plan.having.as_ref().expect("HAVING plan"); + + assert_eq!(aggregate.aggregate().group_by().len(), 1); + assert_eq!(aggregate.aggregate().targets().len(), 2); + } + + #[test] + fn having_avg_receives_independent_count_helper() { + let ast = pg_raw_parse::parse( + "SELECT customer_id, AVG(price) FROM purchases \ + GROUP BY customer_id HAVING AVG(price) > 10", + ) + .unwrap(); + let rewrite = build(&ast, &Schema::default(), false) + .unwrap() + .expect("HAVING requires a rewrite"); + + assert_eq!( + rewrite + .plan + .as_deref() + .expect("projection plan") + .having + .as_ref() + .expect("HAVING plan") + .aggregate() + .targets() + .len(), + 2 + ); + assert_eq!( + rewrite + .plan + .as_deref() + .expect("projection plan") + .aggregate_helpers + .len(), + 2, + "selected and hidden AVG each require their own COUNT" + ); + } + + #[test] + fn aggregate_used_only_by_having_is_merged_and_hidden() { + let ast = pg_raw_parse::parse( + "SELECT customer_id, 'total' FROM purchases \ + GROUP BY 1 HAVING COUNT(*) > 2", + ) + .unwrap(); + let rewrite = build(&ast, &Schema::default(), false) + .unwrap() + .expect("HAVING requires a rewrite"); + let plan = rewrite.plan.expect("projection plan"); + let aggregate = plan.having.as_ref().expect("HAVING plan"); + + assert_eq!(aggregate.aggregate().targets().len(), 1); + assert_eq!(aggregate.aggregate().group_by().len(), 1); + } + + #[test] + fn unsupported_having_queries_fail_closed() { + for sql in [ + "SELECT customer_id, SUM(price) FROM purchases \ + GROUP BY customer_id HAVING SUM(price) / COUNT(*) > 2", + "SELECT customer_id FROM purchases \ + GROUP BY customer_id HAVING customer_id > 2", + "SELECT lower(customer_id), COUNT(*) FROM purchases \ + GROUP BY lower(customer_id) HAVING COUNT(*) > 2", + "SELECT * FROM purchases GROUP BY customer_id HAVING COUNT(*) > 2", + "SELECT p.*, COUNT(*) FROM purchases p GROUP BY customer_id HAVING COUNT(*) > 2", + "SELECT (p).*, COUNT(*) FROM purchases p GROUP BY customer_id HAVING COUNT(*) > 2", + "SELECT customer_id, COUNT(*) FROM purchases GROUP BY 'x' HAVING COUNT(*) > 2", + "SELECT name AS n, COUNT(*) FROM purchases GROUP BY n HAVING COUNT(*) > 2", + "SELECT customer_id, SUM(price) / COUNT(*) FROM purchases \ + GROUP BY customer_id HAVING COUNT(*) > 2", + "SELECT customer_id FROM purchases GROUP BY customer_id \ + HAVING COUNT(*) > 2 LIMIT 1", + "SELECT customer_id FROM purchases GROUP BY customer_id \ + HAVING COUNT(*) > 2 OFFSET 1", + "SELECT DISTINCT ON (customer_id) customer_id FROM purchases \ + GROUP BY customer_id HAVING COUNT(*) > 2", + ] { + let ast = pg_raw_parse::parse(sql).unwrap(); + let error = build(&ast, &Schema::default(), false).unwrap_err(); + assert!(matches!(error, Error::UnsupportedHaving(_)), "{sql}"); + } + } } diff --git a/pgdog/src/frontend/router/parser/route.rs b/pgdog/src/frontend/router/parser/route.rs index 5e19e0351..897e7f176 100644 --- a/pgdog/src/frontend/router/parser/route.rs +++ b/pgdog/src/frontend/router/parser/route.rs @@ -1,8 +1,10 @@ use std::{fmt::Display, ops::Deref, sync::Arc}; use super::{ - Aggregate, DistinctBy, Limit, OrderBy, StatementType, advisory_lock::AdvisoryLocks, - explain_trace::ExplainTrace, rewrite::statement::projection::ProjectionRewritePlan, + Aggregate, DistinctBy, Limit, OrderBy, StatementType, + advisory_lock::AdvisoryLocks, + explain_trace::ExplainTrace, + rewrite::statement::projection::{HavingRewriteResult, ProjectionRewritePlan}, }; use crate::frontend::{client::query_engine::TempTableChange, router::sharding::PendingLookup}; use lazy_static::lazy_static; @@ -270,7 +272,11 @@ impl Route { } pub(crate) fn aggregate(&self) -> &Aggregate { - &self.aggregate + self.projection_rewrite_plan + .as_deref() + .and_then(|plan| plan.having.as_deref()) + .map(HavingRewriteResult::aggregate) + .unwrap_or(&self.aggregate) } /// Set shard on this route, along with reasoning