From 951a2ddf46b4835dff8f407b12f8011a4071c62c Mon Sep 17 00:00:00 2001 From: Jason Zhang Date: Tue, 28 Feb 2023 16:10:37 +0000 Subject: SERVER-73689 Fix shard key update check in update_stage.cpp to exclude _id queries --- src/mongo/db/ops/update_request.h | 21 ++++++++++++++++++++- src/mongo/db/ops/write_ops.cpp | 28 +++++++++++++++++++--------- src/mongo/db/ops/write_ops.h | 3 ++- src/mongo/db/ops/write_ops.idl | 13 ++++++++++++- src/mongo/db/ops/write_ops_exec_test.cpp | 12 ++++++++++++ 5 files changed, 65 insertions(+), 12 deletions(-) (limited to 'src/mongo/db/ops') diff --git a/src/mongo/db/ops/update_request.h b/src/mongo/db/ops/update_request.h index e635ce7f156..bd9c2e04a73 100644 --- a/src/mongo/db/ops/update_request.h +++ b/src/mongo/db/ops/update_request.h @@ -76,7 +76,10 @@ public: }; UpdateRequest(const write_ops::UpdateOpEntry& updateOp = write_ops::UpdateOpEntry()) - : _updateOp(updateOp), _sampleId(updateOp.getSampleId()) {} + : _updateOp(updateOp), + _sampleId(updateOp.getSampleId()), + _allowShardKeyUpdatesWithoutFullShardKeyInQuery( + updateOp.getAllowShardKeyUpdatesWithoutFullShardKeyInQuery()) {} void setNamespaceString(const NamespaceString& nsString) { _nsString = nsString; @@ -265,6 +268,16 @@ public: return _sampleId; } + void setAllowShardKeyUpdatesWithoutFullShardKeyInQuery( + OptionalBool allowShardKeyUpdatesWithoutFullShardKeyInQuery) { + _allowShardKeyUpdatesWithoutFullShardKeyInQuery = + allowShardKeyUpdatesWithoutFullShardKeyInQuery; + } + + const OptionalBool& getAllowShardKeyUpdatesWithoutFullShardKeyInQuery() const { + return _allowShardKeyUpdatesWithoutFullShardKeyInQuery; + } + std::string toString() const { StringBuilder builder; builder << " query: " << getQuery(); @@ -298,6 +311,8 @@ public: builder << " multi: " << isMulti(); builder << " fromOplogApplication: " << _fromOplogApplication; builder << " isExplain: " << static_cast(_explain); + builder << " $_allowShardKeyUpdatesWithoutFullShardKeyInQuery: " + << _allowShardKeyUpdatesWithoutFullShardKeyInQuery; return builder.str(); } @@ -325,6 +340,10 @@ private: // The unique sample id for this request if it has been chosen for sampling. boost::optional _sampleId; + // True if this update is allowed to modify the shard key without the specifying the full shard + // key. + OptionalBool _allowShardKeyUpdatesWithoutFullShardKeyInQuery; + // Flags controlling the update. // God bypasses _id checking and index generation. It is only used on behalf of system diff --git a/src/mongo/db/ops/write_ops.cpp b/src/mongo/db/ops/write_ops.cpp index b80d1fc43e0..c57f65ccdf8 100644 --- a/src/mongo/db/ops/write_ops.cpp +++ b/src/mongo/db/ops/write_ops.cpp @@ -152,7 +152,8 @@ int getUpdateSizeEstimate(const BSONObj& q, const boost::optional& collation, const boost::optional>& arrayFilters, const mongo::BSONObj& hint, - const boost::optional& sampleId) { + const boost::optional& sampleId, + const bool includeAllowShardKeyUpdatesWithoutFullShardKeyInQuery) { using UpdateOpEntry = write_ops::UpdateOpEntry; // This constant accounts for the null terminator in each field name and the BSONType byte for @@ -210,6 +211,12 @@ int getUpdateSizeEstimate(const BSONObj& q, estSize += UpdateOpEntry::kSampleIdFieldName.size() + kUUIDSize + kPerElementOverhead; } + // Add the size of the '$_allowShardKeyUpdatesWithoutFullShardKeyInQuery' field, if present. + if (includeAllowShardKeyUpdatesWithoutFullShardKeyInQuery) { + estSize += UpdateOpEntry::kAllowShardKeyUpdatesWithoutFullShardKeyInQueryFieldName.size() + + kBoolSize + kPerElementOverhead; + } + return estSize; } @@ -248,14 +255,17 @@ int getDeleteSizeEstimate(const BSONObj& q, } bool verifySizeEstimate(const write_ops::UpdateOpEntry& update) { - return write_ops::getUpdateSizeEstimate(update.getQ(), - update.getU(), - update.getC(), - update.getUpsertSupplied().has_value(), - update.getCollation(), - update.getArrayFilters(), - update.getHint(), - update.getSampleId()) >= update.toBSON().objsize(); + return write_ops::getUpdateSizeEstimate( + update.getQ(), + update.getU(), + update.getC(), + update.getUpsertSupplied().has_value(), + update.getCollation(), + update.getArrayFilters(), + update.getHint(), + update.getSampleId(), + update.getAllowShardKeyUpdatesWithoutFullShardKeyInQuery().has_value()) >= + update.toBSON().objsize(); } bool verifySizeEstimate(const write_ops::DeleteOpEntry& deleteOp) { diff --git a/src/mongo/db/ops/write_ops.h b/src/mongo/db/ops/write_ops.h index e2e546cf3df..2545c8eb78e 100644 --- a/src/mongo/db/ops/write_ops.h +++ b/src/mongo/db/ops/write_ops.h @@ -114,7 +114,8 @@ int getUpdateSizeEstimate(const BSONObj& q, const boost::optional& collation, const boost::optional>& arrayFilters, const mongo::BSONObj& hint, - const boost::optional& sampleId); + const boost::optional& sampleId, + bool includeAllowShardKeyUpdatesWithoutFullShardKeyInQuery); int getDeleteSizeEstimate(const BSONObj& q, const boost::optional& collation, const mongo::BSONObj& hint, diff --git a/src/mongo/db/ops/write_ops.idl b/src/mongo/db/ops/write_ops.idl index 52f3dd1a9e9..a2d80295d03 100644 --- a/src/mongo/db/ops/write_ops.idl +++ b/src/mongo/db/ops/write_ops.idl @@ -224,7 +224,6 @@ structs: cpp_name: originalCollation stability: internal - UpdateOpEntry: description: "Parser for the entries in the 'updates' array of an update command." strict: true @@ -283,6 +282,12 @@ structs: type: uuid optional: true stability: unstable + $_allowShardKeyUpdatesWithoutFullShardKeyInQuery: + description: "Set to true if shard key updates are allowed without the full shard + key in the query." + type: optionalBool + cpp_name: allowShardKeyUpdatesWithoutFullShardKeyInQuery + stability: internal DeleteOpEntry: description: "Parser for the entries in the 'deletes' array of a delete command." @@ -587,4 +592,10 @@ commands: optional: true cpp_name: originalCollation stability: internal + $_allowShardKeyUpdatesWithoutFullShardKeyInQuery: + description: "Set to true if shard key updates are allowed without the full shard + key in the query." + type: optionalBool + cpp_name: allowShardKeyUpdatesWithoutFullShardKeyInQuery + stability: internal diff --git a/src/mongo/db/ops/write_ops_exec_test.cpp b/src/mongo/db/ops/write_ops_exec_test.cpp index c5f646db73a..f977d7878dc 100644 --- a/src/mongo/db/ops/write_ops_exec_test.cpp +++ b/src/mongo/db/ops/write_ops_exec_test.cpp @@ -87,6 +87,18 @@ TEST_F(WriteOpsExecTest, TestUpdateSizeEstimationLogic) { // Add a sampleId. updateOpEntry.setSampleId(UUID::gen()); ASSERT(write_ops::verifySizeEstimate(updateOpEntry)); + + // Add $_allowShardKeyUpdatesWithoutFullShardKeyInQuery. + updateOpEntry.setAllowShardKeyUpdatesWithoutFullShardKeyInQuery(OptionalBool(false)); + ASSERT(write_ops::verifySizeEstimate(updateOpEntry)); + + // Set '$_allowShardKeyUpdatesWithoutFullShardKeyInQuery' to true. + updateOpEntry.setAllowShardKeyUpdatesWithoutFullShardKeyInQuery(OptionalBool(true)); + ASSERT(write_ops::verifySizeEstimate(updateOpEntry)); + + // Set '$_allowShardKeyUpdatesWithoutFullShardKeyInQuery' to boost::none. + updateOpEntry.setAllowShardKeyUpdatesWithoutFullShardKeyInQuery(OptionalBool(boost::none)); + ASSERT(write_ops::verifySizeEstimate(updateOpEntry)); } TEST_F(WriteOpsExecTest, TestDeleteSizeEstimationLogic) { -- cgit v1.2.1