From 141ff826fd18ccbb1f93cee13bcc1a6a2f7a85da Mon Sep 17 00:00:00 2001
From: Maddie Zechar <mez2113@columbia.edu>
Date: Wed, 23 Nov 2022 17:39:10 +0000
Subject: SERVER-71386 Telemetry collection fails for queries with non-object
 fields

---
 .../redact_queries_with_nonobject_fields.js        | 75 ++++++++++++++++++++++
 src/mongo/db/query/telemetry.cpp                   | 43 ++++++++-----
 2 files changed, 103 insertions(+), 15 deletions(-)
 create mode 100644 jstests/telemetry/redact_queries_with_nonobject_fields.js

diff --git a/jstests/telemetry/redact_queries_with_nonobject_fields.js b/jstests/telemetry/redact_queries_with_nonobject_fields.js
new file mode 100644
index 00000000000..8bc7aad5a89
--- /dev/null
+++ b/jstests/telemetry/redact_queries_with_nonobject_fields.js
@@ -0,0 +1,75 @@
+/**
+ * Test that telemetry key generation works for queries with non-object fields.
+ */
+load('jstests/libs/analyze_plan.js');
+
+(function() {
+"use strict";
+
+// Turn on the collecting of telemetry metrics.
+let options = {
+    setParameter: "internalQueryConfigureTelemetrySamplingRate=2147483647",
+};
+
+const conn = MongoRunner.runMongod(options);
+const testDB = conn.getDB('test');
+var collA = testDB[jsTestName()];
+var collB = db[jsTestName() + 'Two'];
+collA.drop();
+collB.drop();
+
+for (var i = 0; i < 200; i++) {
+    collA.insert({foo: 0, bar: Math.floor(Math.random() * 3)});
+    collA.insert({foo: 1, bar: Math.floor(Math.random() * -2)});
+    collB.insert({foo: Math.floor(Math.random() * 2), bar: Math.floor(Math.random() * 2)});
+}
+
+function confirmAggSuccess(collName, pipeline) {
+    const command = {aggregate: collName, cursor: {}};
+    command.pipeline = pipeline;
+    assert.commandWorked(testDB.runCommand(command));
+}
+// Test with non-object fields $limit and $skip.
+confirmAggSuccess(collA.getName(), [{$sort: {bar: -1}}, {$limit: 2}, {$match: {foo: {$lte: 2}}}]);
+confirmAggSuccess(collA.getName(), [{$sort: {bar: -1}}, {$skip: 50}, {$match: {foo: {$lte: 2}}}]);
+confirmAggSuccess(collA.getName(),
+                  [{$sort: {bar: -1}}, {$limit: 2}, {$skip: 50}, {$match: {foo: 0}}]);
+
+// Test non-object field, $unionWith.
+confirmAggSuccess(collA.getName(), [{$unionWith: collB.getName()}]);
+
+// Test $limit in $setWindowFields for good measure.
+confirmAggSuccess(collA.getName(), [
+    {$_internalInhibitOptimization: {}},
+    {
+        $setWindowFields: {
+            sortBy: {foo: 1},
+            output: {sum: {$sum: "$bar", window: {documents: ["unbounded", "current"]}}}
+        }
+    },
+    {$sort: {foo: 1}},
+    {$limit: 5}
+]);
+// Test find commands containing non-object fields
+assert.commandWorked(testDB.runCommand({find: collA.getName(), limit: 20}));
+assert.commandWorked(testDB.runCommand({find: collA.getName(), skip: 199}));
+collA.find().skip(100);
+
+// findOne has a nonobject field, $limit.
+collB.findOne();
+collB.findOne({foo: 1});
+
+// Test non-object field $unwind
+confirmAggSuccess(
+    collA.getName(), [{
+        "$facet": {
+            "productOfJoin": [
+                {"$lookup": {"from": collB.getName(), "pipeline": [{"$match": {}}], "as": "join"}},
+                {"$unwind": "$join"},
+                {"$project": {"str": 1}}
+            ]
+        }
+    }]);
+
+MongoRunner.stopMongod(conn);
+}());
diff --git a/src/mongo/db/query/telemetry.cpp b/src/mongo/db/query/telemetry.cpp
index 7b8c2c36c89..05b9a46b351 100644
--- a/src/mongo/db/query/telemetry.cpp
+++ b/src/mongo/db/query/telemetry.cpp
@@ -290,6 +290,26 @@ std::string fleSafeFieldNameRedactor(const BSONElement& e) {
     return e.fieldNameStringData().toString();
 }
 
+/**
+ * Append the element to the builder and redact any literals within the element. The element may be
+ * of any type.
+ */
+void appendWithRedactedLiterals(BSONObjBuilder& builder, const BSONElement& el) {
+    if (el.type() == Object) {
+        builder.append(el.fieldNameStringData(), el.Obj().redact(false, fleSafeFieldNameRedactor));
+    } else if (el.type() == Array) {
+        BSONObjBuilder arrayBuilder = builder.subarrayStart(fleSafeFieldNameRedactor(el));
+        for (auto&& arrayElem : el.Obj()) {
+            appendWithRedactedLiterals(arrayBuilder, arrayElem);
+        }
+        arrayBuilder.done();
+    } else {
+        auto fieldName = fleSafeFieldNameRedactor(el);
+        builder.append(fieldName, "###"_sd);
+    }
+    builder.done();
+}
+
 }  // namespace
 
 const BSONObj& TelemetryMetrics::redactKey(const BSONObj& key) const {
@@ -380,11 +400,8 @@ void registerAggRequest(const AggregateCommandRequest& request, OperationContext
     BSONObjBuilder pipelineBuilder = telemetryKey.subarrayStart("pipeline"_sd);
     try {
         for (auto&& stage : request.getPipeline()) {
-            auto el = stage.firstElement();
             BSONObjBuilder stageBuilder = pipelineBuilder.subobjStart("stage"_sd);
-            stageBuilder.append(el.fieldNameStringData(),
-                                el.Obj().redact(false, fleSafeFieldNameRedactor));
-            stageBuilder.done();
+            appendWithRedactedLiterals(stageBuilder, stage.firstElement());
         }
         pipelineBuilder.done();
         telemetryKey.append("namespace", request.getNamespace().toString());
@@ -433,18 +450,14 @@ void registerFindRequest(const FindCommandRequest& request,
     }
 
     BSONObjBuilder telemetryKey;
-    BSONObjBuilder findBuilder = telemetryKey.subobjStart("find"_sd);
     try {
-        auto findBson = request.toBSON({});
-        for (auto&& findEntry : findBson) {
-            if (findEntry.isABSONObj()) {
-                telemetryKey.append(findEntry.fieldNameStringData(),
-                                    findEntry.Obj().redact(false, fleSafeFieldNameRedactor));
-            } else {
-                telemetryKey.append(findEntry.fieldNameStringData(), "###"_sd);
-            }
-        }
-        findBuilder.done();
+        // Serialize the request.
+        BSONObjBuilder serializedRequest;
+        BSONObjBuilder asElement = serializedRequest.subobjStart("find");
+        request.serialize({}, &asElement);
+        asElement.done();
+        // And append as an element to the telemetry key.
+        appendWithRedactedLiterals(telemetryKey, serializedRequest.obj().firstElement());
         telemetryKey.append("namespace", collection.toString());
         if (request.getReadConcern()) {
             telemetryKey.append("readConcern", *request.getReadConcern());
-- 
cgit v1.2.1