SERVER-64069 Relax restriction that changing a document's shard key value must have a transaction number

9 files changed, 283 insertions, 83 deletions

diff --git a/jstests/sharding/resharding_replicate_updates_as_insert_delete.js b/jstests/sharding/resharding_replicate_updates_as_insert_delete.js
index 41ce0efd0f4..ccb43892c90 100644
--- a/jstests/sharding/resharding_replicate_updates_as_insert_delete.js
+++ b/jstests/sharding/resharding_replicate_updates_as_insert_delete.js
@@ -52,10 +52,18 @@ reshardingTest.withReshardingInBackground(  //
         const tempColl = mongos.getCollection(tempNs);
         assert.soon(() => tempColl.findOne(docToUpdate) !== null);
 
-        assert.commandFailedWithCode(
-            testColl.update({_id: 0, x: 2, s: 2}, {$set: {y: 10}}),
-            ErrorCodes.IllegalOperation,
-            'was able to update value under new shard key as ordinary write');
+        // When the updateDocumentShardKeyUsingTransactionApi feature flag is enabled, ordinary
+        // updates that modify a document's shard key will complete.
+        assert.commandWorked(testColl.insert({_id: 1, x: 2, s: 2, y: 2}));
+        const updateRes = testColl.update({_id: 1, x: 2, s: 2}, {$set: {y: 10}});
+        if (updateDocumentShardKeyUsingTransactionApiEnabled) {
+            assert.commandWorked(updateRes);
+        } else {
+            assert.commandFailedWithCode(
+                updateRes,
+                ErrorCodes.IllegalOperation,
+                'was able to update value under new shard key as ordinary write');
+        }
 
         const session = testColl.getMongo().startSession({retryWrites: true});
         const sessionColl =
diff --git a/jstests/sharding/update_compound_shard_key.js b/jstests/sharding/update_compound_shard_key.js
index e4a9d72d7c9..be331f2bd3c 100644
--- a/jstests/sharding/update_compound_shard_key.js
+++ b/jstests/sharding/update_compound_shard_key.js
@@ -14,7 +14,8 @@ load("jstests/sharding/libs/update_shard_key_helpers.js");
 
 const st = new ShardingTest({mongos: 1, shards: 3});
 
-enableCoordinateCommitReturnImmediatelyAfterPersistingDecision(st);
+const updateDocumentShardKeyUsingTransactionApiEnabled =
+    isUpdateDocumentShardKeyUsingTransactionApiEnabled(st.s);
 
 const kDbName = 'update_compound_sk';
 const ns = kDbName + '.coll';
@@ -85,6 +86,22 @@ function assertUpdateWorkedWithNoMatchingDoc(query, update, isUpsert, inTransact
               st.s.getDB(kDbName).coll.find(update["$set"] ? update["$set"] : update).itcount());
 }
 
+/**
+ * Updates to a document's shard key that would change the document's owning shard will succeed if
+ * the updateDocumentShardKeyUsingTransactionApi feature flag is enabled.
+ */
+function assertWouldChangeOwningShardUpdateResult(res, expectedUpdatedDoc) {
+    const numDocsUpdated = st.s.getDB(kDbName).coll.find(expectedUpdatedDoc).itcount();
+
+    if (updateDocumentShardKeyUsingTransactionApiEnabled) {
+        assert.commandWorked(res);
+        assert.eq(1, numDocsUpdated);
+    } else {
+        assert.commandFailedWithCode(res, ErrorCodes.IllegalOperation);
+        assert.eq(0, numDocsUpdated);
+    }
+}
+
 //
 // Update Type Replacement-style.
 //
