SERVER-73045 Preserve the original filter in the write phase for updateOne/deleteOne/findAndModify without shard key

6 files changed, 181 insertions, 74 deletions

diff --git a/jstests/sharding/analyze_shard_key/read_and_write_distribution.js b/jstests/sharding/analyze_shard_key/read_and_write_distribution.js
index 892cccac739..75254d54d77 100644
--- a/jstests/sharding/analyze_shard_key/read_and_write_distribution.js
+++ b/jstests/sharding/analyze_shard_key/read_and_write_distribution.js
@@ -3,7 +3,7 @@
  * distribution metrics, but on replica sets it does not since query sampling is only supported on
  * sharded clusters at this point.
  *
- * @tags: [requires_fcv_63, featureFlagAnalyzeShardKey]
+ * @tags: [requires_fcv_63, featureFlagAnalyzeShardKey, featureFlagUpdateOneWithoutShardKey]
  */
 (function() {
 "use strict";
@@ -283,67 +283,63 @@ function makeTestCase(collName, isShardedColl, {shardKeyField, isHashed, minVal,
         writeDistribution.numShardKeyUpdates++;
     }
 
-    // TODO (SERVER-73045): Remove the if below to add test coverage for sampling of single write
-    // without shard key.
-    if (!isShardedColl) {
-        // Below are writes targeting a variable number of shards.
-
-        for (let i = 0; i < getRandomCount(); i++) {
-            cmdObjs.push({
-                update: collName,
-                updates: [{q: {[shardKeyField]: {$gte: getNextVal()}}, u: {$set: {z: 0}}}]
-            });
-            writeDistribution.sampleSize.update++;
-            writeDistribution.sampleSize.total++;
-            if (isHashed) {
-                // For hashed sharding, range queries on the shard key target all shards.
-                writeDistribution.numScatterGather++;
-            } else {
-                writeDistribution.numVariableShard++;
-            }
-            writeDistribution.numSingleWritesWithoutShardKey++;
-        }
+    // Below are writes targeting a variable number of shards.
 
-        for (let i = 0; i < getRandomCount(); i++) {
-            cmdObjs.push(
-                {delete: collName, deletes: [{q: {[shardKeyField]: {$lte: minVal++}}, limit: 0}]});
-            writeDistribution.sampleSize.delete ++;
-            writeDistribution.sampleSize.total++;
-            if (isHashed) {
-                // For hashed sharding, range queries on the shard key target all shards.
-                writeDistribution.numScatterGather++;
-            } else {
-                writeDistribution.numVariableShard++;
-            }
-            writeDistribution.numMultiWritesWithoutShardKey++;
+    for (let i = 0; i < getRandomCount(); i++) {
+        cmdObjs.push({
+            update: collName,
+            updates: [{q: {[shardKeyField]: {$gte: getNextVal()}}, u: {$set: {z: 0}}}]
+        });
+        writeDistribution.sampleSize.update++;
+        writeDistribution.sampleSize.total++;
+        if (isHashed) {
+            // For hashed sharding, range queries on the shard key target all shards.
+            writeDistribution.numScatterGather++;
+        } else {
+            writeDistribution.numVariableShard++;
         }
+        writeDistribution.numSingleWritesWithoutShardKey++;
+    }
 
-        for (let i = 0; i < getRandomCount(); i++) {
-            cmdObjs.push({
-                findAndModify: collName,
-                query: {[shardKeyField]: {$lte: getNextVal()}},
-                update: {$set: {z: 0}}
-            });
-            writeDistribution.sampleSize.findAndModify++;
-            writeDistribution.sampleSize.total++;
-            if (isHashed) {
-                // For hashed sharding, range queries on the shard key target all shards.
-                writeDistribution.numScatterGather++;
-            } else {
-                writeDistribution.numVariableShard++;
-            }
-            writeDistribution.numSingleWritesWithoutShardKey++;
+    for (let i = 0; i < getRandomCount(); i++) {
+        cmdObjs.push(
+            {delete: collName, deletes: [{q: {[shardKeyField]: {$lte: minVal++}}, limit: 0}]});
+        writeDistribution.sampleSize.delete ++;
+        writeDistribution.sampleSize.total++;
+        if (isHashed) {
+            // For hashed sharding, range queries on the shard key target all shards.
+            writeDistribution.numScatterGather++;
+        } else {
+            writeDistribution.numVariableShard++;
         }
+        writeDistribution.numMultiWritesWithoutShardKey++;
+    }
 
-        // Below are writes targeting all shards.
-
-        for (let i = 0; i < getRandomCount(); i++) {
-            cmdObjs.push({findAndModify: collName, query: {}, update: {$set: {z: 0}}});
-            writeDistribution.sampleSize.findAndModify++;
-            writeDistribution.sampleSize.total++;
+    for (let i = 0; i < getRandomCount(); i++) {
+        cmdObjs.push({
+            findAndModify: collName,
+            query: {[shardKeyField]: {$lte: getNextVal()}},
+            update: {$set: {z: 0}}
+        });
+        writeDistribution.sampleSize.findAndModify++;
+        writeDistribution.sampleSize.total++;
+        if (isHashed) {
+            // For hashed sharding, range queries on the shard key target all shards.
             writeDistribution.numScatterGather++;
-            writeDistribution.numSingleWritesWithoutShardKey++;
+        } else {
+            writeDistribution.numVariableShard++;
         }
+        writeDistribution.numSingleWritesWithoutShardKey++;
+    }
+
+    // Below are writes targeting all shards.
+
+    for (let i = 0; i < getRandomCount(); i++) {
+        cmdObjs.push({findAndModify: collName, query: {}, update: {$set: {z: 0}}});
+        writeDistribution.sampleSize.findAndModify++;
+        writeDistribution.sampleSize.total++;
+        writeDistribution.numScatterGather++;
+        writeDistribution.numSingleWritesWithoutShardKey++;
     }
 
     return {cmdObjs, metrics: {readDistribution, writeDistribution}};
@@ -516,13 +512,12 @@ const analyzeShardKeyNumRanges = 10;
 
     runTest({isShardedColl: false, shardKeyField: "x", isHashed: false});
     runTest({isShardedColl: false, shardKeyField: "x", isHashed: true});
+
     // Note that {x: 1} is the current shard key for the sharded collection being tested.
     runTest({isShardedColl: true, shardKeyField: "x", isHashed: false});
     runTest({isShardedColl: true, shardKeyField: "x", isHashed: true});
-    // TODO (SERVER-73045): Uncomment the tests below to add test coverage for sampling of single
-    // writes without shard key.
-    // runTest({isShardedColl: true, shardKeyField: "y", isHashed: false});
-    // runTest({isShardedColl: true, shardKeyField: "y", isHashed: true});
+    runTest({isShardedColl: true, shardKeyField: "y", isHashed: false});
+    runTest({isShardedColl: true, shardKeyField: "y", isHashed: true});
 
     st.stop();
 }
diff --git a/jstests/sharding/updateOne_without_shard_key/cluster_write_without_shard_key_basic.js b/jstests/sharding/updateOne_without_shard_key/cluster_write_without_shard_key_basic.js
index aa9c54b2d6d..4f623a3b3b1 100644
--- a/jstests/sharding/updateOne_without_shard_key/cluster_write_without_shard_key_basic.js
+++ b/jstests/sharding/updateOne_without_shard_key/cluster_write_without_shard_key_basic.js
@@ -94,9 +94,12 @@ function runAndVerifyCommand(testCase) {
                 writeCmd: {
                     update: collName,
                     updates: [
-                        {q: {}, u: {$set: {a: aFieldValue}}},
+                        {
+                            q: {},
+                            u: {$set: {a: aFieldValue}},
+                            collation: {locale: "simple"},
+                        },
                     ],
-                    collation: {locale: "simple"},
                     writeConcern: {w: "majority"},
                 },
                 shardId: shardConn,
diff --git a/src/mongo/db/commands/find_and_modify.cpp b/src/mongo/db/commands/find_and_modify.cpp
index bf0189073a5..2d0082285eb 100644
--- a/src/mongo/db/commands/find_and_modify.cpp
+++ b/src/mongo/db/commands/find_and_modify.cpp
@@ -684,8 +684,18 @@ write_ops::FindAndModifyCommandReply CmdFindAndModify::Invocation::typedRun(
     }
 
     if (analyze_shard_key::supportsPersistingSampledQueries() && req.getSampleId()) {
-        analyze_shard_key::QueryAnalysisWriter::get(opCtx)->addFindAndModifyQuery(req).getAsync(
-            [](auto) {});
+        auto findAndModifyOp = req;
+
+        // If the initial query was a write without shard key, the two phase write protocol
+        // modifies the query in the write phase. In order to get correct metrics, we need to
+        // reconstruct the original query prior to sampling.
+        if (req.getOriginalQuery()) {
+            findAndModifyOp.setQuery(*req.getOriginalQuery());
+            findAndModifyOp.setCollation(req.getOriginalCollation());
+        }
+        analyze_shard_key::QueryAnalysisWriter::get(opCtx)
+            ->addFindAndModifyQuery(findAndModifyOp)
+            .getAsync([](auto) {});
     }
 
     if (MONGO_unlikely(failAllFindAndModify.shouldFail())) {
diff --git a/src/mongo/db/ops/write_ops.idl b/src/mongo/db/ops/write_ops.idl
index 5f793cb9780..4fe3d0d7650 100644
--- a/src/mongo/db/ops/write_ops.idl
+++ b/src/mongo/db/ops/write_ops.idl
@@ -205,6 +205,25 @@ structs:
                 type: EncryptionInformation
                 optional: true
                 stability: unstable
+            $_originalQuery:
+                description: "The original write query. This is used for updateOne/deleteOne 
+                              without shard key during the write phase of the two phase protocol in 
+                              order to make sure the shard key query analysis stores the correct 
+                              client query."
+                type: object
+                optional: true
+                cpp_name: originalQuery
+                stability: internal
+            $_originalCollation:
+                description: "The original write query. This is used for updateOne/deleteOne 
+                              without shard key during the write phase of the two phase protocol in 
+                              order to make sure the shard key query analysis stores the correct 
+                              client collation."
+                type: object
+                optional: true
+                cpp_name: originalCollation
+                stability: internal
+
 
     UpdateOpEntry:
         description: "Parser for the entries in the 'updates' array of an update command."
@@ -550,3 +569,22 @@ commands:
                 type: uuid
                 optional: true
                 stability: unstable
+            $_originalQuery:
+                description: "The original write query. This is used for findAndModify without shard 
+                              key during the write phase of the two phase protocol in order to make 
+                              sure the shard key query analysis stores the correct client 
+                              query."
+                type: object
+                optional: true
+                cpp_name: originalQuery
+                stability: internal
+            $_originalCollation:
+                description: "The original collation. This is used for findAndModify without shard 
+                              key during the write phase of the two phase protocol in order to make 
+                              sure the shard key query analysis stores the correct client 
+                              collation."
+                type: object
+                optional: true
+                cpp_name: originalCollation
+                stability: internal
+            
diff --git a/src/mongo/db/ops/write_ops_exec.cpp b/src/mongo/db/ops/write_ops_exec.cpp
index dcb72a4ca93..71d6531de1a 100644
--- a/src/mongo/db/ops/write_ops_exec.cpp
+++ b/src/mongo/db/ops/write_ops_exec.cpp
@@ -1106,8 +1106,18 @@ WriteResult performUpdates(OperationContext* opCtx,
         });
 
         if (analyze_shard_key::supportsPersistingSampledQueries() && singleOp.getSampleId()) {
+            auto updateOp = wholeOp;
+
+            // If the initial query was a write without shard key, the two phase write protocol
+            // modifies the query in the write phase. In order to get correct metrics, we need to
+            // reconstruct the original query prior to sampling.
+            if (wholeOp.getOriginalQuery()) {
+                updateOp.getUpdates().front().setQ(*wholeOp.getOriginalQuery());
+                updateOp.getUpdates().front().setCollation(wholeOp.getOriginalCollation());
+            }
+
             analyze_shard_key::QueryAnalysisWriter::get(opCtx)
-                ->addUpdateQuery(wholeOp, currentOpIndex)
+                ->addUpdateQuery(updateOp, currentOpIndex)
                 .getAsync([](auto) {});
         }
 
@@ -1354,8 +1364,18 @@ WriteResult performDeletes(OperationContext* opCtx,
         });
 
         if (analyze_shard_key::supportsPersistingSampledQueries() && singleOp.getSampleId()) {
+            auto deleteOp = wholeOp;
+
+            // If the initial query was a write without shard key, the two phase write protocol
+            // modifies the query in the write phase. In order to get correct metrics, we need to
+            // reconstruct the original query prior to sampling.
+            if (wholeOp.getOriginalQuery()) {
+                deleteOp.getDeletes().front().setQ(*wholeOp.getOriginalQuery());
+                deleteOp.getDeletes().front().setCollation(wholeOp.getOriginalCollation());
+            }
+
             analyze_shard_key::QueryAnalysisWriter::get(opCtx)
-                ->addDeleteQuery(wholeOp, currentOpIndex)
+                ->addDeleteQuery(deleteOp, currentOpIndex)
                 .getAsync([](auto) {});
         }
 
@@ -1389,7 +1409,7 @@ WriteResult performDeletes(OperationContext* opCtx,
     }
 
     return out;
-}
+}  // namespace mongo::write_ops_exec
 
 Status performAtomicTimeseriesWrites(
     OperationContext* opCtx,
diff --git a/src/mongo/s/commands/cluster_write_without_shard_key_cmd.cpp b/src/mongo/s/commands/cluster_write_without_shard_key_cmd.cpp
index 3250f90b0af..7d9a470dc4a 100644
--- a/src/mongo/s/commands/cluster_write_without_shard_key_cmd.cpp
+++ b/src/mongo/s/commands/cluster_write_without_shard_key_cmd.cpp
@@ -49,30 +49,71 @@ BSONObj _createCmdObj(const BSONObj& writeCmd,
                       const StringData& commandName,
                       const BSONObj& targetDocId,
                       const NamespaceString& nss) {
-    // Drop collation and writeConcern as
-    // targeting by _id uses default collation and writeConcern cannot be specified for
-    // commands run in internal transactions. This object will be used to construct the command
-    // request used by clusterWriteWithoutShardKey.
+
+    // Drop the writeConcern as it cannot be specified for commands run in internal transactions.
+    // This object will be used to construct the command request used by
+    // _clusterWriteWithoutShardKey.
     BSONObjBuilder writeCmdObjBuilder(
-        writeCmd.removeFields(std::set<std::string>{"collation", "writeConcern"}));
+        writeCmd.removeField(WriteConcernOptions::kWriteConcernField));
     writeCmdObjBuilder.appendElementsUnique(BSON("$db" << nss.dbName().toString()));
     auto writeCmdObj = writeCmdObjBuilder.obj();
 
     // Parse original write command and set _id as query filter for new command object.
     if (commandName == "update") {
         auto parsedUpdateRequest = write_ops::UpdateCommandRequest::parse(
-            IDLParserContext("_clusterWriteWithoutShardKey"), writeCmdObj);
+            IDLParserContext("_clusterWriteWithoutShardKeyForUpdate"), writeCmdObj);
+
+        // The original query and collation are sent along with the modified command for the
+        // purposes of query sampling.
+        if (parsedUpdateRequest.getUpdates().front().getSampleId()) {
+            auto writeCommandRequestBase = write_ops::WriteCommandRequestBase(
+                parsedUpdateRequest.getWriteCommandRequestBase());
+            writeCommandRequestBase.setOriginalQuery(
+                parsedUpdateRequest.getUpdates().front().getQ());
+            writeCommandRequestBase.setOriginalCollation(
+                parsedUpdateRequest.getUpdates().front().getCollation());
+            parsedUpdateRequest.setWriteCommandRequestBase(writeCommandRequestBase);
+        }
+
         parsedUpdateRequest.getUpdates().front().setQ(targetDocId);
+        // Unset the collation because targeting by _id uses default collation.
+        parsedUpdateRequest.getUpdates().front().setCollation(boost::none);
         return parsedUpdateRequest.toBSON(BSONObj());
     } else if (commandName == "delete") {
         auto parsedDeleteRequest = write_ops::DeleteCommandRequest::parse(
-            IDLParserContext("_clusterWriteWithoutShardKey"), writeCmdObj);
+            IDLParserContext("_clusterWriteWithoutShardKeyForDelete"), writeCmdObj);
+
+        // The original query and collation are sent along with the modified command for the
+        // purposes of query sampling.
+        if (parsedDeleteRequest.getDeletes().front().getSampleId()) {
+            auto writeCommandRequestBase = write_ops::WriteCommandRequestBase(
+                parsedDeleteRequest.getWriteCommandRequestBase());
+            writeCommandRequestBase.setOriginalQuery(
+                parsedDeleteRequest.getDeletes().front().getQ());
+            writeCommandRequestBase.setOriginalCollation(
+                parsedDeleteRequest.getDeletes().front().getCollation());
+            parsedDeleteRequest.setWriteCommandRequestBase(writeCommandRequestBase);
+        }
+
         parsedDeleteRequest.getDeletes().front().setQ(targetDocId);
+        // Unset the collation because targeting by _id uses default collation.
+        parsedDeleteRequest.getDeletes().front().setCollation(boost::none);
         return parsedDeleteRequest.toBSON(BSONObj());
     } else if (commandName == "findandmodify" || commandName == "findAndModify") {
         auto parsedFindAndModifyRequest = write_ops::FindAndModifyCommandRequest::parse(
-            IDLParserContext("_clusterWriteWithoutShardKey"), writeCmdObj);
+            IDLParserContext("_clusterWriteWithoutShardKeyForFindAndModify"), writeCmdObj);
+
+        // The original query and collation are sent along with the modified command for the
+        // purposes of query sampling.
+        if (parsedFindAndModifyRequest.getSampleId()) {
+            parsedFindAndModifyRequest.setOriginalQuery(parsedFindAndModifyRequest.getQuery());
+            parsedFindAndModifyRequest.setOriginalCollation(
+                parsedFindAndModifyRequest.getCollation());
+        }
+
         parsedFindAndModifyRequest.setQuery(targetDocId);
+        // Unset the collation because targeting by _id uses default collation.
+        parsedFindAndModifyRequest.setCollation(boost::none);
         return parsedFindAndModifyRequest.toBSON(BSONObj());
     } else {
         uasserted(ErrorCodes::InvalidOptions,