New upstream version 6.0.20upstream

1 files changed, 125 insertions, 0 deletions

diff --git a/src/mongo/db/query/query_shape/agg_cmd_shape.cpp b/src/mongo/db/query/query_shape/agg_cmd_shape.cpp
new file mode 100644
index 00000000000..e997150ecc6
--- /dev/null
+++ b/src/mongo/db/query/query_shape/agg_cmd_shape.cpp
@@ -0,0 +1,125 @@
+/**
+ *    Copyright (C) 2023-present MongoDB, Inc.
+ *
+ *    This program is free software: you can redistribute it and/or modify
+ *    it under the terms of the Server Side Public License, version 1,
+ *    as published by MongoDB, Inc.
+ *
+ *    This program is distributed in the hope that it will be useful,
+ *    but WITHOUT ANY WARRANTY; without even the implied warranty of
+ *    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ *    Server Side Public License for more details.
+ *
+ *    You should have received a copy of the Server Side Public License
+ *    along with this program. If not, see
+ *    <http://www.mongodb.com/licensing/server-side-public-license>.
+ *
+ *    As a special exception, the copyright holders give permission to link the
+ *    code of portions of this program with the OpenSSL library under certain
+ *    conditions as described in each individual source file and distribute
+ *    linked combinations including the program with the OpenSSL library. You
+ *    must comply with the Server Side Public License in all respects for
+ *    all of the code used other than as permitted herein. If you modify file(s)
+ *    with this exception, you may extend this exception to your version of the
+ *    file(s), but you are not obligated to do so. If you do not wish to do so,
+ *    delete this exception statement from your version. If you delete this
+ *    exception statement from all source files in the program, then also delete
+ *    it in the license file.
+ */
+
+#include "mongo/db/query/query_shape/agg_cmd_shape.h"
+
+#include "mongo/db/query/query_shape/shape_helpers.h"
+
+namespace mongo::query_shape {
+
+AggCmdShapeComponents::AggCmdShapeComponents(
+    const AggregateCommandRequest& aggRequest,
+    stdx::unordered_set<NamespaceString> involvedNamespaces_,
+    std::vector<BSONObj> pipeline)
+    : allowDiskUse(aggRequest.getAllowDiskUse()),
+      involvedNamespaces(std::move(involvedNamespaces_)),
+      representativePipeline(std::move(pipeline)) {}
+
+AggCmdShapeComponents::AggCmdShapeComponents(
+    OptionalBool allowDiskUse,
+    stdx::unordered_set<NamespaceString> involvedNamespaces_,
+    std::vector<BSONObj> pipeline)
+    : allowDiskUse(allowDiskUse),
+      involvedNamespaces(std::move(involvedNamespaces_)),
+      representativePipeline(std::move(pipeline)) {}
+
+void AggCmdShapeComponents::HashValue(absl::HashState state) const {
+    state = absl::HashState::combine(std::move(state), allowDiskUse);
+    for (auto&& shapifiedStage : representativePipeline) {
+        state = absl::HashState::combine(std::move(state), simpleHash(shapifiedStage));
+    }
+}
+
+void AggCmdShape::appendLetCmdSpecificShapeComponents(
+    BSONObjBuilder& bob,
+    const boost::intrusive_ptr<ExpressionContext>& expCtx,
+    const SerializationOptions& opts) const {
+    tassert(7633000,
+            "We don't support serializing to the unmodified shape here, since we have already "
+            "shapified and stored the representative query - we've lost the original literals",
+            opts.literalPolicy != LiteralSerializationPolicy::kUnchanged);
+
+    if (opts == SerializationOptions::kRepresentativeQueryShapeSerializeOptions) {
+        // We have this copy stored already!
+        return _components.appendTo(bob);
+    } else {
+        // The cached pipeline shape doesn't match the requested options, so we have to
+        // re-parse the pipeline from the initial request.
+        expCtx->inMongos = _inMongos;
+        expCtx->addResolvedNamespaces(_components.involvedNamespaces);
+        auto reparsed = Pipeline::parse(_components.representativePipeline, expCtx);
+        auto serializedPipeline = reparsed->serializeToBson(opts);
+        AggCmdShapeComponents{
+            _components.allowDiskUse, _components.involvedNamespaces, serializedPipeline}
+            .appendTo(bob);
+    }
+}
+
+void AggCmdShapeComponents::appendTo(BSONObjBuilder& bob) const {
+    bob.append("command", "aggregate");
+
+    // pipeline
+    bob.append(AggregateCommandRequest::kPipelineFieldName, representativePipeline);
+
+    // allowDiskUse
+    if (allowDiskUse.has_value()) {
+        bob.append(AggregateCommandRequest::kAllowDiskUseFieldName, bool(allowDiskUse));
+    }
+}
+
+// As part of the size, we must track the allocation of elements in the representative
+// pipeline, as well as the elements in the unordered set of involved namespaces.
+size_t AggCmdShapeComponents::size() const {
+    return sizeof(AggCmdShapeComponents) + shape_helpers::containerSize(representativePipeline) +
+        shape_helpers::containerSize(involvedNamespaces);
+}
+
+AggCmdShape::AggCmdShape(const AggregateCommandRequest& aggregateCommand,
+                         NamespaceString origNss,
+                         stdx::unordered_set<NamespaceString> involvedNamespaces_,
+                         const Pipeline& pipeline,
+                         const boost::intrusive_ptr<ExpressionContext>& expCtx)
+    : CmdWithLetShape(aggregateCommand.getLet(),
+                      expCtx,
+                      _components,
+                      std::move(origNss),
+                      aggregateCommand.getCollation().value_or(BSONObj())),
+      _components(aggregateCommand,
+                  std::move(involvedNamespaces_),
+                  pipeline.serializeToBson(
+                      SerializationOptions::kRepresentativeQueryShapeSerializeOptions)),
+      _inMongos(expCtx->inMongos) {}
+
+size_t AggCmdShape::extraSize() const {
+    // To account for possible padding, we calculate the extra space with the difference instead of
+    // using sizeof(bool);
+    return sizeof(AggCmdShape) - sizeof(CmdWithLetShape) - sizeof(AggCmdShapeComponents);
+}
+
+}  // namespace mongo::query_shape