SERVER-73689 Fix shard key update check in update_stage.cpp to exclude _id queries

5 files changed, 65 insertions, 12 deletions

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<bool>(_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<UUID> _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<mongo::BSONObj>& collation,
                           const boost::optional<std::vector<mongo::BSONObj>>& arrayFilters,
                           const mongo::BSONObj& hint,
-                          const boost::optional<UUID>& sampleId) {
+                          const boost::optional<UUID>& 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<mongo::BSONObj>& collation,
                           const boost::optional<std::vector<mongo::BSONObj>>& arrayFilters,
                           const mongo::BSONObj& hint,
-                          const boost::optional<UUID>& sampleId);
+                          const boost::optional<UUID>& sampleId,
+                          bool includeAllowShardKeyUpdatesWithoutFullShardKeyInQuery);
 int getDeleteSizeEstimate(const BSONObj& q,
                           const boost::optional<mongo::BSONObj>& 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) {