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

8 files changed, 89 insertions, 29 deletions

diff --git a/src/mongo/db/commands/find_and_modify.cpp b/src/mongo/db/commands/find_and_modify.cpp
index d7dbad7dade..cdc7965ddfa 100644
--- a/src/mongo/db/commands/find_and_modify.cpp
+++ b/src/mongo/db/commands/find_and_modify.cpp
@@ -703,6 +703,9 @@ write_ops::FindAndModifyCommandReply CmdFindAndModify::Invocation::typedRun(
                 }
                 updateRequest.setSampleId(req.getSampleId());
 
+                updateRequest.setAllowShardKeyUpdatesWithoutFullShardKeyInQuery(
+                    req.getAllowShardKeyUpdatesWithoutFullShardKeyInQuery());
+
                 const ExtensionsCallbackReal extensionsCallback(
                     opCtx, &updateRequest.getNamespaceString());
                 ParsedUpdate parsedUpdate(opCtx, &updateRequest, extensionsCallback);
diff --git a/src/mongo/db/exec/update_stage.cpp b/src/mongo/db/exec/update_stage.cpp
index 692dc6790d4..99c7c28c96b 100644
--- a/src/mongo/db/exec/update_stage.cpp
+++ b/src/mongo/db/exec/update_stage.cpp
@@ -634,19 +634,21 @@ void UpdateStage::_checkRestrictionsOnUpdatingShardKeyAreNotViolated(
     const auto& shardKeyPathsVector = collDesc.getKeyPatternFields();
     pathsupport::EqualityMatches equalities;
 
-    // We can now update a document shard key without needing to specify the full shard key in the
-    // query. The protocol to perform the write uses _id with a proper shardVersion to target the
-    // specific document, but due to current constraints that _id writes are manually broadcast with
-    // ShardVersion::Ignored, we differentiate user queries that originally only contained _id and
-    // the protocol's use of _id in its queries by the shardVersion set.
-    auto sentShardVersion =
-        OperationShardingState::get(opCtx()).getShardVersion(_params.request->getNamespaceString());
+    if (_params.request->getAllowShardKeyUpdatesWithoutFullShardKeyInQuery()) {
+        bool isInternalClient =
+            !cc().session() || (cc().session()->getTags() & transport::Session::kInternalClient);
+        uassert(ErrorCodes::InvalidOptions,
+                "$_allowShardKeyUpdatesWithoutFullShardKeyInQuery is an internal parameter",
+                isInternalClient);
+    }
+
+    // If the incoming update request came from a _clusterWriteWithoutShardKey command, we allow
+    // shard key updates without the full shard key specified.
     bool allowShardKeyUpdatesWithoutFullShardKeyInQuery =
         feature_flags::gFeatureFlagUpdateOneWithoutShardKey.isEnabled(
             serverGlobalParams.featureCompatibility) &&
-        sentShardVersion && !ShardVersion::isIgnoredVersion(*sentShardVersion);
+        _params.request->getAllowShardKeyUpdatesWithoutFullShardKeyInQuery();
 
-    // TODO: SERVER-73689 Fix shard key update check in update_stage.cpp to exclude _id queries.
     uassert(31025,
             "Shard key update is not allowed without specifying the full shard key in the "
             "query",
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) {
diff --git a/src/mongo/db/pipeline/process_interface/common_process_interface.h b/src/mongo/db/pipeline/process_interface/common_process_interface.h
index 4082b152db2..b3b5f26468f 100644
--- a/src/mongo/db/pipeline/process_interface/common_process_interface.h
+++ b/src/mongo/db/pipeline/process_interface/common_process_interface.h
@@ -76,14 +76,16 @@ public:
 
         int estimateUpdateSizeBytes(const BatchObject& batchObject,
                                     UpsertType type) const override {
-            return getUpdateSizeEstimate(std::get<BSONObj>(batchObject),
-                                         std::get<write_ops::UpdateModification>(batchObject),
-                                         std::get<boost::optional<BSONObj>>(batchObject),
-                                         type != UpsertType::kNone /* includeUpsertSupplied */,
-                                         boost::none /* collation */,
-                                         boost::none /* arrayFilters */,
-                                         BSONObj() /* hint*/,
-                                         boost::none /* sampleId */) +
+            return getUpdateSizeEstimate(
+                       std::get<BSONObj>(batchObject),
+                       std::get<write_ops::UpdateModification>(batchObject),
+                       std::get<boost::optional<BSONObj>>(batchObject),
+                       type != UpsertType::kNone /* includeUpsertSupplied */,
+                       boost::none /* collation */,
+                       boost::none /* arrayFilters */,
+                       BSONObj() /* hint*/,
+                       boost::none /* sampleId */,
+                       false /* $_allowShardKeyUpdatesWithoutFullShardKeyInQuery */) +
                 write_ops::kWriteCommandBSONArrayPerElementOverheadBytes;
         }
     };