hasura · hallettj · Aug 14, 2024 · Aug 13, 2024 · Aug 13, 2024 · Aug 13, 2024
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,8 @@ This changelog documents the changes between release versions.
 
 ## [Unreleased]
 - Accept predicate arguments in native mutations and native queries ([#92](https://github.com/hasura/ndc-mongodb/pull/92))
+- Serialize aggregate results as simple JSON (instead of Extended JSON) for
+  consistency with non-aggregate result serialization ([#96](https://github.com/hasura/ndc-mongodb/pull/96))
 
 ## [1.0.0] - 2024-07-09
 

diff --git a/crates/integration-tests/src/tests/aggregation.rs b/crates/integration-tests/src/tests/aggregation.rs
@@ -0,0 +1,39 @@
+use insta::assert_yaml_snapshot;
+use serde_json::json;
+
+use crate::graphql_query;
+
+#[tokio::test]
+async fn runs_aggregation_over_top_level_fields() -> anyhow::Result<()> {
+    assert_yaml_snapshot!(
+        graphql_query(
+            r#"
+                query($albumId: Int!) {
+                  track(order_by: { id: Asc }, where: { albumId: { _eq: $albumId } }) {
+                    milliseconds
+                    unitPrice
+                  }
+                  trackAggregate(
+                    filter_input: { order_by: { id: Asc }, where: { albumId: { _eq: $albumId } } }
+                  ) {
+                    _count
+                    milliseconds {
+                      _avg
+                      _max
+                      _min
+                      _sum
+                    }
+                    unitPrice {
+                      _count
+                      _count_distinct
+                    }
+                  }
+                }
+            "#
+        )
+        .variables(json!({ "albumId": 9 }))
+        .run()
+        .await?
+    );
+    Ok(())
+}
diff --git a/crates/integration-tests/src/tests/mod.rs b/crates/integration-tests/src/tests/mod.rs
@@ -7,6 +7,7 @@
 //     rust-analyzer.cargo.allFeatures = true
 //
 
+mod aggregation;
 mod basic;
 mod local_relationship;
 mod native_mutation;

diff --git a/...pshots/integration_tests__tests__aggregation__runs_aggregation_over_top_level_fields.snap b/...pshots/integration_tests__tests__aggregation__runs_aggregation_over_top_level_fields.snap
@@ -0,0 +1,33 @@
+---
+source: crates/integration-tests/src/tests/aggregation.rs
+expression: "graphql_query(r#\"\n                query($albumId: Int!) {\n                  track(order_by: { id: Asc }, where: { albumId: { _eq: $albumId } }) {\n                    milliseconds\n                    unitPrice\n                  }\n                  trackAggregate(\n                    filter_input: { order_by: { id: Asc }, where: { albumId: { _eq: $albumId } } }\n                  ) {\n                    _count\n                    milliseconds {\n                      _avg\n                      _max\n                      _min\n                      _sum\n                    }\n                    unitPrice {\n                      _count\n                      _count_distinct\n                    }\n                  }\n                }\n            \"#).variables(json!({\n                        \"albumId\": 9\n                    })).run().await?"
+---
+data:
+  track:
+    - milliseconds: 221701
+      unitPrice: "0.99"
+    - milliseconds: 436453
+      unitPrice: "0.99"
+    - milliseconds: 374543
+      unitPrice: "0.99"
+    - milliseconds: 322925
+      unitPrice: "0.99"
+    - milliseconds: 288208
+      unitPrice: "0.99"
+    - milliseconds: 308035
+      unitPrice: "0.99"
+    - milliseconds: 369345
+      unitPrice: "0.99"
+    - milliseconds: 350197
+      unitPrice: "0.99"
+  trackAggregate:
+    _count: 8
+    milliseconds:
+      _avg: 333925.875
+      _max: 436453
+      _min: 221701
+      _sum: 2671407
+    unitPrice:
+      _count: 8
+      _count_distinct: 1
+errors: ~
diff --git a/crates/mongodb-agent-common/src/query/foreach.rs b/crates/mongodb-agent-common/src/query/foreach.rs
@@ -261,25 +261,17 @@ mod tests {
         ]);
 
         let expected_response = query_response()
-            .row_set(
-                row_set()
-                    .aggregates([("count", json!({ "$numberInt": "2" }))])
-                    .rows([
-                        [
-                            ("albumId", json!(1)),
-                            ("title", json!("For Those About To Rock We Salute You")),
-                        ],
-                        [("albumId", json!(4)), ("title", json!("Let There Be Rock"))],
-                    ]),
-            )
-            .row_set(
-                row_set()
-                    .aggregates([("count", json!({ "$numberInt": "2" }))])
-                    .rows([
-                        [("albumId", json!(2)), ("title", json!("Balls to the Wall"))],
-                        [("albumId", json!(3)), ("title", json!("Restless and Wild"))],
-                    ]),
-            )
+            .row_set(row_set().aggregates([("count", json!(2))]).rows([
+                [
+                    ("albumId", json!(1)),
+                    ("title", json!("For Those About To Rock We Salute You")),
+                ],
+                [("albumId", json!(4)), ("title", json!("Let There Be Rock"))],
+            ]))
+            .row_set(row_set().aggregates([("count", json!(2))]).rows([
+                [("albumId", json!(2)), ("title", json!("Balls to the Wall"))],
+                [("albumId", json!(3)), ("title", json!("Restless and Wild"))],
+            ]))
             .build();
 
         let db = mock_aggregate_response_for_pipeline(
@@ -307,7 +299,7 @@ mod tests {
         );
 
         let result = execute_query_request(db, &music_config(), query_request).await?;
-        assert_eq!(expected_response, result);
+        assert_eq!(result, expected_response);
 
         Ok(())
     }
@@ -370,8 +362,8 @@ mod tests {
         ]);
 
         let expected_response = query_response()
-            .row_set(row_set().aggregates([("count", json!({ "$numberInt": "2" }))]))
-            .row_set(row_set().aggregates([("count", json!({ "$numberInt": "2" }))]))
+            .row_set(row_set().aggregates([("count", json!(2))]))
+            .row_set(row_set().aggregates([("count", json!(2))]))
             .build();
 
         let db = mock_aggregate_response_for_pipeline(

diff --git a/crates/mongodb-agent-common/src/query/mod.rs b/crates/mongodb-agent-common/src/query/mod.rs
@@ -103,10 +103,7 @@ mod tests {
             .into();
 
         let expected_response = row_set()
-            .aggregates([
-                ("count", json!({ "$numberInt": "11" })),
-                ("avg", json!({ "$numberInt": "3" })),
-            ])
+            .aggregates([("count", json!(11)), ("avg", json!(3))])
             .into_response();
 
         let expected_pipeline = bson!([
@@ -175,7 +172,7 @@ mod tests {
             .into();
 
         let expected_response = row_set()
-            .aggregates([("avg", json!({ "$numberDouble": "3.1" }))])
+            .aggregates([("avg", json!(3.1))])
             .row([("student_gpa", 3.1)])
             .into_response();
 

diff --git a/crates/mongodb-agent-common/src/query/relations.rs b/crates/mongodb-agent-common/src/query/relations.rs
@@ -613,7 +613,7 @@ mod tests {
                 "students_aggregate",
                 json!({
                     "aggregates": {
-                        "aggregate_count": { "$numberInt": "2" }
+                        "aggregate_count": 2
                     }
                 }),
             )])

diff --git a/crates/mongodb-agent-common/src/query/response.rs b/crates/mongodb-agent-common/src/query/response.rs
@@ -135,10 +135,10 @@ fn serialize_row_set_with_aggregates(
 fn serialize_aggregates(
     mode: ExtendedJsonMode,
     path: &[&str],
-    _query_aggregates: &IndexMap<ndc_models::FieldName, Aggregate>,
+    query_aggregates: &IndexMap<ndc_models::FieldName, Aggregate>,
     value: Bson,
 ) -> Result<IndexMap<ndc_models::FieldName, serde_json::Value>> {
-    let aggregates_type = type_for_aggregates()?;
+    let aggregates_type = type_for_aggregates(query_aggregates);
     let json = bson_to_json(mode, &aggregates_type, value)?;
 
     // The NDC type uses an IndexMap for aggregate values; we need to convert the map
@@ -184,8 +184,8 @@ fn type_for_row_set(
 ) -> Result<Type> {
     let mut type_fields = BTreeMap::new();
 
-    if aggregates.is_some() {
-        type_fields.insert("aggregates".into(), type_for_aggregates()?);
+    if let Some(aggregates) = aggregates {
+        type_fields.insert("aggregates".into(), type_for_aggregates(aggregates));
     }
 
     if let Some(query_fields) = fields {
@@ -199,9 +199,25 @@ fn type_for_row_set(
     }))
 }
 
-// TODO: infer response type for aggregates MDB-130
-fn type_for_aggregates() -> Result<Type> {
-    Ok(Type::Scalar(MongoScalarType::ExtendedJSON))
+fn type_for_aggregates(query_aggregates: &IndexMap<ndc_models::FieldName, Aggregate>) -> Type {
+    let fields = query_aggregates
+        .iter()
+        .map(|(field_name, aggregate)| {
+            (
+                field_name.to_string().into(),
+                match aggregate {
+                    Aggregate::ColumnCount { .. } => {
+                        Type::Scalar(MongoScalarType::Bson(mongodb_support::BsonScalarType::Int))
+                    }
+                    Aggregate::StarCount => {
+                        Type::Scalar(MongoScalarType::Bson(mongodb_support::BsonScalarType::Int))
+                    }
+                    Aggregate::SingleColumn { result_type, .. } => result_type.clone(),
+                },
+            )
+        })
+        .collect();
+    Type::Object(ObjectType { fields, name: None })
 }
 
 fn type_for_row(