Error when encountering an unsupported aggregate function (#1020)

When handling aggregate functions in cross-shard queries, we can only
return a correct result for functions that we have explicit support for.
No matter what the aggregate function is doing, we need to know how to
combine the results from the separate shards.

Prior to this change, we would just silently do the wrong thing,
treating this as any other non-aggregate expression and just returning a
union of the rows from each query. We now explicitly check if an
aggregate function with that name exists and error if we don't recognize
it.

This is future proofed against new functions in later postgres versions,
as well as user defined aggregates (which we can likely never support).
This will produce a false positive if a function exists and is defined
as aggregate for some argument types but not others. But frankly anyone
doing that is asking for trouble.

This logic is objectively in the wrong place. What we are doing here is
validation, not parsing. However as I'm familiarizing myself with these
code paths and preparing a larger rearchitecture, I'm slowly hammering
things into a shape that's easier to move around. I do not intend on
leaving this logic here long term.

Author: sgrif

integration/ci/cache-key.sh

@@ -13,6 +13,7 @@ REPO_ROOT="$( cd "${SCRIPT_DIR}/../.." && pwd )"
 cd "$REPO_ROOT"
 
 files=(
+    rust-toolchain.toml
     Cargo.lock
     Cargo.toml
     .cargo/config.toml

integration/go/go_pgx/sharded_test.go

@@ -152,11 +152,11 @@ func TestShardedTwoPc(t *testing.T) {
 		assert.NoError(t, err)
 	}
 
-	// +4 is for schema sync
+	// +5 is for schema sync
 	assertShowField(t, "SHOW STATS", "total_xact_2pc_count", 200, "pgdog_2pc", "pgdog_sharded", 0, "primary")
 	assertShowField(t, "SHOW STATS", "total_xact_2pc_count", 200, "pgdog_2pc", "pgdog_sharded", 1, "primary")
-	assertShowField(t, "SHOW STATS", "total_xact_count", 401+4, "pgdog_2pc", "pgdog_sharded", 0, "primary") // PREPARE, COMMIT for each transaction + TRUNCATE
-	assertShowField(t, "SHOW STATS", "total_xact_count", 401+4, "pgdog_2pc", "pgdog_sharded", 1, "primary")
+	assertShowField(t, "SHOW STATS", "total_xact_count", 401+5, "pgdog_2pc", "pgdog_sharded", 0, "primary") // PREPARE, COMMIT for each transaction + TRUNCATE
+	assertShowField(t, "SHOW STATS", "total_xact_count", 401+5, "pgdog_2pc", "pgdog_sharded", 1, "primary")
 
 	for i := range 200 {
 		rows, err := conn.Query(

integration/python/test_session_mode.py

@@ -356,3 +356,14 @@ async def test_no_search_path_session_mode():
 
         async with conn.transaction():
             await conn.execute("SELECT * FROM shard_0.py_test_no_search_path_session_mode LIMIT 1")
+
+@pytest.mark.asyncio
+async def test_unrecognized_aggregate_function_works_on_schema_based_sharding():
+    conn = await async_session_conn("shard_0")
+    await conn.execute("DROP AGGREGATE IF EXISTS pgdog_sum_python(int4)")
+    await conn.execute("CREATE AGGREGATE pgdog_sum_python (int4) (sfunc = int4_sum, stype = bigint)")
+    await conn.execute("DROP TABLE IF EXISTS unrecognized_agg_test_python")
+    await conn.execute("CREATE TABLE unrecognized_agg_test_python (lol int4)")
+
+    async with conn.transaction():
+        await conn.execute("SELECT pgdog_sum_python(lol) FROM unrecognized_agg_test_python")

integration/rust/tests/integration/mod.rs

@@ -36,3 +36,4 @@ pub mod timestamp_sorting;
 pub mod tls_enforced;
 pub mod tls_reload;
 pub mod transaction_state;
+pub mod unrecognized_aggregate;

integration/rust/tests/integration/unrecognized_aggregate.rs

@@ -0,0 +1,46 @@
+use rust::setup::{admin_sqlx, connections_sqlx};
+use sqlx::Executor;
+use std::assert_matches;
+
+async fn define_custom_aggregate_fn() {
+    for connection in connections_sqlx().await {
+        connection
+            .execute("DROP AGGREGATE IF EXISTS pgdog_sum (int4)")
+            .await
+            .unwrap();
+        connection
+            .execute("CREATE AGGREGATE pgdog_sum (int4) (sfunc = int4_sum, stype = bigint)")
+            .await
+            .unwrap();
+        connection
+            .execute("DROP TABLE IF EXISTS unrecognized_agg_test")
+            .await
+            .unwrap();
+        connection
+            .execute("CREATE TABLE unrecognized_agg_test (lol int4, customer_id bigint)")
+            .await
+            .unwrap();
+    }
+    admin_sqlx().await.execute("RELOAD").await.unwrap();
+}
+
+#[tokio::test]
+async fn unrecognized_aggregate_function_errors_only_on_cross_shard_queries() {
+    define_custom_aggregate_fn().await;
+    let mut connections = connections_sqlx().await;
+    let sharded = connections.pop().unwrap();
+    let unsharded = connections.pop().unwrap();
+
+    let unsharded_query = unsharded
+        .fetch_one("SELECT pgdog_sum(lol) FROM unrecognized_agg_test")
+        .await;
+    assert_matches!(unsharded_query, Ok(_));
+
+    let sharded_query = sharded
+        .fetch_one("SELECT pgdog_sum(lol) FROM unrecognized_agg_test")
+        .await;
+    let err = sharded_query
+        .err()
+        .expect("unrecognized aggregate executed successfully");
+    assert!(err.to_string().contains("pgdog_sum is not yet supported"));
+}

mise.toml

@@ -1,4 +1,4 @@
 [tools]
-"rust" = { version = "1.91.0" }
+"rust" = { version = "1.96.0" }
 "cargo:cargo-nextest" = "latest"
-"cargo:cargo-watch" = "latest"
+"cargo:cargo-watch" = "latest"

pgdog-stats/src/schema.rs

@@ -1,6 +1,11 @@
 use indexmap::IndexMap;
 use serde::{Deserialize, Serialize};
-use std::{collections::HashMap, hash::Hash, ops::Deref, sync::Arc};
+use std::{
+    collections::{HashMap, HashSet},
+    hash::Hash,
+    ops::Deref,
+    sync::Arc,
+};
 
 /// Schema name -> Table name -> Relation
 pub type Relations = HashMap<String, HashMap<String, Relation>>;
@@ -134,6 +139,7 @@ impl Relation {
 pub struct SchemaInner {
     pub search_path: Vec<String>,
     pub relations: Relations,
+    pub aggregate_functions: HashSet<String>,
 }
 
 impl Hash for SchemaInner {

pgdog/src/backend/pool/connection/aggregate.rs

@@ -176,6 +176,7 @@ impl<'a> Accumulator<'a> {
                     }
                 }
             }
+            AggregateFunction::Unrecognized(..) => return Ok(false),
         }
 
         Ok(true)
@@ -512,6 +513,13 @@ impl<'a> Aggregates<'a> {
                             false
                         })
                 }
+                AggregateFunction::Unrecognized(fname) => {
+                    unsupported.get_or_insert(UnsupportedAggregate {
+                        function: fname.clone(),
+                        reason: format!("{fname} is not yet supported"),
+                    });
+                    false
+                }
                 _ => true,
             }
         });
