feat: rewrite cross-shard HAVING queries and enable post merge filtering

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<dyn std::error::Error>> {
+    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::<i64, _>("customer_id"), 1);
+    assert_eq!(rows[0].get::<i64, _>("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::<i64, _>("customer_id"), 2);
+
+    cleanup(&sharded).await;
+    Ok(())
+}
+
+#[tokio::test]
+async fn prepared_having_parameter_typed_and_untyped() -> Result<(), Box<dyn std::error::Error>> {
+    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::<i64, _>("customer_id"), 1);
+    assert_eq!(rows[0].get::<i64, _>("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<dyn std::error::Error>> {
+    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();
+    }
+}

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;

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();

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(*)",

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<Self> {
+        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<Self> {
-        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()))
         })
     }
 }

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)]

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),