SERVER-75284 Add validation for originalQuery and originalCollation

6 files changed, 174 insertions, 49 deletions

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 e2ef51ad8f9..9afeaf3a71d 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
@@ -369,6 +369,90 @@ function runAndVerifyCommand(testCase) {
         targetDocId: {_id: _id},
     };
     assert.commandFailedWithCode(mongosConn.runCommand(cmdObj), ErrorCodes.IllegalOperation);
+
+    // Cannot pass $_originalQuery as an external client in an update command.
+    cmdObj = {
+        update: collName,
+        updates: [
+            {
+                q: {},
+                u: {$set: {a: aFieldValue}},
+                sampleId: UUID(),
+            },
+        ],
+        $_originalQuery: {},
+    };
+    assert.commandFailedWithCode(shardConn.getCollection(collName).runCommand(cmdObj),
+                                 ErrorCodes.InvalidOptions);
+
+    // Cannot pass $_originalCollation as an external client in an update command.
+    cmdObj = {
+        update: collName,
+        updates: [
+            {
+                q: {},
+                u: {$set: {a: aFieldValue}},
+                sampleId: UUID(),
+                collation: {locale: "simple"},
+            },
+        ],
+        $_originalCollation: {},
+    };
+    assert.commandFailedWithCode(shardConn.getCollection(collName).runCommand(cmdObj),
+                                 ErrorCodes.InvalidOptions);
+
+    // Cannot pass $_originalQuery as an external client in a delete command.
+    cmdObj = {
+        delete: collName,
+        deletes: [
+            {
+                q: {},
+                limit: 1,
+                sampleId: UUID(),
+            },
+        ],
+        $_originalQuery: {},
+    };
+    assert.commandFailedWithCode(shardConn.getCollection(collName).runCommand(cmdObj),
+                                 ErrorCodes.InvalidOptions);
+
+    // Cannot pass $_originalCollation as an external client in a delete command.
+    cmdObj = {
+        delete: collName,
+        deletes: [
+            {
+                q: {},
+                limit: 1,
+                sampleId: UUID(),
+                collation: {locale: "simple"},
+            },
+        ],
+        $_originalCollation: {},
+    };
+    assert.commandFailedWithCode(shardConn.getCollection(collName).runCommand(cmdObj),
+                                 ErrorCodes.InvalidOptions);
+
+    // Cannot pass $_originalQuery as an external client in a findandmodify command.
+    cmdObj = {
+        findandmodify: collName,
+        query: {},
+        update: [{$set: {a: aFieldValue}}],
+        sampleId: UUID(),
+        $_originalQuery: {},
+    };
+    assert.commandFailedWithCode(shardConn.getCollection(collName).runCommand(cmdObj),
+                                 ErrorCodes.InvalidOptions);
+
+    // Cannot pass $_originalCollation as an external client in a findandmodify command.
+    cmdObj = {
+        findandmodify: collName,
+        query: {},
+        update: [{$set: {a: aFieldValue}}],
+        sampleId: UUID(),
+        $_originalCollation: {},
+    };
+    assert.commandFailedWithCode(shardConn.getCollection(collName).runCommand(cmdObj),
+                                 ErrorCodes.InvalidOptions);
 })();
 
 st.stop();
diff --git a/src/mongo/db/commands/find_and_modify.cpp b/src/mongo/db/commands/find_and_modify.cpp
index 7042a844da0..aa6b1a691ad 100644
--- a/src/mongo/db/commands/find_and_modify.cpp
+++ b/src/mongo/db/commands/find_and_modify.cpp
@@ -474,7 +474,7 @@ write_ops::FindAndModifyCommandReply CmdFindAndModify::Invocation::typedRun(
         opCtx, ns(), analyze_shard_key::SampledCommandNameEnum::kFindAndModify, req);
     if (sampleId) {
         analyze_shard_key::QueryAnalysisWriter::get(opCtx)
-            ->addFindAndModifyQuery(*sampleId, req)
+            ->addFindAndModifyQuery(opCtx, *sampleId, req)
             .getAsync([](auto) {});
     }
 