@@ -803,7 +811,7 @@ mod test {
     }
 
     fn parse(stmt: &str) -> Aggregate {
-        Aggregate::parse(&select(stmt))
+        Aggregate::parse(&select(stmt), &Default::default())
     }
 
     #[test]

pgdog/src/backend/schema/mod.rs

@@ -60,9 +60,18 @@ impl Schema {
             .map(|p| p.trim().replace("\"", ""))
             .collect();
 
+        let aggregate_functions = server
+            .fetch_all::<String>(
+                "SELECT DISTINCT proname FROM pg_proc INNER JOIN pg_aggregate ON oid = aggfnoid",
+            )
+            .await?
+            .into_iter()
+            .collect();
+
         let inner = SchemaInner {
             search_path,
             relations,
+            aggregate_functions,
         };
 
         Ok(Self {
@@ -79,6 +88,15 @@ impl Schema {
     pub(crate) fn from_parts(
         search_path: Vec<String>,
         relations: HashMap<(String, String), Relation>,
+    ) -> Self {
+        Self::from_parts_with_agg(search_path, relations, Vec::new())
+    }
+
+    #[cfg(test)]
+    pub(crate) fn from_parts_with_agg(
+        search_path: Vec<String>,
+        relations: HashMap<(String, String), Relation>,
+        aggregate_functions: Vec<String>,
     ) -> Self {
         let mut nested: StatsRelations = HashMap::new();
         for ((schema, name), relation) in relations {
@@ -91,6 +109,7 @@ impl Schema {
             inner: StatsSchema::new(SchemaInner {
                 search_path,
                 relations: nested,
+                aggregate_functions: aggregate_functions.into_iter().collect(),
             }),
         }
     }
@@ -555,4 +574,24 @@ mod test {
             );
         }
     }
+
+    #[tokio::test]
+    async fn test_loading_aggregate_functions() {
+        let mut server = test_server().await;
+        server.execute_checked("BEGIN").await.unwrap();
+
+        let schema = Schema::load(&mut server).await.unwrap();
+        assert!(!schema
+            .aggregate_functions
+            .contains(&String::from("pgdog_sum")));
+
+        server
+            .execute_checked("CREATE AGGREGATE pgdog_sum (int4) (sfunc = int4_sum, stype = bigint)")
+            .await
+            .unwrap();
+        let schema = Schema::load(&mut server).await.unwrap();
+        assert!(schema
+            .aggregate_functions
+            .contains(&String::from("pgdog_sum")));
+    }
 }

pgdog/src/frontend/router/parser/aggregate.rs

@@ -2,7 +2,8 @@ use pg_query::protobuf::Integer;
 use pg_query::protobuf::{a_const::Val, Node, SelectStmt, String as PgQueryString};
 use pg_query::NodeEnum;
 
-use crate::frontend::router::parser::{ExpressionRegistry, Function};
+use super::{ExpressionRegistry, Function};
+use crate::backend::schema::Schema;
 
 #[derive(Debug, Clone, PartialEq)]
 pub struct AggregateTarget {
@@ -41,10 +42,11 @@ pub enum AggregateFunction {
     StddevSamp,
     VarPop,
     VarSamp,
+    Unrecognized(String),
 }
 
 impl AggregateFunction {
-    pub fn as_str(&self) -> &'static str {
+    pub fn as_str(&self) -> &str {
         match self {
             AggregateFunction::Count => "count",
             AggregateFunction::Max => "max",
@@ -55,6 +57,7 @@ impl AggregateFunction {
             AggregateFunction::StddevSamp => "stddev_samp",
             AggregateFunction::VarPop => "var_pop",
             AggregateFunction::VarSamp => "var_samp",
+            AggregateFunction::Unrecognized(s) => &*s,
         }
     }
 }
@@ -121,7 +124,7 @@ fn columns_match(group_by_names: &[&String], select_names: &[&String]) -> bool {
 
 impl Aggregate {
     /// Figure out what aggregates are present and which ones PgDog supports.
-    pub fn parse(stmt: &SelectStmt) -> Self {
+    pub fn parse(stmt: &SelectStmt, schema: &Schema) -> Self {
         let mut targets = vec![];
         let mut registry = ExpressionRegistry::new();
         let group_by = stmt
@@ -168,7 +171,13 @@ impl Aggregate {
                             "stddev_pop" => Some(AggregateFunction::StddevPop),
                             "variance" | "var_samp" => Some(AggregateFunction::VarSamp),
                             "var_pop" => Some(AggregateFunction::VarPop),
-                            _ => None,
+                            fname => {
+                                if schema.aggregate_functions.contains(fname) {
+                                    Some(AggregateFunction::Unrecognized(fname.to_owned()))
+                                } else {
+                                    None
+                                }
+                            }
                         };
 
                         if let Some(function) = function {
@@ -259,7 +268,7 @@ mod test {
     }
 
     fn parse(stmt: &str) -> Aggregate {
-        Aggregate::parse(&select(stmt))
+        Aggregate::parse(&select(stmt), &Default::default())
     }
 
     #[test]
@@ -455,4 +464,34 @@ mod test {
         assert_eq!(aggr.group_by(), &[0]);
         assert_eq!(aggr.targets().len(), 1);
     }
+
+    #[test]
+    fn test_unrecognized_aggregate_function_errors() {
+        let schema_with_agg = Schema::from_parts_with_agg(
+            Vec::new(),
+            Default::default(),
+            vec![String::from("mysum")],
+        );
+        let schema_without_agg = Default::default();
+        let query = select("SELECT mysum(lol) FROM example");
+
+        // A random function that isn't listed as aggregate in the schema
+        // doesn't require special support on our end, so we should be fine.
+        let aggregate = Aggregate::parse(&query, &schema_without_agg);
+        assert_eq!(aggregate.targets, Vec::new());
+
+        // If we see an aggregate function we don't recognize, we can't
+        // process the query correctly, since we need to combine the
+        // results from each shard.
+        let aggregate = Aggregate::parse(&query, &schema_with_agg);
+        let funcs = aggregate
+            .targets
+            .into_iter()
+            .map(|t| t.function)
+            .collect::<Vec<_>>();
+        assert_eq!(
+            funcs,
+            vec![AggregateFunction::Unrecognized("mysum".to_owned())]
+        );
+    }
 }