@@ -153,23 +170,29 @@ assert.commandFailedWithCode(
 //
 
 // Case when upsert needs to insert a new document and the new document should belong in a shard
-// other than the one targeted by the update. These upserts can only succeed in a
-// multi-statement transaction or with retryWrites: true.
+// other than the one targeted by the update. These upserts can only succeed when the
+// updateDocumentShardKeyUsingTransactionApi feature flag is enabled or if not enabled, when the
+// upsert is run in a multi-statement transaction or with retryWrites: true.
 const updateDoc = {
     x: 1110,
     y: 55,
     z: 3,
     replStyleUpdate: true
 };
-assert.commandFailedWithCode(
-    st.s.getDB(kDbName).coll.update({x: 4, y: 0, z: 0}, updateDoc, {upsert: true}),
-    ErrorCodes.IllegalOperation);
+const updateRes = st.s.getDB(kDbName).coll.update({x: 4, y: 0, z: 0}, updateDoc, {upsert: true});
+assertWouldChangeOwningShardUpdateResult(updateRes, updateDoc);
 
-// The above upsert works with transactions.
+const updateDocTxn = {
+    x: 1110,
+    y: 55,
+    z: 4,
+    replStyleUpdate: true
+};
 session.startTransaction();
-assertUpdateWorkedWithNoMatchingDoc({x: 4, y: 0, z: 0}, updateDoc, true, true);
+assertUpdateWorkedWithNoMatchingDoc(
+    {x: 4, y: 0, z: 0}, updateDocTxn, true /*isUpsert*/, true /*inTransaction*/);
 assert.commandWorked(session.commitTransaction_forTesting());
-assert.eq(1, sessionDB.coll.find(updateDoc).itcount());
+assert.eq(1, sessionDB.coll.find(updateDocTxn).itcount());
 
 // Full shard key not specified in query.
 
@@ -257,20 +280,26 @@ assert.commandFailedWithCode(st.s.getDB(kDbName).coll.update(
 // Test upsert-specific behaviours.
 
 // Case when upsert needs to insert a new document and the new document should belong in a shard
-// other than the one targeted by the update. These upserts can only succeed in a
-// multi-statement transaction or with retryWrites: true.
-const update = {
-    "$set": {x: 2110, y: 55, z: 3, opStyle: true}
+// other than the one targeted by the update. These upserts can only succeed when the
+// updateDocumentShardKeyUsingTransactionApi feature flag is enabled or if not enabled, when the
+// upsert is run in a multi-statement transaction or with retryWrites: true.
+const upsertDoc = {
+    x: 2110,
+    y: 55,
+    z: 3,
+    opStyle: true
 };
-assert.commandFailedWithCode(
-    st.s.getDB(kDbName).coll.update({x: 4, y: 0, z: 0, opStyle: true}, update, {upsert: true}),
-    ErrorCodes.IllegalOperation);
+const upsertRes =
+    st.s.getDB(kDbName).coll.update({x: 4, y: 0, z: 0, opStyle: true}, upsertDoc, {upsert: true});
+assertWouldChangeOwningShardUpdateResult(upsertRes, upsertDoc);
 
-// The above upsert works with transactions.
+const upsertDocTxn = {
+    "$set": {x: 2110, y: 55, z: 4, opStyle: true}
+};
 session.startTransaction();
-assertUpdateWorkedWithNoMatchingDoc({x: 4, y: 0, z: 0, opStyle: true}, update, true, true);
+assertUpdateWorkedWithNoMatchingDoc({x: 4, y: 0, z: 0, opStyle: true}, upsertDocTxn, true, true);
 assert.commandWorked(session.commitTransaction_forTesting());
-assert.eq(1, sessionDB.coll.find(update["$set"]).itcount());
+assert.eq(1, sessionDB.coll.find(upsertDocTxn["$set"]).itcount());
 
 // Full shard key not specified in query.
 
@@ -372,35 +401,44 @@ assert.commandFailedWithCode(
 // Test upsert-specific behaviours.
 
 // Case when upsert needs to insert a new document and the new document should belong in a shard
-// other than the one targeted by the update. These upserts can only succeed in a
-// multi-statement transaction or with retryWrites: true.
-assert.commandFailedWithCode(
-    st.s.getDB(kDbName).coll.update({x: 4, y: 0, z: 0},
-                                    [{
-                                        "$project": {
-                                            x: {$literal: 2111},
-                                            y: {$literal: 55},
-                                            z: {$literal: 3},
-                                            pipelineUpdate: {$literal: true}
-                                        }
-                                    }],
-                                    {upsert: true}),
-    ErrorCodes.IllegalOperation);
-
-// The above upsert works with transactions.
+// other than the one targeted by the update. These upserts can only succeed when run in a
+// multi-statement transaction or with retryWrites: true or if the
+// updateDocumentShardKeyUsingTransactionApi feature flag is enabled.
+const upsertProjectDoc = {
+    x: 2111,
+    y: 55,
+    z: 3
+};
+const upsertProjectRes = st.s.getDB(kDbName).coll.update({x: 4, y: 0, z: 0},
+                                                         [{
+                                                             "$project": {
+                                                                 x: {$literal: 2111},
+                                                                 y: {$literal: 55},
+                                                                 z: {$literal: 3},
+                                                                 pipelineUpdate: {$literal: true}
+                                                             }
+                                                         }],
+                                                         {upsert: true});
+assertWouldChangeOwningShardUpdateResult(upsertProjectRes, upsertProjectDoc);
+
+const upsertProjectTxnDoc = {
+    x: 2111,
+    y: 55,
+    z: 4
+};
 session.startTransaction();
 assertUpdateWorkedWithNoMatchingDoc({x: 4, y: 0, z: 0, pipelineUpdate: true},
                                     [{
                                         "$project": {
                                             x: {$literal: 2111},
                                             y: {$literal: 55},
-                                            z: {$literal: 3},
+                                            z: {$literal: 4},
                                             pipelineUpdate: {$literal: true}
                                         }
                                     }],
                                     true);
 assert.commandWorked(session.commitTransaction_forTesting());
-assert.eq(1, sessionDB.coll.find({x: 2111, y: 55, z: 3, pipelineUpdate: true}).itcount());
+assert.eq(1, sessionDB.coll.find(upsertProjectTxnDoc).itcount());
 
 // Full shard key not specified in query.
 assert.commandFailedWithCode(st.s.getDB(kDbName).coll.update(
diff --git a/jstests/sharding/update_shard_key_doc_moves_shards_without_txn_number.js b/jstests/sharding/update_shard_key_doc_moves_shards_without_txn_number.js
new file mode 100644
index 00000000000..884a940607e
--- /dev/null
+++ b/jstests/sharding/update_shard_key_doc_moves_shards_without_txn_number.js
@@ -0,0 +1,147 @@
+/**
+ * Test that a client not running a retryable write or transaction can update a document's shard
+ * key if the updateDocumentShardKeyUsingTransactionApi feature flag is enabled.
+ *
+ * @tags: [
+ *    requires_sharding,
+ *    uses_transactions,
+ *    uses_multi_shard_transaction,
+ * ]
+ */
+(function() {
+"use strict";
+
+load("jstests/libs/fail_point_util.js");
+load("jstests/sharding/libs/sharded_transactions_helpers.js");
+
+const st = new ShardingTest({shards: 2});
+const shard0Primary = st.rs0.getPrimary();
+
+const dbName = "testDb";
+const collName = "testColl";
+const ns = dbName + "." + collName;
+const db = st.getDB(dbName);
+const testColl = db.getCollection(collName);
+
+const updateDocumentShardKeyUsingTransactionApiEnabled =
+    isUpdateDocumentShardKeyUsingTransactionApiEnabled(st.s);
+
+// Set up a sharded collection with two shards:
+// shard0: [MinKey, 0]
+// shard1: [0, MaxKey]
+assert.commandWorked(st.s.adminCommand({enableSharding: dbName}));
+assert.commandWorked(st.s.adminCommand({shardCollection: ns, key: {x: 1}}));
+assert.commandWorked(st.s.adminCommand({split: ns, middle: {x: 0}}));
+assert.commandWorked(
+    st.s.adminCommand({moveChunk: ns, find: {x: MinKey}, to: st.shard0.shardName}));
+assert.commandWorked(st.s.adminCommand({moveChunk: ns, find: {x: 1}, to: st.shard1.shardName}));
+
+/**
+ * Asserts that when the updateDocumentShardKeyUsingTransactionApi feature flag is enabled, ordinary
+ * updates that modify a document's shard key will complete.
+ */
+function runAndVerifyCommandChangingOwningShard(cmdObj, expectedUpdatedDoc) {
+    const updateRes = db.runCommand(cmdObj);
+    const numDocsUpdated = testColl.find(expectedUpdatedDoc).itcount();
+
+    if (updateDocumentShardKeyUsingTransactionApiEnabled) {
+        assert.commandWorked(updateRes);
+        assert.eq(1, numDocsUpdated);
+    } else {
+        assert.commandFailedWithCode(updateRes, ErrorCodes.IllegalOperation);
+        assert.eq(0, numDocsUpdated);
+    }
+}
+
+// This test is meant to test updating document shard keys without a session.
+TestData.disableImplicitSessions = true;
+
+assert.commandWorked(db.runCommand({insert: collName, documents: [{x: -1}]}));
+
+/**
+ * Verify that a client with the feature flag updateDocumentShardKeyUsingTransactionApi enabled can
+ * update a shard key that change's the document's owning shard without a session.
+ * Test 1: update()
+ * Test 2: findAndModify()
+ */
+(() => {
+    assert.commandWorked(db.runCommand({insert: collName, documents: [{x: -1}]}));
+    let cmdObj = {update: collName, updates: [{q: {x: -1}, u: {"$set": {x: 1}}, upsert: false}]};
+    runAndVerifyCommandChangingOwningShard(cmdObj, {x: 1});
+
+    assert.commandWorked(db.runCommand({insert: collName, documents: [{x: -2}]}));
+    cmdObj = {findAndModify: collName, query: {x: -2}, update: {"$set": {x: 2}}, upsert: false};
+    runAndVerifyCommandChangingOwningShard(cmdObj, {x: 2});
+})();
+
+/**
+ * Verify that a client can update a document's shard key without a session, where the update
+ * upserts a document.
+ * Test 1: update() with upsert: true
+ * Test 2: findAndModify() with upsert: true
+ */
+(() => {
+    let cmdObj = {update: collName, updates: [{q: {x: -3}, u: {"$set": {x: 3}}, upsert: true}]};
+    assert.eq(0, testColl.find({x: -3}).itcount());
+    runAndVerifyCommandChangingOwningShard(cmdObj, {x: 3});
+
+    cmdObj = {findAndModify: collName, query: {x: -4}, update: {"$set": {x: 4}}, upsert: true};
+    assert.eq(0, testColl.find({x: -4}).itcount());
+    runAndVerifyCommandChangingOwningShard(cmdObj, {x: 4});
+})();
+
+/**
+ * Verify that a client update a document's shard key inside a session, that would cause an existing
+ * doc to move owning shard.
+ * Test 1: update() inside session
+ * Test 2: findAndModify() inside session
+ */
+const lsid = ({id: UUID()});
+(() => {
+    assert.commandWorked(db.runCommand({insert: collName, documents: [{x: -5}]}));
+    let cmdObj = {
+        update: collName,
+        updates: [{q: {x: -5}, u: {"$set": {x: 5}}, upsert: false}],
+        lsid: lsid
+    };
+    runAndVerifyCommandChangingOwningShard(cmdObj, {x: 5});
+
+    assert.commandWorked(db.runCommand({insert: collName, documents: [{x: -6}]}));
+    cmdObj = {
+        findAndModify: collName,
+        query: {x: -6},
+        update: {"$set": {x: 6}},
+        upsert: false,
+        lsid: lsid
+    };
+    runAndVerifyCommandChangingOwningShard(cmdObj, {x: 6});
+})();
+
+/**
+ * Verify that a client can update a document's shard key inside a session with an update
+ * that would cause the document to move from shard0 to shard1, where the update upserts a document.
+ * Test 1: update() with upsert: true inside session
+ * Test 2: findAndModify() with upsert: inside session
+ */
+(() => {
+    let cmdObj = {
+        update: collName,
+        updates: [{q: {x: -7}, u: {"$set": {x: 7}}, upsert: true}],
+        lsid: lsid
+    };
+    assert.eq(0, testColl.find({x: -7}).itcount());
+    runAndVerifyCommandChangingOwningShard(cmdObj, {x: 7});
+
+    cmdObj = {
+        findAndModify: collName,
+        query: {x: -8},
+        update: {"$set": {x: 8}},
+        upsert: true,
+        lsid: lsid
+    };
+    assert.eq(0, testColl.find({x: -8}).itcount());
+    runAndVerifyCommandChangingOwningShard(cmdObj, {x: 8});
+})();
+
+st.stop();
+})();
diff --git a/jstests/sharding/update_shard_key_doc_on_same_shard.js b/jstests/sharding/update_shard_key_doc_on_same_shard.js
index 883c4a63d7c..0bf2e6682bc 100644
--- a/jstests/sharding/update_shard_key_doc_on_same_shard.js
+++ b/jstests/sharding/update_shard_key_doc_on_same_shard.js
@@ -29,26 +29,6 @@ enableCoordinateCommitReturnImmediatelyAfterPersistingDecision(st);
 assert.commandWorked(mongos.adminCommand({enableSharding: kDbName}));
 st.ensurePrimaryShard(kDbName, shard0);
 
-// -----------------------------------------
-// Updates to the shard key are not allowed if write is not retryable and not in a multi-stmt
-// txn
-// -----------------------------------------
-
-let docsToInsert = [{"x": 4, "a": 3}, {"x": 100}, {"x": 300, "a": 3}, {"x": 500, "a": 6}];
-shardCollectionMoveChunks(st, kDbName, ns, {"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
-
-assert.writeError(mongos.getDB(kDbName).foo.update({"x": 300}, {"x": 600}));
-assert.eq(1, mongos.getDB(kDbName).foo.find({"x": 300}).itcount());
-assert.eq(0, mongos.getDB(kDbName).foo.find({"x": 600}).itcount());
-
-assert.throws(function() {
-    mongos.getDB(kDbName).foo.findAndModify({query: {"x": 300}, update: {$set: {"x": 600}}});
-});
-assert.eq(1, mongos.getDB(kDbName).foo.find({"x": 300}).itcount());
-assert.eq(0, mongos.getDB(kDbName).foo.find({"x": 600}).itcount());
-
-mongos.getDB(kDbName).foo.drop();
-
 // ---------------------------------
 // Update shard key retryable write
 // ---------------------------------
@@ -794,7 +774,7 @@ assertCanUpdateInBulkOpWhenDocsRemainOnSameShard(st, kDbName, ns, session, sessi
 assertCanUpdateInBulkOpWhenDocsRemainOnSameShard(st, kDbName, ns, session, sessionDB, true, true);
 
 // Update two docs, updating one twice
-docsToInsert = [{"x": 4, "a": 3}, {"x": 100}, {"x": 300, "a": 3}, {"x": 500, "a": 6}];
+let docsToInsert = [{"x": 4, "a": 3}, {"x": 100}, {"x": 300, "a": 3}, {"x": 500, "a": 6}];
 shardCollectionMoveChunks(st, kDbName, ns, {"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
 
 session.startTransaction();
diff --git a/src/mongo/db/SConscript b/src/mongo/db/SConscript
index e1bbc943142..26950156e51 100644
--- a/src/mongo/db/SConscript
+++ b/src/mongo/db/SConscript
@@ -1484,6 +1484,7 @@ env.Library(
         '$BUILD_DIR/mongo/db/catalog/local_oplog_info',
         '$BUILD_DIR/mongo/db/commands/server_status_core',
         '$BUILD_DIR/mongo/db/concurrency/exception_util',
+        '$BUILD_DIR/mongo/db/internal_transactions_feature_flag',
         '$BUILD_DIR/mongo/db/stats/resource_consumption_metrics',
         '$BUILD_DIR/mongo/db/storage/record_store_base',
         '$BUILD_DIR/mongo/db/timeseries/timeseries_options',
diff --git a/src/mongo/db/exec/update_stage.cpp b/src/mongo/db/exec/update_stage.cpp
index bd559549ac7..04d7e35bc25 100644
--- a/src/mongo/db/exec/update_stage.cpp
+++ b/src/mongo/db/exec/update_stage.cpp
@@ -43,6 +43,7 @@
 #include "mongo/db/exec/shard_filterer_impl.h"
 #include "mongo/db/exec/working_set_common.h"
 #include "mongo/db/exec/write_stage_common.h"
+#include "mongo/db/internal_transactions_feature_flag_gen.h"
 #include "mongo/db/op_observer/op_observer.h"
 #include "mongo/db/query/collection_query_info.h"
 #include "mongo/db/query/explain.h"
@@ -718,11 +719,18 @@ void UpdateStage::_checkRestrictionsOnUpdatingShardKeyAreNotViolated(
 
     // If this node is a replica set primary node, an attempted update to the shard key value must
     // either be a retryable write or inside a transaction.
+    // An update without a transaction number is legal if
+    // gFeatureFlagUpdateDocumentShardKeyUsingTransactionApi is enabled because mongos
+    // will be able to start an internal transaction to handle the wouldChangeOwningShard error
+    // thrown below.
     // If this node is a replica set secondary node, we can skip validation.
-    uassert(ErrorCodes::IllegalOperation,
-            "Must run update to shard key field in a multi-statement transaction or with "
-            "retryWrites: true.",
-            opCtx()->getTxnNumber() || !opCtx()->writesAreReplicated());
+    if (!feature_flags::gFeatureFlagUpdateDocumentShardKeyUsingTransactionApi.isEnabled(
+            serverGlobalParams.featureCompatibility)) {
+        uassert(ErrorCodes::IllegalOperation,
+                "Must run update to shard key field in a multi-statement transaction or with "
+                "retryWrites: true.",
+                opCtx()->getTxnNumber());
+    }
 }
 
 
diff --git a/src/mongo/db/exec/upsert_stage.cpp b/src/mongo/db/exec/upsert_stage.cpp
index 06235a9ac52..a8afd618fd1 100644
--- a/src/mongo/db/exec/upsert_stage.cpp
+++ b/src/mongo/db/exec/upsert_stage.cpp
@@ -33,6 +33,7 @@
 #include "mongo/db/catalog/local_oplog_info.h"
 #include "mongo/db/concurrency/exception_util.h"
 #include "mongo/db/curop_failpoint_helpers.h"
+#include "mongo/db/internal_transactions_feature_flag_gen.h"
 #include "mongo/db/query/query_feature_flags_gen.h"
 #include "mongo/db/s/collection_sharding_state.h"
 #include "mongo/db/s/operation_sharding_state.h"
@@ -138,12 +139,20 @@ void UpsertStage::_performInsert(BSONObj newDocument) {
             if (!collFilter.keyBelongsToMe(newShardKey)) {
                 // An attempt to upsert a document with a shard key value that belongs on another
                 // shard must either be a retryable write or inside a transaction.
-                uassert(ErrorCodes::IllegalOperation,
+                // An upsert without a transaction number is legal if
+                // gFeatureFlagUpdateDocumentShardKeyUsingTransactionApi is enabled because mongos
+                // will be able to start an internal transaction to handle the
+                // wouldChangeOwningShard error thrown below.
+                if (!feature_flags::gFeatureFlagUpdateDocumentShardKeyUsingTransactionApi.isEnabled(
+                        serverGlobalParams.featureCompatibility)) {
+                    uassert(
+                        ErrorCodes::IllegalOperation,
                         "The upsert document could not be inserted onto the shard targeted by the "
                         "query, since its shard key belongs on a different shard. Cross-shard "
                         "upserts are only allowed when running in a transaction or with "
                         "retryWrites: true.",
                         opCtx()->getTxnNumber());
+                }
                 uasserted(WouldChangeOwningShardInfo(_params.request->getQuery(),
                                                      newDocument,
                                                      true /* upsert */,
diff --git a/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp b/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp
index bcd2d5a746f..657e1e9086d 100644
--- a/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp
+++ b/src/mongo/s/commands/cluster_find_and_modify_cmd.cpp
@@ -144,7 +144,7 @@ BSONObj getShardKey(OperationContext* opCtx,
     return shardKey;
 }
 
-void handleWouldChangeOwningShardErrorRetryableWrite(
+void handleWouldChangeOwningShardErrorNonTransaction(
     OperationContext* opCtx,
     const ShardId& shardId,
     const NamespaceString& nss,
@@ -519,7 +519,8 @@ private:
                             const NamespaceString& nss,
                             const BSONObj& cmdObj,
                             BSONObjBuilder* result) {
-        bool isRetryableWrite = opCtx->getTxnNumber() && !TransactionRouter::get(opCtx);
+        auto txnRouter = TransactionRouter::get(opCtx);
+        bool isRetryableWrite = opCtx->getTxnNumber() && !txnRouter;
 
         const auto response = [&] {
             std::vector<AsyncRequestsSender::Request> requests;
@@ -574,13 +575,15 @@ private:
                 // Strip runtime constants because they will be added again when this command is
                 // recursively sent through the service entry point.
                 parsedRequest.setLegacyRuntimeConstants(boost::none);
-                if (isRetryableWrite) {
-                    parsedRequest.setStmtId(0);
-                    handleWouldChangeOwningShardErrorRetryableWrite(
-                        opCtx, shardId, nss, parsedRequest, result);
-                } else {
+                if (txnRouter) {
                     handleWouldChangeOwningShardErrorTransaction(
                         opCtx, nss, responseStatus, parsedRequest, result);
+                } else {
+                    if (isRetryableWrite) {
+                        parsedRequest.setStmtId(0);
+                    }
+                    handleWouldChangeOwningShardErrorNonTransaction(
+                        opCtx, shardId, nss, parsedRequest, result);
                 }
             } else {
                 // TODO SERVER-67429: Remove this branch.
diff --git a/src/mongo/s/commands/cluster_write_cmd.cpp b/src/mongo/s/commands/cluster_write_cmd.cpp
index a43e919b217..358615c1993 100644
--- a/src/mongo/s/commands/cluster_write_cmd.cpp
+++ b/src/mongo/s/commands/cluster_write_cmd.cpp
@@ -144,7 +144,7 @@ boost::optional<WouldChangeOwningShardInfo> getWouldChangeOwningShardErrorInfo(
     }
 }
 
-void handleWouldChangeOwningShardErrorRetryableWrite(OperationContext* opCtx,
+void handleWouldChangeOwningShardErrorNonTransaction(OperationContext* opCtx,
                                                      BatchedCommandRequest* request,
                                                      BatchedCommandResponse* response) {
     // Strip write concern because this command will be sent as part of a
@@ -294,18 +294,24 @@ bool handleWouldChangeOwningShardError(OperationContext* opCtx,
 
     if (feature_flags::gFeatureFlagUpdateDocumentShardKeyUsingTransactionApi.isEnabled(
             serverGlobalParams.featureCompatibility)) {
-        if (isRetryableWrite) {
-            if (MONGO_unlikely(hangAfterThrowWouldChangeOwningShardRetryableWrite.shouldFail())) {
-                LOGV2(5918603, "Hit hangAfterThrowWouldChangeOwningShardRetryableWrite failpoint");
-                hangAfterThrowWouldChangeOwningShardRetryableWrite.pauseWhileSet(opCtx);
-            }
-
-            handleWouldChangeOwningShardErrorRetryableWrite(opCtx, request, response);
-        } else {
+        if (txnRouter) {
             auto updateResult = handleWouldChangeOwningShardErrorTransaction(
                 opCtx, request, response, *wouldChangeOwningShardErrorInfo);
             updatedShardKey = updateResult.updatedShardKey;
             upsertedId = std::move(updateResult.upsertedId);
+        } else {
+            // Updating a document's shard key such that its owning shard changes must be run in a
+            // transaction. If this update is not already in a transaction, complete the update
+            // using an internal transaction.
+            if (isRetryableWrite) {
+                if (MONGO_unlikely(
+                        hangAfterThrowWouldChangeOwningShardRetryableWrite.shouldFail())) {
+                    LOGV2(5918603,
+                          "Hit hangAfterThrowWouldChangeOwningShardRetryableWrite failpoint");
+                    hangAfterThrowWouldChangeOwningShardRetryableWrite.pauseWhileSet(opCtx);
+                }
+            }
+            handleWouldChangeOwningShardErrorNonTransaction(opCtx, request, response);
         }
     } else {
         // TODO SERVER-67429: Delete this branch.