From eb8987017c794a74e30f95c03c3a7525c073b3dd Mon Sep 17 00:00:00 2001
From: Jesse Hallett <jesse.hallett@hasura.io>
Date: Tue, 10 Dec 2024 17:39:32 -0800
Subject: [PATCH] remove double ifNull-wrapping on count results, update tests

---
 crates/mongodb-agent-common/src/query/mod.rs  | 42 ++++++++++++-------
 .../src/query/pipeline.rs                     | 19 ++++-----
 2 files changed, 37 insertions(+), 24 deletions(-)

diff --git a/crates/mongodb-agent-common/src/query/mod.rs b/crates/mongodb-agent-common/src/query/mod.rs
index 3353b572..d6094ca6 100644
--- a/crates/mongodb-agent-common/src/query/mod.rs
+++ b/crates/mongodb-agent-common/src/query/mod.rs
@@ -110,11 +110,11 @@ mod tests {
             {
                 "$facet": {
                     "avg": [
-                        { "$match": { "gpa": { "$exists": true, "$ne": null } } },
+                        { "$match": { "gpa": { "$ne": null } } },
                         { "$group": { "_id": null, "result": { "$avg": "$gpa" } } },
                     ],
                     "count": [
-                        { "$match": { "gpa": { "$exists": true, "$ne": null } } },
+                        { "$match": { "gpa": { "$ne": null } } },
                         { "$group": { "_id": "$gpa" } },
                         { "$count": "result" },
                     ],
@@ -123,10 +123,17 @@ mod tests {
             {
                 "$replaceWith": {
                     "aggregates": {
-                        "avg": { "$getField": {
-                            "field": "result",
-                            "input": { "$first": { "$getField": { "$literal": "avg" } } },
-                        } },
+                        "avg": {
+                            "$ifNull": [
+                                {
+                                    "$getField": {
+                                        "field": "result",
+                                        "input": { "$first": { "$getField": { "$literal": "avg" } } },
+                                    }
+                                },
+                                null
+                            ]
+                        },
                         "count": {
                             "$ifNull": [
                                 {
@@ -180,24 +187,31 @@ mod tests {
             { "$match": { "gpa": { "$lt": 4.0 } } },
             {
                 "$facet": {
-                    "avg": [
-                        { "$match": { "gpa": { "$exists": true, "$ne": null } } },
-                        { "$group": { "_id": null, "result": { "$avg": "$gpa" } } },
-                    ],
                     "__ROWS__": [{
                         "$replaceWith": {
                             "student_gpa": { "$ifNull": ["$gpa", null] },
                         },
                     }],
+                    "avg": [
+                        { "$match": { "gpa": { "$ne": null } } },
+                        { "$group": { "_id": null, "result": { "$avg": "$gpa" } } },
+                    ],
                 },
             },
             {
                 "$replaceWith": {
                     "aggregates": {
-                        "avg": { "$getField": {
-                            "field": "result",
-                            "input": { "$first": { "$getField": { "$literal": "avg" } } },
-                        } },
+                        "avg": {
+                            "$ifNull": [
+                                {
+                                    "$getField": {
+                                        "field": "result",
+                                        "input": { "$first": { "$getField": { "$literal": "avg" } } },
+                                    }
+                                },
+                                null
+                            ]
+                        },
                     },
                     "rows": "$__ROWS__",
                 },
diff --git a/crates/mongodb-agent-common/src/query/pipeline.rs b/crates/mongodb-agent-common/src/query/pipeline.rs
index c45c78ad..f89d2c8f 100644
--- a/crates/mongodb-agent-common/src/query/pipeline.rs
+++ b/crates/mongodb-agent-common/src/query/pipeline.rs
@@ -192,15 +192,10 @@ fn facet_pipelines_for_query(
             // has a field called `result`. This code selects each facet result by name, and pulls
             // out the `result` value.
             let value_expr = doc! {
-                "$ifNull": [
-                    {
-                        "$getField": {
-                            "field": RESULT_FIELD, // evaluates to the value of this field
-                            "input": { "$first": get_field(key.as_str()) }, // field is accessed from this document
-                        },
-                    },
-                    null,
-                ]
+                "$getField": {
+                    "field": RESULT_FIELD, // evaluates to the value of this field
+                    "input": { "$first": get_field(key.as_str()) }, // field is accessed from this document
+                },
             };
 
             // Matching SQL semantics, if a **count** aggregation does not match any rows we want
@@ -209,8 +204,12 @@ fn facet_pipelines_for_query(
                 doc! {
                     "$ifNull": [value_expr, 0],
                 }
+            // Otherwise if the aggregate value is missing because the aggregation applied to an
+            // empty document set then provide an explicit `null` value.
             } else {
-                value_expr
+                doc! {
+                    "$ifNull": [value_expr, null]
+                }
             };
 
             (key.to_string(), value_expr.into())