3 files changed, 54 insertions, 31 deletions

diff --git a/src/mongo/s/commands/strategy.cpp b/src/mongo/s/commands/strategy.cpp
index 55081a9a04b..4ff6bdfe3e9 100644
--- a/src/mongo/s/commands/strategy.cpp
+++ b/src/mongo/s/commands/strategy.cpp
@@ -257,38 +257,7 @@ void execCommandClient(OperationContext* opCtx,
     }
 
     auto& readConcernArgs = repl::ReadConcernArgs::get(opCtx);
-
     if (readConcernArgs.getLevel() == repl::ReadConcernLevel::kSnapshotReadConcern) {
-        // TODO SERVER-33708.
-        if (!invocation->supportsReadConcern(readConcernArgs.getLevel())) {
-            auto body = result->getBodyBuilder();
-            CommandHelpers::appendCommandStatusNoThrow(
-                body,
-                Status(ErrorCodes::InvalidOptions,
-                       str::stream()
-                           << "read concern snapshot is not supported on mongos for the command "
-                           << c->getName()));
-            return;
-        }
-
-        if (!opCtx->getTxnNumber()) {
-            auto body = result->getBodyBuilder();
-            CommandHelpers::appendCommandStatusNoThrow(
-                body,
-                Status(ErrorCodes::InvalidOptions,
-                       "read concern snapshot is supported only in a transaction"));
-            return;
-        }
-
-        if (readConcernArgs.getArgsAtClusterTime()) {
-            auto body = result->getBodyBuilder();
-            CommandHelpers::appendCommandStatusNoThrow(
-                body,
-                Status(ErrorCodes::InvalidOptions,
-                       "read concern snapshot is not supported with atClusterTime on mongos"));
-            return;
-        }
-
         uassert(ErrorCodes::InvalidOptions,
                 "read concern snapshot is only supported in a multi-statement transaction",
                 TransactionRouter::get(opCtx));
@@ -399,6 +368,16 @@ void runCommand(OperationContext* opCtx,
         return;
     }
 
+    if (readConcernArgs.getLevel() == repl::ReadConcernLevel::kSnapshotReadConcern) {
+        uassert(ErrorCodes::InvalidOptions,
+                str::stream() << "read concern snapshot is not supported on mongos for the command "
+                              << commandName,
+                invocation->supportsReadConcern(readConcernArgs.getLevel()));
+        uassert(ErrorCodes::InvalidOptions,
+                "read concern snapshot is not supported with atClusterTime on mongos",
+                !readConcernArgs.getArgsAtClusterTime());
+    }
+
     boost::optional<ScopedRouterSession> scopedSession;
     auto osi =
         initializeOperationSessionInfo(opCtx, request.body, command->requiresAuth(), true, true);
diff --git a/src/mongo/s/transaction_router.cpp b/src/mongo/s/transaction_router.cpp
index 3b44b615eab..013f212148d 100644
--- a/src/mongo/s/transaction_router.cpp
+++ b/src/mongo/s/transaction_router.cpp
@@ -495,6 +495,8 @@ void TransactionRouter::beginOrContinueTxn(OperationContext* opCtx,
         uassert(ErrorCodes::InvalidOptions,
                 "Only the first command in a transaction may specify a readConcern",
                 repl::ReadConcernArgs::get(opCtx).isEmpty());
+
+        repl::ReadConcernArgs::get(opCtx) = _readConcernArgs;
     }
 
     if (_txnNumber == txnNumber) {
diff --git a/src/mongo/s/transaction_router_test.cpp b/src/mongo/s/transaction_router_test.cpp
index b0da4b30af0..cb211627533 100644
--- a/src/mongo/s/transaction_router_test.cpp
+++ b/src/mongo/s/transaction_router_test.cpp
@@ -30,6 +30,7 @@
 
 #include "mongo/client/remote_command_targeter_mock.h"
 #include "mongo/db/logical_clock.h"
+#include "mongo/db/repl/read_concern_args.h"
 #include "mongo/s/sharding_router_test_fixture.h"
 #include "mongo/s/transaction_router.h"
 #include "mongo/unittest/death_test.h"
@@ -1366,5 +1367,46 @@ TEST_F(TransactionRouterTest, ImplicitAbortIgnoresErrors) {
     future.timed_get(kFutureTimeout);
 }
 
+TEST_F(TransactionRouterTest, ContinuingTransactionPlacesItsReadConcernOnOpCtx) {
+    TxnNumber txnNum{3};
+
+    TransactionRouter txnRouter({});
+    txnRouter.checkOut();
+    txnRouter.beginOrContinueTxn(operationContext(), txnNum, true);
+    txnRouter.setAtClusterTimeToLatestTime(operationContext());
+
+    repl::ReadConcernArgs::get(operationContext()) = repl::ReadConcernArgs();
+    txnRouter.beginOrContinueTxn(operationContext(), txnNum, false);
+
+    ASSERT(repl::ReadConcernArgs::get(operationContext()).getLevel() ==
+           repl::ReadConcernLevel::kSnapshotReadConcern);
+}
+
+TEST_F(TransactionRouterTest, SubsequentStatementCanSelectAtClusterTimeIfNotSelectedYet) {
+    TxnNumber txnNum{3};
+
+    TransactionRouter txnRouter({});
+    txnRouter.checkOut();
+    txnRouter.beginOrContinueTxn(operationContext(), txnNum, true);
+
+    // First statement does not select an atClusterTime, but does not target any participants.
+
+    repl::ReadConcernArgs::get(operationContext()) = repl::ReadConcernArgs();
+    txnRouter.beginOrContinueTxn(operationContext(), txnNum, false);
+
+    // Subsequent statement does select an atClusterTime and does target a participant.
+    txnRouter.setAtClusterTimeToLatestTime(operationContext());
+
+    BSONObj expectedReadConcern = BSON("level"
+                                       << "snapshot"
+                                       << "atClusterTime"
+                                       << kInMemoryLogicalTime.asTimestamp());
+
+    auto newCmd = txnRouter.attachTxnFieldsIfNeeded(shard1,
+                                                    BSON("insert"
+                                                         << "test"));
+    ASSERT_BSONOBJ_EQ(expectedReadConcern, newCmd["readConcern"].Obj());
+}
+
 }  // unnamed namespace
 }  // namespace mongo