diff --git a/src/mongo/db/ops/write_ops_exec.cpp b/src/mongo/db/ops/write_ops_exec.cpp
index 8fc25577e3a..9aec8d49943 100644
--- a/src/mongo/db/ops/write_ops_exec.cpp
+++ b/src/mongo/db/ops/write_ops_exec.cpp
@@ -1401,7 +1401,7 @@ WriteResult performUpdates(OperationContext* opCtx,
             opCtx, ns, analyze_shard_key::SampledCommandNameEnum::kUpdate, singleOp);
         if (sampleId) {
             analyze_shard_key::QueryAnalysisWriter::get(opCtx)
-                ->addUpdateQuery(*sampleId, wholeOp, currentOpIndex)
+                ->addUpdateQuery(opCtx, *sampleId, wholeOp, currentOpIndex)
                 .getAsync([](auto) {});
         }
 
@@ -1672,7 +1672,7 @@ WriteResult performDeletes(OperationContext* opCtx,
         if (auto sampleId = analyze_shard_key::getOrGenerateSampleId(
                 opCtx, ns, analyze_shard_key::SampledCommandNameEnum::kDelete, singleOp)) {
             analyze_shard_key::QueryAnalysisWriter::get(opCtx)
-                ->addDeleteQuery(*sampleId, wholeOp, currentOpIndex)
+                ->addDeleteQuery(opCtx, *sampleId, wholeOp, currentOpIndex)
                 .getAsync([](auto) {});
         }
 
diff --git a/src/mongo/db/s/query_analysis_writer.cpp b/src/mongo/db/s/query_analysis_writer.cpp
index 0289bec8053..7404066b9c9 100644
--- a/src/mongo/db/s/query_analysis_writer.cpp
+++ b/src/mongo/db/s/query_analysis_writer.cpp
@@ -94,6 +94,11 @@ BSONObj createIndex(OperationContext* opCtx, const NamespaceString& nss, const B
     return resObj;
 }
 
+bool isInternalClient(const OperationContext* opCtx) {
+    return !opCtx->getClient()->session() ||
+        (opCtx->getClient()->session()->getTags() & transport::Session::kInternalClient);
+}
+
 struct SampledCommandRequest {
     UUID sampleId;
     NamespaceString nss;
@@ -119,7 +124,10 @@ SampledCommandRequest makeSampledReadCommand(const UUID& sampleId,
  * Returns a sampled update command for the update at 'opIndex' in the given update command.
  */
 SampledCommandRequest makeSampledUpdateCommandRequest(
-    const UUID& sampleId, const write_ops::UpdateCommandRequest& originalCmd, int opIndex) {
+    const OperationContext* opCtx,
+    const UUID& sampleId,
+    const write_ops::UpdateCommandRequest& originalCmd,
+    int opIndex) {
     auto op = originalCmd.getUpdates()[opIndex];
     if (op.getSampleId()) {
         tassert(ErrorCodes::IllegalOperation,
@@ -131,10 +139,14 @@ SampledCommandRequest makeSampledUpdateCommandRequest(
     // 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 here.
-    if (originalCmd.getOriginalQuery()) {
+    if (originalCmd.getOriginalQuery() || originalCmd.getOriginalCollation()) {
         tassert(7406500,
                 "Found a _clusterWithoutShardKey command with batch size > 1",
                 originalCmd.getUpdates().size() == 1);
+        uassert(ErrorCodes::InvalidOptions,
+                "Cannot specify '$_originalQuery' or '$_originalCollation' since they are internal "
+                "fields",
+                isInternalClient(opCtx));
         op.setQ(*originalCmd.getOriginalQuery());
         op.setCollation(originalCmd.getOriginalCollation());
     }
@@ -151,7 +163,10 @@ SampledCommandRequest makeSampledUpdateCommandRequest(
  * Returns a sampled delete command for the delete at 'opIndex' in the given delete command.
  */
 SampledCommandRequest makeSampledDeleteCommandRequest(
-    const UUID& sampleId, const write_ops::DeleteCommandRequest& originalCmd, int opIndex) {
+    const OperationContext* opCtx,
+    const UUID& sampleId,
+    const write_ops::DeleteCommandRequest& originalCmd,
+    int opIndex) {
     auto op = originalCmd.getDeletes()[opIndex];
     if (op.getSampleId()) {
         tassert(ErrorCodes::IllegalOperation,
@@ -163,10 +178,14 @@ SampledCommandRequest makeSampledDeleteCommandRequest(
     // 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 here.
-    if (originalCmd.getOriginalQuery()) {
+    if (originalCmd.getOriginalQuery() || originalCmd.getOriginalCollation()) {
         tassert(7406501,
                 "Found a _clusterWithoutShardKey command with batch size > 1",
                 originalCmd.getDeletes().size() == 1);
+        uassert(ErrorCodes::InvalidOptions,
+                "Cannot specify '$_originalQuery' or '$_originalCollation' since they are internal "
+                "fields",
+                isInternalClient(opCtx));
         op.setQ(*originalCmd.getOriginalQuery());
         op.setCollation(originalCmd.getOriginalCollation());
     }
@@ -183,7 +202,9 @@ SampledCommandRequest makeSampledDeleteCommandRequest(
  * Returns a sampled findAndModify command for the given findAndModify command.
  */
 SampledCommandRequest makeSampledFindAndModifyCommandRequest(
-    const UUID& sampleId, const write_ops::FindAndModifyCommandRequest& originalCmd) {
+    OperationContext* opCtx,
+    const UUID& sampleId,
+    const write_ops::FindAndModifyCommandRequest& originalCmd) {
     write_ops::FindAndModifyCommandRequest sampledCmd(originalCmd.getNamespace());
     if (sampledCmd.getSampleId()) {
         tassert(ErrorCodes::IllegalOperation,
@@ -195,7 +216,11 @@ SampledCommandRequest makeSampledFindAndModifyCommandRequest(
     // 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 here.
-    if (originalCmd.getOriginalQuery()) {
+    if (originalCmd.getOriginalQuery() || originalCmd.getOriginalCollation()) {
+        uassert(ErrorCodes::InvalidOptions,
+                "Cannot specify '$_originalQuery' or '$_originalCollation' since they are internal "
+                "fields",
+                isInternalClient(opCtx));
         sampledCmd.setQuery(*originalCmd.getOriginalQuery());
         sampledCmd.setCollation(originalCmd.getOriginalCollation());
     } else {
@@ -626,12 +651,16 @@ ExecutorFuture<void> QueryAnalysisWriter::_addReadQuery(
 }
 
 ExecutorFuture<void> QueryAnalysisWriter::addUpdateQuery(
-    const UUID& sampleId, const write_ops::UpdateCommandRequest& updateCmd, int opIndex) {
+    OperationContext* originalOpCtx,
+    const UUID& sampleId,
+    const write_ops::UpdateCommandRequest& updateCmd,
+    int opIndex) {
     invariant(_executor);
 
     return ExecutorFuture<void>(_executor)
         .then([this,
-               sampledUpdateCmd = makeSampledUpdateCommandRequest(sampleId, updateCmd, opIndex)]() {
+               sampledUpdateCmd =
+                   makeSampledUpdateCommandRequest(originalOpCtx, sampleId, updateCmd, opIndex)]() {
             auto opCtxHolder = cc().makeOperationContext();
             auto opCtx = opCtxHolder.get();
 
@@ -673,19 +702,23 @@ ExecutorFuture<void> QueryAnalysisWriter::addUpdateQuery(
 }
 
 ExecutorFuture<void> QueryAnalysisWriter::addUpdateQuery(
-    const write_ops::UpdateCommandRequest& updateCmd, int opIndex) {
+    OperationContext* opCtx, const write_ops::UpdateCommandRequest& updateCmd, int opIndex) {
     auto sampleId = updateCmd.getUpdates()[opIndex].getSampleId();
     invariant(sampleId);
-    return addUpdateQuery(*sampleId, updateCmd, opIndex);
+    return addUpdateQuery(opCtx, *sampleId, updateCmd, opIndex);
 }
 
 ExecutorFuture<void> QueryAnalysisWriter::addDeleteQuery(
-    const UUID& sampleId, const write_ops::DeleteCommandRequest& deleteCmd, int opIndex) {
+    OperationContext* originalOpCtx,
+    const UUID& sampleId,
+    const write_ops::DeleteCommandRequest& deleteCmd,
+    int opIndex) {
     invariant(_executor);
 
     return ExecutorFuture<void>(_executor)
         .then([this,
-               sampledDeleteCmd = makeSampledDeleteCommandRequest(sampleId, deleteCmd, opIndex)]() {
+               sampledDeleteCmd =
+                   makeSampledDeleteCommandRequest(originalOpCtx, sampleId, deleteCmd, opIndex)]() {
             auto opCtxHolder = cc().makeOperationContext();
             auto opCtx = opCtxHolder.get();
 
@@ -727,20 +760,22 @@ ExecutorFuture<void> QueryAnalysisWriter::addDeleteQuery(
 }
 
 ExecutorFuture<void> QueryAnalysisWriter::addDeleteQuery(
-    const write_ops::DeleteCommandRequest& deleteCmd, int opIndex) {
+    OperationContext* opCtx, const write_ops::DeleteCommandRequest& deleteCmd, int opIndex) {
     auto sampleId = deleteCmd.getDeletes()[opIndex].getSampleId();
     invariant(sampleId);
-    return addDeleteQuery(*sampleId, deleteCmd, opIndex);
+    return addDeleteQuery(opCtx, *sampleId, deleteCmd, opIndex);
 }
 
 ExecutorFuture<void> QueryAnalysisWriter::addFindAndModifyQuery(
-    const UUID& sampleId, const write_ops::FindAndModifyCommandRequest& findAndModifyCmd) {
+    OperationContext* originalOpCtx,
+    const UUID& sampleId,
+    const write_ops::FindAndModifyCommandRequest& findAndModifyCmd) {
     invariant(_executor);
 
     return ExecutorFuture<void>(_executor)
         .then([this,
-               sampledFindAndModifyCmd =
-                   makeSampledFindAndModifyCommandRequest(sampleId, findAndModifyCmd)]() {
+               sampledFindAndModifyCmd = makeSampledFindAndModifyCommandRequest(
+                   originalOpCtx, sampleId, findAndModifyCmd)]() {
             auto opCtxHolder = cc().makeOperationContext();
             auto opCtx = opCtxHolder.get();
 
@@ -782,10 +817,10 @@ ExecutorFuture<void> QueryAnalysisWriter::addFindAndModifyQuery(
 }
 
 ExecutorFuture<void> QueryAnalysisWriter::addFindAndModifyQuery(
-    const write_ops::FindAndModifyCommandRequest& findAndModifyCmd) {
+    OperationContext* opCtx, const write_ops::FindAndModifyCommandRequest& findAndModifyCmd) {
     auto sampleId = findAndModifyCmd.getSampleId();
     invariant(sampleId);
-    return addFindAndModifyQuery(*sampleId, findAndModifyCmd);
+    return addFindAndModifyQuery(opCtx, *sampleId, findAndModifyCmd);
 }
 
 ExecutorFuture<void> QueryAnalysisWriter::addDiff(const UUID& sampleId,
diff --git a/src/mongo/db/s/query_analysis_writer.h b/src/mongo/db/s/query_analysis_writer.h
index 18f980dd20f..0b5d5f02ac9 100644
--- a/src/mongo/db/s/query_analysis_writer.h
+++ b/src/mongo/db/s/query_analysis_writer.h
@@ -173,22 +173,28 @@ public:
                                           const BSONObj& filter,
                                           const BSONObj& collation);
 
-    ExecutorFuture<void> addUpdateQuery(const UUID& sampleId,
+    ExecutorFuture<void> addUpdateQuery(OperationContext* opCtx,
+                                        const UUID& sampleId,
                                         const write_ops::UpdateCommandRequest& updateCmd,
                                         int opIndex);
-    ExecutorFuture<void> addUpdateQuery(const write_ops::UpdateCommandRequest& updateCmd,
+    ExecutorFuture<void> addUpdateQuery(OperationContext* opCtx,
+                                        const write_ops::UpdateCommandRequest& updateCmd,
                                         int opIndex);
 
-    ExecutorFuture<void> addDeleteQuery(const UUID& sampleId,
+    ExecutorFuture<void> addDeleteQuery(OperationContext* opCtx,
+                                        const UUID& sampleId,
                                         const write_ops::DeleteCommandRequest& deleteCmd,
                                         int opIndex);
-    ExecutorFuture<void> addDeleteQuery(const write_ops::DeleteCommandRequest& deleteCmd,
+    ExecutorFuture<void> addDeleteQuery(OperationContext* opCtx,
+                                        const write_ops::DeleteCommandRequest& deleteCmd,
                                         int opIndex);
 
     ExecutorFuture<void> addFindAndModifyQuery(
-        const UUID& sampleId, const write_ops::FindAndModifyCommandRequest& findAndModifyCmd);
-    ExecutorFuture<void> addFindAndModifyQuery(
+        OperationContext* opCtx,
+        const UUID& sampleId,
         const write_ops::FindAndModifyCommandRequest& findAndModifyCmd);
+    ExecutorFuture<void> addFindAndModifyQuery(
+        OperationContext* opCtx, const write_ops::FindAndModifyCommandRequest& findAndModifyCmd);
 
     ExecutorFuture<void> addDiff(const UUID& sampleId,
                                  const NamespaceString& nss,
diff --git a/src/mongo/db/s/query_analysis_writer_test.cpp b/src/mongo/db/s/query_analysis_writer_test.cpp
index c21217c8e04..c40e596a7d3 100644
--- a/src/mongo/db/s/query_analysis_writer_test.cpp
+++ b/src/mongo/db/s/query_analysis_writer_test.cpp
@@ -740,7 +740,7 @@ TEST_F(QueryAnalysisWriterTest, AggregateQuery) {
 DEATH_TEST_F(QueryAnalysisWriterTest, UpdateQueryNotMarkedForSampling, "invariant") {
     auto& writer = *QueryAnalysisWriter::get(operationContext());
     auto [originalCmd, _] = makeUpdateCommandRequest(nss0, 1, {} /* markForSampling */);
-    writer.addUpdateQuery(originalCmd, 0).get();
+    writer.addUpdateQuery(operationContext(), originalCmd, 0).get();
 }
 
 TEST_F(QueryAnalysisWriterTest, UpdateQueriesMarkedForSampling) {
@@ -750,8 +750,8 @@ TEST_F(QueryAnalysisWriterTest, UpdateQueriesMarkedForSampling) {
         makeUpdateCommandRequest(nss0, 3, {0, 2} /* markForSampling */);
     ASSERT_EQ(expectedSampledCmds.size(), 2U);
 
-    writer.addUpdateQuery(originalCmd, 0).get();
-    writer.addUpdateQuery(originalCmd, 2).get();
+    writer.addUpdateQuery(operationContext(), originalCmd, 0).get();
+    writer.addUpdateQuery(operationContext(), originalCmd, 2).get();
     ASSERT_EQ(writer.getQueriesCountForTest(), 2);
     writer.flushQueriesForTest(operationContext());
     ASSERT_EQ(writer.getQueriesCountForTest(), 0);
@@ -768,7 +768,7 @@ TEST_F(QueryAnalysisWriterTest, UpdateQueriesMarkedForSampling) {
 DEATH_TEST_F(QueryAnalysisWriterTest, DeleteQueryNotMarkedForSampling, "invariant") {
     auto& writer = *QueryAnalysisWriter::get(operationContext());
     auto [originalCmd, _] = makeDeleteCommandRequest(nss0, 1, {} /* markForSampling */);
-    writer.addDeleteQuery(originalCmd, 0).get();
+    writer.addDeleteQuery(operationContext(), originalCmd, 0).get();
 }
 
 TEST_F(QueryAnalysisWriterTest, DeleteQueriesMarkedForSampling) {
@@ -778,8 +778,8 @@ TEST_F(QueryAnalysisWriterTest, DeleteQueriesMarkedForSampling) {
         makeDeleteCommandRequest(nss0, 3, {1, 2} /* markForSampling */);
     ASSERT_EQ(expectedSampledCmds.size(), 2U);
 
-    writer.addDeleteQuery(originalCmd, 1).get();
-    writer.addDeleteQuery(originalCmd, 2).get();
+    writer.addDeleteQuery(operationContext(), originalCmd, 1).get();
+    writer.addDeleteQuery(operationContext(), originalCmd, 2).get();
     ASSERT_EQ(writer.getQueriesCountForTest(), 2);
     writer.flushQueriesForTest(operationContext());
     ASSERT_EQ(writer.getQueriesCountForTest(), 0);
@@ -797,7 +797,7 @@ DEATH_TEST_F(QueryAnalysisWriterTest, FindAndModifyQueryNotMarkedForSampling, "i
     auto& writer = *QueryAnalysisWriter::get(operationContext());
     auto [originalCmd, _] =
         makeFindAndModifyCommandRequest(nss0, true /* isUpdate */, false /* markForSampling */);
-    writer.addFindAndModifyQuery(originalCmd).get();
+    writer.addFindAndModifyQuery(operationContext(), originalCmd).get();
 }
 
 TEST_F(QueryAnalysisWriterTest, FindAndModifyQueryUpdateMarkedForSampling) {
@@ -808,7 +808,7 @@ TEST_F(QueryAnalysisWriterTest, FindAndModifyQueryUpdateMarkedForSampling) {
     ASSERT_EQ(expectedSampledCmds.size(), 1U);
     auto [sampleId, expectedSampledCmd] = *expectedSampledCmds.begin();
 
-    writer.addFindAndModifyQuery(originalCmd).get();
+    writer.addFindAndModifyQuery(operationContext(), originalCmd).get();
     ASSERT_EQ(writer.getQueriesCountForTest(), 1);
     writer.flushQueriesForTest(operationContext());
     ASSERT_EQ(writer.getQueriesCountForTest(), 0);
@@ -828,7 +828,7 @@ TEST_F(QueryAnalysisWriterTest, FindAndModifyQueryRemoveMarkedForSampling) {
     ASSERT_EQ(expectedSampledCmds.size(), 1U);
     auto [sampleId, expectedSampledCmd] = *expectedSampledCmds.begin();
 
-    writer.addFindAndModifyQuery(originalCmd).get();
+    writer.addFindAndModifyQuery(operationContext(), originalCmd).get();
     ASSERT_EQ(writer.getQueriesCountForTest(), 1);
     writer.flushQueriesForTest(operationContext());
     ASSERT_EQ(writer.getQueriesCountForTest(), 0);
@@ -859,8 +859,8 @@ TEST_F(QueryAnalysisWriterTest, MultipleQueriesAndCollections) {
     auto originalCountFilter = makeNonEmptyFilter();
     auto originalCountCollation = makeNonEmptyCollation();
 
-    writer.addDeleteQuery(originalDeleteCmd, 1).get();
-    writer.addUpdateQuery(originalUpdateCmd, 0).get();
+    writer.addDeleteQuery(operationContext(), originalDeleteCmd, 1).get();
+    writer.addUpdateQuery(operationContext(), originalUpdateCmd, 0).get();
     writer.addCountQuery(countSampleId, nss1, originalCountFilter, originalCountCollation).get();
     ASSERT_EQ(writer.getQueriesCountForTest(), 3);
     writer.flushQueriesForTest(operationContext());
@@ -917,7 +917,7 @@ TEST_F(QueryAnalysisWriterTest, DuplicateQueries) {
                                    originalFindFilter,
                                    originalFindCollation);
 
-    writer.addUpdateQuery(originalUpdateCmd, 0).get();
+    writer.addUpdateQuery(operationContext(), originalUpdateCmd, 0).get();
     writer
         .addFindQuery(findSampleId,
                       nss0,
@@ -1067,10 +1067,10 @@ TEST_F(QueryAnalysisWriterTest, FlushAfterAddUpdateIfExceedsSizeLimit) {
                                  std::string(maxMemoryUsageBytes / 2, 'a') /* filterFieldName */);
     ASSERT_EQ(expectedSampledCmds.size(), 2U);
 
-    writer.addUpdateQuery(originalCmd, 0).get();
+    writer.addUpdateQuery(operationContext(), originalCmd, 0).get();
     ASSERT_EQ(writer.getQueriesCountForTest(), 1);
     // Adding the next query causes the size to exceed the limit.
-    writer.addUpdateQuery(originalCmd, 2).get();
+    writer.addUpdateQuery(operationContext(), originalCmd, 2).get();
     ASSERT_EQ(writer.getQueriesCountForTest(), 0);
 
     ASSERT_EQ(getSampledQueryDocumentsCount(nss0), 2);
@@ -1095,10 +1095,10 @@ TEST_F(QueryAnalysisWriterTest, FlushAfterAddDeleteIfExceedsSizeLimit) {
                                  std::string(maxMemoryUsageBytes / 2, 'a') /* filterFieldName */);
     ASSERT_EQ(expectedSampledCmds.size(), 2U);
 
-    writer.addDeleteQuery(originalCmd, 0).get();
+    writer.addDeleteQuery(operationContext(), originalCmd, 0).get();
     ASSERT_EQ(writer.getQueriesCountForTest(), 1);
     // Adding the next query causes the size to exceed the limit.
-    writer.addDeleteQuery(originalCmd, 1).get();
+    writer.addDeleteQuery(operationContext(), originalCmd, 1).get();
     ASSERT_EQ(writer.getQueriesCountForTest(), 0);
 
     ASSERT_EQ(getSampledQueryDocumentsCount(nss0), 2);
@@ -1133,10 +1133,10 @@ TEST_F(QueryAnalysisWriterTest, FlushAfterAddFindAndModifyIfExceedsSizeLimit) {
     ASSERT_EQ(expectedSampledCmds0.size(), 1U);
     auto [sampleId1, expectedSampledCmd1] = *expectedSampledCmds1.begin();
 
-    writer.addFindAndModifyQuery(originalCmd0).get();
+    writer.addFindAndModifyQuery(operationContext(), originalCmd0).get();
     ASSERT_EQ(writer.getQueriesCountForTest(), 1);
     // Adding the next query causes the size to exceed the limit.
-    writer.addFindAndModifyQuery(originalCmd1).get();
+    writer.addFindAndModifyQuery(operationContext(), originalCmd1).get();
     ASSERT_EQ(writer.getQueriesCountForTest(), 0);
 
     ASSERT_EQ(getSampledQueryDocumentsCount(nss0), 1);
@@ -1539,16 +1539,16 @@ void QueryAnalysisWriterTest::assertNoSampling(const NamespaceString& nss, const
     ASSERT_EQ(writer.getQueriesCountForTest(), 0);
 
     auto originalUpdateCmd = makeUpdateCommandRequest(nss, 1, {0} /* markForSampling */).first;
-    writer.addUpdateQuery(originalUpdateCmd, 0).get();
+    writer.addUpdateQuery(operationContext(), originalUpdateCmd, 0).get();
     ASSERT_EQ(writer.getQueriesCountForTest(), 0);
 
     auto originalDeleteCmd = makeDeleteCommandRequest(nss, 1, {0} /* markForSampling */).first;
-    writer.addDeleteQuery(originalDeleteCmd, 0).get();
+    writer.addDeleteQuery(operationContext(), originalDeleteCmd, 0).get();
     ASSERT_EQ(writer.getQueriesCountForTest(), 0);
 
     auto originalFindAndModifyCmd =
         makeFindAndModifyCommandRequest(nss, true /* isUpdate */, true /* markForSampling */).first;
-    writer.addFindAndModifyQuery(originalFindAndModifyCmd).get();
+    writer.addFindAndModifyQuery(operationContext(), originalFindAndModifyCmd).get();
 
     writer
         .addDiff(UUID::gen() /* sampleId */,