diff options
Diffstat (limited to 'src/mongo/db/query/query_shape.cpp')
| -rw-r--r-- | src/mongo/db/query/query_shape.cpp | 292 |
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 |