SERVER-76367: Abstract request-specific shapifying logic into RequestShapifiers

1 files changed, 80 insertions, 24 deletions

diff --git a/src/mongo/db/query/query_shape.cpp b/src/mongo/db/query/query_shape.cpp
index ec5179ab96d..05d91f94279 100644
--- a/src/mongo/db/query/query_shape.cpp
+++ b/src/mongo/db/query/query_shape.cpp
@@ -28,14 +28,12 @@
  */
 
 #include "mongo/db/query/query_shape.h"
-#include "query_request_helper.h"
-#include "sort_pattern.h"
 
 #include "mongo/base/status.h"
 #include "mongo/db/query/find_command_gen.h"
-#include "mongo/db/query/plan_explainer.h"
 #include "mongo/db/query/projection_ast_util.h"
 #include "mongo/db/query/projection_parser.h"
+#include "mongo/db/query/query_request_helper.h"
 #include "mongo/db/query/sort_pattern.h"
 
 namespace mongo::query_shape {
@@ -70,9 +68,9 @@ BSONObj representativePredicateShape(
     return predicate->serialize(opts);
 }
 
-BSONObj sortShape(const BSONObj& sortSpec,
-                  const boost::intrusive_ptr<ExpressionContext>& expCtx,
-                  const SerializationOptions& opts) {
+BSONObj extractSortShape(const BSONObj& sortSpec,
+                         const boost::intrusive_ptr<ExpressionContext>& expCtx,
+                         const SerializationOptions& opts) {
     if (sortSpec.isEmpty()) {
         return sortSpec;
     }
@@ -160,7 +158,8 @@ void addRemainingFindCommandFields(BSONObjBuilder* bob,
         opts.appendLiteral(bob, FindCommandRequest::kCollationFieldName, collation);
     }
 }
-BSONObj redactHintComponent(BSONObj obj, const SerializationOptions& opts, bool redactValues) {
+
+BSONObj extractHintShape(BSONObj obj, const SerializationOptions& opts, bool redactValues) {
     BSONObjBuilder bob;
     for (BSONElement elem : obj) {
         if (hintSpecialField.compare(elem.fieldName()) == 0) {
@@ -191,9 +190,9 @@ BSONObj redactHintComponent(BSONObj obj, const SerializationOptions& opts, bool
  * In a let specification all field names are variable names, and all values are either expressions
  * or constants.
  */
-BSONObj redactLetSpec(BSONObj letSpec,
-                      const SerializationOptions& opts,
-                      boost::intrusive_ptr<ExpressionContext> expCtx) {
+BSONObj extractLetSpecShape(BSONObj letSpec,
+                            const SerializationOptions& opts,
+                            boost::intrusive_ptr<ExpressionContext> expCtx) {
 
     BSONObjBuilder bob;
     for (BSONElement elem : letSpec) {
@@ -206,26 +205,30 @@ BSONObj redactLetSpec(BSONObj letSpec,
     return bob.obj();
 }
 
+BSONObj extractNamespaceShape(NamespaceString nss, const SerializationOptions& opts) {
+    BSONObjBuilder bob;
+    if (nss.tenantId()) {
+        bob.append("tenantId", opts.serializeIdentifier(nss.tenantId().value().toString()));
+    }
+    bob.append("db", opts.serializeIdentifier(nss.db()));
+    bob.append("coll", opts.serializeIdentifier(nss.coll()));
+    return bob.obj();
+}
+
 BSONObj extractQueryShape(const FindCommandRequest& findCommand,
                           const SerializationOptions& opts,
                           const boost::intrusive_ptr<ExpressionContext>& expCtx) {
     BSONObjBuilder bob;
     // Serialize the namespace as part of the query shape.
     {
-        BSONObjBuilder cmdNs = bob.subobjStart("cmdNs");
         auto ns = findCommand.getNamespaceOrUUID();
         if (ns.nss()) {
-            auto nss = ns.nss().value();
-            if (nss.tenantId()) {
-                cmdNs.append("tenantId",
-                             opts.serializeIdentifier(nss.tenantId().value().toString()));
-            }
-            cmdNs.append("db", opts.serializeIdentifier(nss.db()));
-            cmdNs.append("coll", opts.serializeIdentifier(nss.coll()));
+            bob.append("cmdNs", extractNamespaceShape(ns.nss().value(), opts));
         } else {
+            BSONObjBuilder cmdNs = bob.subobjStart("cmdNs");
             cmdNs.append("uuid", opts.serializeIdentifier(ns.uuid()->toString()));
+            cmdNs.done();
         }
-        cmdNs.done();
     }
 
     // Redact the namespace of the command.
@@ -256,7 +259,7 @@ BSONObj extractQueryShape(const FindCommandRequest& findCommand,
 
     // Let Spec.
     if (auto letSpec = findCommand.getLet()) {
-        auto redactedObj = redactLetSpec(letSpec.get(), opts, expCtx);
+        auto redactedObj = extractLetSpecShape(letSpec.get(), opts, expCtx);
         auto ownedObj = redactedObj.getOwned();
         bob.append(FindCommandRequest::kLetFieldName, std::move(ownedObj));
     }
@@ -279,22 +282,22 @@ BSONObj extractQueryShape(const FindCommandRequest& findCommand,
     // Hint, max, and min won't serialize if the object is empty.
     if (!findCommand.getHint().isEmpty()) {
         bob.append(FindCommandRequest::kHintFieldName,
-                   redactHintComponent(findCommand.getHint(), opts, false));
+                   extractHintShape(findCommand.getHint(), opts, false));
         // Max/Min aren't valid without hint.
         if (!findCommand.getMax().isEmpty()) {
             bob.append(FindCommandRequest::kMaxFieldName,
-                       redactHintComponent(findCommand.getMax(), opts, true));
+                       extractHintShape(findCommand.getMax(), opts, true));
         }
         if (!findCommand.getMin().isEmpty()) {
             bob.append(FindCommandRequest::kMinFieldName,
-                       redactHintComponent(findCommand.getMin(), opts, true));
+                       extractHintShape(findCommand.getMin(), opts, true));
         }
     }
 
     // Sort.
     if (!findCommand.getSort().isEmpty()) {
         bob.append(FindCommandRequest::kSortFieldName,
-                   query_shape::sortShape(findCommand.getSort(), expCtx, opts));
+                   query_shape::extractSortShape(findCommand.getSort(), expCtx, opts));
     }
 
     // Fields for literal redaction. Adds limit, skip, batchSize, maxTimeMS, and noCursorTimeOut
@@ -305,4 +308,57 @@ BSONObj extractQueryShape(const FindCommandRequest& findCommand,
 
     return bob.obj();
 }
+
+BSONObj extractQueryShape(const AggregateCommandRequest& aggregateCommand,
+                          const std::vector<BSONObj>& serializedPipeline,
+                          const SerializationOptions& opts,
+                          const boost::intrusive_ptr<ExpressionContext>& expCtx) {
+    BSONObjBuilder bob;
+
+    // TODO SERVER-73152 update to newest query shape definition
+
+    // namespace
+    bob.append("ns", extractNamespaceShape(aggregateCommand.getNamespace(), opts));
+    bob.append(AggregateCommandRequest::kCommandName,
+               opts.serializeIdentifier(aggregateCommand.getNamespace().coll()));
+
+    // pipeline
+    {
+        BSONArrayBuilder pipelineBab(
+            bob.subarrayStart(AggregateCommandRequest::kPipelineFieldName));
+        for (const auto& stage : serializedPipeline) {
+            pipelineBab.append(stage);
+        }
+        pipelineBab.doneFast();
+    }
+
+    // explain
+    if (aggregateCommand.getExplain().has_value()) {
+        bob.append(AggregateCommandRequest::kExplainFieldName, true);
+    }
+
+    // allowDiskUse
+    if (auto param = aggregateCommand.getAllowDiskUse(); param.has_value()) {
+        bob.append(AggregateCommandRequest::kAllowDiskUseFieldName, param.value_or(false));
+    }
+
+    // collation
+    if (auto param = aggregateCommand.getCollation()) {
+        bob.append(AggregateCommandRequest::kCollationFieldName, param.get());
+    }
+
+    // hint
+    if (auto hint = aggregateCommand.getHint()) {
+        bob.append(AggregateCommandRequest::kHintFieldName,
+                   extractHintShape(hint.get(), opts, false));
+    }
+
+    // let
+    if (auto letSpec = aggregateCommand.getLet()) {
+        auto redactedObj = extractLetSpecShape(letSpec.get(), opts, expCtx);
+        auto ownedObj = redactedObj.getOwned();
+        bob.append(FindCommandRequest::kLetFieldName, std::move(ownedObj));
+    }
+    return bob.obj();
+}
 }  // namespace mongo::query_shape