New upstream version 6.0.20upstream

1 files changed, 292 insertions, 0 deletions

diff --git a/src/mongo/db/query/query_shape.cpp b/src/mongo/db/query/query_shape.cpp
new file mode 100644
index 00000000000..02d4c97d25f
--- /dev/null
+++ b/src/mongo/db/query/query_shape.cpp
@@ -0,0 +1,292 @@
+/**
+ *    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.h"
+
+#include "mongo/base/status.h"
+#include "mongo/db/query/find_command_gen.h"
+#include "mongo/db/query/projection_ast_util.h"
+#include "mongo/db/query/query_request_helper.h"
+#include "mongo/db/query/query_shape_gen.h"
+#include "mongo/db/query/query_stats/query_stats.h"
+#include "mongo/db/query/shape_helpers.h"
+#include "mongo/db/query/sort_pattern.h"
+
+namespace mongo::query_shape {
+
+BSONObj debugPredicateShape(const MatchExpression* predicate) {
+    SerializationOptions opts;
+    opts.literalPolicy = LiteralSerializationPolicy::kToDebugTypeString;
+    return predicate->serialize(opts);
+}
+BSONObj representativePredicateShape(const MatchExpression* predicate) {
+    SerializationOptions opts;
+    opts.literalPolicy = LiteralSerializationPolicy::kToRepresentativeParseableValue;
+    return predicate->serialize(opts);
+}
+
+BSONObj debugPredicateShape(const MatchExpression* predicate,
+                            std::function<std::string(StringData)> transformIdentifiersCallback) {
+    SerializationOptions opts;
+    opts.literalPolicy = LiteralSerializationPolicy::kToDebugTypeString;
+    opts.transformIdentifiersCallback = transformIdentifiersCallback;
+    opts.transformIdentifiers = true;
+    return predicate->serialize(opts);
+}
+
+BSONObj representativePredicateShape(
+    const MatchExpression* predicate,
+    std::function<std::string(StringData)> transformIdentifiersCallback) {
+    SerializationOptions opts;
+    opts.literalPolicy = LiteralSerializationPolicy::kToRepresentativeParseableValue;
+    opts.transformIdentifiersCallback = transformIdentifiersCallback;
+    opts.transformIdentifiers = true;
+    return predicate->serialize(opts);
+}
+
+BSONObj extractSortShape(const BSONObj& sortSpec,
+                         const boost::intrusive_ptr<ExpressionContext>& expCtx,
+                         const SerializationOptions& opts) {
+    if (sortSpec.isEmpty()) {
+        return sortSpec;
+    }
+    auto natural = sortSpec[query_request_helper::kNaturalSortField];
+
+    if (!natural) {
+        return SortPattern{sortSpec, expCtx}
+            .serialize(SortPattern::SortKeySerialization::kForPipelineSerialization, opts)
+            .toBson();
+    }
+    // This '$natural' will fail to parse as a valid SortPattern since it is not a valid field
+    // path - it is usually considered and converted into a hint. For the query shape, we'll
+    // keep it unmodified.
+    BSONObjBuilder bob;
+    for (auto&& elem : sortSpec) {
+        if (elem.isABSONObj()) {
+            // We expect this won't work or parse on the main command path, but for shapification we
+            // don't really care, just treat it as a literal and don't bother parsing.
+            opts.appendLiteral(
+                &bob, opts.serializeFieldPathFromString(elem.fieldNameStringData()), elem);
+        } else if (elem.fieldNameStringData() == natural.fieldNameStringData()) {
+            bob.append(elem);
+        } else {
+            bob.appendAs(elem, opts.serializeFieldPathFromString(elem.fieldNameStringData()));
+        }
+    }
+    return bob.obj();
+}
+
+void addShapeLiterals(BSONObjBuilder* bob,
+                      const FindCommandRequest& findCommand,
+                      const SerializationOptions& opts) {
+    if (auto limit = findCommand.getLimit()) {
+        opts.appendLiteral(
+            bob, FindCommandRequest::kLimitFieldName, static_cast<long long>(*limit));
+    }
+    if (auto skip = findCommand.getSkip()) {
+        opts.appendLiteral(bob, FindCommandRequest::kSkipFieldName, static_cast<long long>(*skip));
+    }
+}
+
+static std::vector<
+    std::pair<StringData, std::function<const OptionalBool(const FindCommandRequest&)>>>
+    boolArgMap = {
+        {FindCommandRequest::kSingleBatchFieldName, &FindCommandRequest::getSingleBatch},
+        {FindCommandRequest::kAllowDiskUseFieldName, &FindCommandRequest::getAllowDiskUse},
+        {FindCommandRequest::kReturnKeyFieldName, &FindCommandRequest::getReturnKey},
+        {FindCommandRequest::kShowRecordIdFieldName, &FindCommandRequest::getShowRecordId},
+        {FindCommandRequest::kTailableFieldName, &FindCommandRequest::getTailable},
+        {FindCommandRequest::kAwaitDataFieldName, &FindCommandRequest::getAwaitData},
+};
+std::vector<std::pair<StringData, std::function<const BSONObj(const FindCommandRequest&)>>>
+    objArgMap = {
+        {FindCommandRequest::kCollationFieldName, &FindCommandRequest::getCollation},
+
+};
+
+void addRemainingFindCommandFields(BSONObjBuilder* bob,
+                                   const FindCommandRequest& findCommand,
+                                   const SerializationOptions& opts) {
+    for (auto [fieldName, getterFunction] : boolArgMap) {
+        auto optBool = getterFunction(findCommand);
+        optBool.serializeToBSON(fieldName, bob);
+    }
+
+    if (auto optOplogReplay = findCommand.getOplogReplay()) {
+        if (optOplogReplay.has_value()) {
+            opts.appendLiteral(
+                bob, FindCommandRequest::kOplogReplayFieldName, optOplogReplay.value_or(false));
+        }
+    }
+
+    auto collation = findCommand.getCollation();
+    if (!collation.isEmpty()) {
+        bob->append(FindCommandRequest::kCollationFieldName, collation);
+    }
+}
+
+
+/**
+ * In a let specification all field names are variable names, and all values are either
+ * expressions or constants.
+ */
+BSONObj extractLetSpecShape(BSONObj letSpec,
+                            const SerializationOptions& opts,
+                            const boost::intrusive_ptr<ExpressionContext>& expCtx) {
+
+    BSONObjBuilder bob;
+    for (BSONElement elem : letSpec) {
+        auto expr = Expression::parseOperand(expCtx.get(), elem, expCtx->variablesParseState);
+        auto redactedValue = expr->serialize(opts);
+        // Note that this will throw on deeply nested let variables.
+        redactedValue.addToBsonObj(&bob, opts.serializeFieldPathFromString(elem.fieldName()));
+    }
+    return bob.obj();
+}
+
+void appendCmdNs(BSONObjBuilder& bob,
+                 const NamespaceString& nss,
+                 const SerializationOptions& opts) {
+    BSONObjBuilder nsObj = bob.subobjStart("cmdNs");
+    shape_helpers::appendNamespaceShape(nsObj, nss, opts);
+    nsObj.doneFast();
+}
+
+BSONObj extractQueryShape(const ParsedFindCommand& findRequest,
+                          const SerializationOptions& opts,
+                          const boost::intrusive_ptr<ExpressionContext>& expCtx) {
+    const auto& findCmd = *findRequest.findCommandRequest;
+    BSONObjBuilder bob;
+    // Serialize the namespace as part of the query shape.
+    {
+        auto ns = findCmd.getNamespaceOrUUID();
+        if (ns.nss().has_value()) {
+            appendCmdNs(bob, *ns.nss(), opts);
+        } else {
+            BSONObjBuilder cmdNs = bob.subobjStart("cmdNs");
+            cmdNs.append("uuid", opts.serializeIdentifier(ns.uuid()->toString()));
+            cmdNs.append("db", opts.serializeIdentifier(ns.db()));
+            cmdNs.doneFast();
+        }
+    }
+
+    bob.append("command", "find");
+    std::unique_ptr<MatchExpression> filterExpr;
+    // Filter.
+    bob.append(FindCommandRequest::kFilterFieldName, findRequest.filter->serialize(opts));
+    // Let Spec.
+    if (auto letSpec = findCmd.getLet()) {
+        auto redactedObj = extractLetSpecShape(letSpec.get(), opts, expCtx);
+        auto ownedObj = redactedObj.getOwned();
+        bob.append(FindCommandRequest::kLetFieldName, std::move(ownedObj));
+    }
+
+    if (findRequest.proj) {
+        bob.append(FindCommandRequest::kProjectionFieldName,
+                   projection_ast::serialize(*findRequest.proj->root(), opts));
+    }
+
+    if (!findCmd.getMax().isEmpty()) {
+        bob.append(FindCommandRequest::kMaxFieldName,
+                   shape_helpers::extractMinOrMaxShape(findCmd.getMax(), opts));
+    }
+    if (!findCmd.getMin().isEmpty()) {
+        bob.append(FindCommandRequest::kMinFieldName,
+                   shape_helpers::extractMinOrMaxShape(findCmd.getMin(), opts));
+    }
+
+    // Sort.
+    if (findRequest.sort) {
+        bob.append(
+            FindCommandRequest::kSortFieldName,
+            findRequest.sort
+                ->serialize(SortPattern::SortKeySerialization::kForPipelineSerialization, opts)
+                .toBson());
+    }
+
+    // Fields for literal redaction. Adds limit and skip.
+    addShapeLiterals(&bob, findCmd, opts);
+
+    // Add the fields that require no redaction.
+    addRemainingFindCommandFields(&bob, findCmd, opts);
+
+    return bob.obj();
+}
+
+BSONObj extractQueryShape(const AggregateCommandRequest& aggregateCommand,
+                          const Pipeline& pipeline,
+                          const SerializationOptions& opts,
+                          const boost::intrusive_ptr<ExpressionContext>& expCtx,
+                          const NamespaceString& nss) {
+    BSONObjBuilder bob;
+
+    // namespace
+    appendCmdNs(bob, nss, opts);
+    bob.append("command", "aggregate");
+
+    // pipeline
+    {
+        BSONArrayBuilder pipelineBab(
+            bob.subarrayStart(AggregateCommandRequest::kPipelineFieldName));
+        auto serializedPipeline = pipeline.serializeToBson(opts);
+        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());
+    }
+
+    // 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();
+}
+
+QueryShapeHash hash(const BSONObj& queryShape) {
+    return QueryShapeHash::computeHash(reinterpret_cast<const uint8_t*>(queryShape.objdata()),
+                                       queryShape.objsize());
+}
+}  // namespace mongo::query_shape