SERVER-39158 Change updates to prefer target by the query

16 files changed, 730 insertions, 138 deletions

diff --git a/jstests/core/bulk_legacy_enforce_gle.js b/jstests/core/bulk_legacy_enforce_gle.js
index 21134a7270d..88b7c51e758 100644
--- a/jstests/core/bulk_legacy_enforce_gle.js
+++ b/jstests/core/bulk_legacy_enforce_gle.js
@@ -24,7 +24,7 @@
 
     coll.drop();
     let bulk = coll.initializeUnorderedBulkOp();
-    bulk.find({none: 1}).upsert().updateOne({_id: 1});
+    bulk.find({_id: 1}).upsert().updateOne({_id: 1});
     assert.writeOK(bulk.execute());
     let gle = assert.gleOK(db.runCommand({getLastError: 1}));
     assert.eq(1, gle.n, tojson(gle));
@@ -48,7 +48,7 @@
     bulk = coll.initializeUnorderedBulkOp();
     bulk.find({none: 1}).upsert().updateOne({_id: 1});
     bulk.find({none: 1}).upsert().updateOne({_id: 1});
-    bulk.find({none: 1}).upsert().updateOne({_id: 0});
+    bulk.find({_id: 0}).upsert().updateOne({_id: 0});
     let res = assert.throws(function() {
         bulk.execute();
     });
@@ -62,9 +62,9 @@
     assert(coll.drop());
     insertDocument({_id: 1});
     bulk = coll.initializeUnorderedBulkOp();
-    bulk.find({none: 1}).upsert().updateOne({_id: 0});
+    bulk.find({_id: 0}).upsert().updateOne({_id: 0});
     bulk.find({none: 1}).upsert().updateOne({_id: 1});
-    bulk.find({none: 1}).upsert().updateOne({_id: 2});
+    bulk.find({_id: 2}).upsert().updateOne({_id: 2});
     res = assert.throws(function() {
         bulk.execute();
     });
@@ -79,8 +79,8 @@
     assert(coll.drop());
     insertDocument({_id: 2});
     bulk = coll.initializeUnorderedBulkOp();
-    bulk.find({none: 1}).upsert().updateOne({_id: 0});
-    bulk.find({none: 1}).upsert().updateOne({_id: 1});
+    bulk.find({_id: 0}).upsert().updateOne({_id: 0});
+    bulk.find({_id: 1}).upsert().updateOne({_id: 1});
     bulk.find({none: 1}).upsert().updateOne({_id: 2});
     res = assert.throws(function() {
         bulk.execute();
@@ -95,8 +95,8 @@
     assert(coll.drop());
     insertDocument({_id: 2});
     bulk = coll.initializeUnorderedBulkOp();
-    bulk.find({none: 1}).upsert().updateOne({_id: 0});
-    bulk.find({none: 1}).upsert().updateOne({_id: 1});
+    bulk.find({_id: 0}).upsert().updateOne({_id: 0});
+    bulk.find({_id: 1}).upsert().updateOne({_id: 1});
     bulk.find({none: 1}).upsert().updateOne({_id: 2});
     res = assert.throws(function() {
         bulk.execute();
@@ -111,8 +111,8 @@
     assert(coll.drop());
     insertDocument({_id: 2});
     bulk = coll.initializeUnorderedBulkOp();
-    bulk.find({none: 1}).upsert().updateOne({_id: 0});
-    bulk.find({none: 1}).upsert().updateOne({_id: 1});
+    bulk.find({_id: 0}).upsert().updateOne({_id: 0});
+    bulk.find({_id: 1}).upsert().updateOne({_id: 1});
     bulk.find({none: 1}).upsert().updateOne({_id: 2});
     res = assert.throws(function() {
         bulk.execute();
diff --git a/jstests/core/insert_illegal_doc.js b/jstests/core/insert_illegal_doc.js
index 31246bef7f5..d91866e8766 100644
--- a/jstests/core/insert_illegal_doc.js
+++ b/jstests/core/insert_illegal_doc.js
@@ -6,7 +6,7 @@ coll.ensureIndex({a: 1, b: 1});
 
 var res;
 // test upsert
-res = coll.update({}, {_id: 1, a: [1, 2, 3], b: [4, 5, 6]}, true);
+res = coll.update({_id: 1}, {_id: 1, a: [1, 2, 3], b: [4, 5, 6]}, true);
 assert.writeError(res);
 assert.eq(res.getWriteError().code, ErrorCodes.CannotIndexParallelArrays);
 assert.eq(0, coll.find().itcount(), "should not be a doc");
diff --git a/jstests/core/max_doc_size.js b/jstests/core/max_doc_size.js
index 3ac7149f485..775121a5c9f 100644
--- a/jstests/core/max_doc_size.js
+++ b/jstests/core/max_doc_size.js
@@ -24,15 +24,16 @@
     assert.eq(coll.find({}).itcount(), 1);
 
     coll.drop();
+    const objectId = new ObjectId();
     assert.commandWorked(db.runCommand({
         update: coll.getName(),
         ordered: true,
-        updates: [{q: {a: 1}, u: {_id: new ObjectId(), x: maxStr}, upsert: true}]
+        updates: [{q: {_id: objectId}, u: {_id: objectId, x: maxStr}, upsert: true}]
     }));
     assert.eq(coll.find({}).itcount(), 1);
 
     coll.drop();
-    const objectId = new ObjectId();
+
     assert.commandWorked(coll.insert({_id: objectId}));
     assert.commandWorked(db.runCommand({
         update: coll.getName(),
@@ -56,7 +57,7 @@
     assert.commandFailedWithCode(db.runCommand({
         update: coll.getName(),
         ordered: true,
-        updates: [{q: {a: 1}, u: {_id: new ObjectId(), x: largerThanMaxString}, upsert: true}]
+        updates: [{q: {_id: objectId}, u: {_id: objectId, x: largerThanMaxString}, upsert: true}]
     }),
                                  17420);
 
diff --git a/jstests/multiVersion/update_shard_key_disallowed_fcv40.js b/jstests/multiVersion/update_shard_key_disallowed_fcv40.js
index 7c4ed6751b5..7bf01243d5f 100644
--- a/jstests/multiVersion/update_shard_key_disallowed_fcv40.js
+++ b/jstests/multiVersion/update_shard_key_disallowed_fcv40.js
@@ -64,9 +64,8 @@
         // Assert that updating the shard key when the doc would move shards fails for both modify
         // and replacement updates.
         assert.writeError(sessionDB.foo.update({x: 80}, {$set: {x: 3}}));
-        // TODO: SERVER-39158. Currently, this update will not fail but will not update the doc.
-        // After SERVER-39158 is finished, this should fail.
-        assert.commandWorked(sessionDB.foo.update({x: 80}, {x: 3}));
+        assert.commandFailedWithCode(sessionDB.foo.update({x: 80}, {x: 3}),
+                                     [ErrorCodes.ImmutableField, ErrorCodes.InvalidOptions]);
         assert.eq(1, mongos.getDB(kDbName).foo.find({x: 80}).itcount());
         assert.eq(0, mongos.getDB(kDbName).foo.find({x: 3}).itcount());
 
diff --git a/jstests/sharding/count1.js b/jstests/sharding/count1.js
index d3b2cf63b3e..2275faed656 100644
--- a/jstests/sharding/count1.js
+++ b/jstests/sharding/count1.js
@@ -38,12 +38,12 @@
 
     assert.eq(1, s.config.chunks.count({"ns": "test.foo"}), "sanity check A");
 
-    db.foo.save({_id: 1, name: "eliot"});
-    db.foo.save({_id: 2, name: "sara"});
-    db.foo.save({_id: 3, name: "bob"});
-    db.foo.save({_id: 4, name: "joe"});
-    db.foo.save({_id: 5, name: "mark"});
-    db.foo.save({_id: 6, name: "allan"});
+    assert.commandWorked(db.foo.insert({_id: 1, name: "eliot"}));
+    assert.commandWorked(db.foo.insert({_id: 2, name: "sara"}));
+    assert.commandWorked(db.foo.insert({_id: 3, name: "bob"}));
+    assert.commandWorked(db.foo.insert({_id: 4, name: "joe"}));
+    assert.commandWorked(db.foo.insert({_id: 5, name: "mark"}));
+    assert.commandWorked(db.foo.insert({_id: 6, name: "allan"}));
 
     assert.eq(6, db.foo.find().count(), "basic count");
 
@@ -153,7 +153,7 @@
 
     // part 6
     for (var i = 0; i < 10; i++) {
-        db.foo.save({_id: 7 + i, name: "zzz" + i});
+        assert.commandWorked(db.foo.insert({_id: 7 + i, name: "zzz" + i}));
     }
 
     assert.eq(10, db.foo.find({name: {$gt: "z"}}).itcount(), "LSF1");
diff --git a/jstests/sharding/exact_shard_key_target.js b/jstests/sharding/exact_shard_key_target.js
index 26d7c2402d1..7ff31a97dda 100644
--- a/jstests/sharding/exact_shard_key_target.js
+++ b/jstests/sharding/exact_shard_key_target.js
@@ -48,15 +48,6 @@ assert.eq(1,
               st.shard1.getCollection(coll.toString()).count({updated: true}));
 
 //
-// Successive upserts (save()-style)
-coll.remove({});
-assert.writeOK(coll.update({_id: 1}, {_id: 1, a: {b: 1}}, {upsert: true}));
-assert.writeOK(coll.update({_id: 1}, {_id: 1, a: {b: 1}}, {upsert: true}));
-assert.eq(1,
-          st.shard0.getCollection(coll.toString()).count() +
-              st.shard1.getCollection(coll.toString()).count());
-
-//
 // Successive upserts (replacement-style)
 coll.remove({});
 assert.writeOK(coll.update({a: {b: 1}}, {a: {b: 1}}, {upsert: true}));
diff --git a/jstests/sharding/libs/update_shard_key_helpers.js b/jstests/sharding/libs/update_shard_key_helpers.js
index ce8ab5023aa..1cc792ed365 100644
--- a/jstests/sharding/libs/update_shard_key_helpers.js
+++ b/jstests/sharding/libs/update_shard_key_helpers.js
@@ -154,7 +154,7 @@ function runUpdateCmdFail(
 
     let updatedVal = update["$set"] ? update["$set"] : update;
     assert.eq(1, st.s.getDB(kDbName).foo.find(query).itcount());
-    if (!update["$unset"]) {
+    if (!update["$unset"] && Object.keys(update).length != 0) {
         assert.eq(0, st.s.getDB(kDbName).foo.find(updatedVal).itcount());
     }
 }
@@ -173,7 +173,7 @@ function runFindAndModifyCmdFail(st, kDbName, session, sessionDB, inTxn, query,
     }
     let updatedVal = update["$set"] ? update["$set"] : update;
     assert.eq(1, st.s.getDB(kDbName).foo.find(query).itcount());
-    if (!update["$unset"]) {
+    if (!update["$unset"] && Object.keys(update).length != 0) {
         assert.eq(0, st.s.getDB(kDbName).foo.find(updatedVal).itcount());
     }
 }
@@ -686,23 +686,23 @@ function assertCanUpdatePrimitiveShardKeyHashedChangeShards(
     assert.gte(docsOnShard0.length, 2);
     assert.gte(docsOnShard1.length, 2);
 
-    // TODO SERVER-39158: Add replacement tests as well. Can change second test to be replacement
-    // rather than modify.
-
     // Since we now know that the value of x in each of docsOnShard0[1] and docsOnShard1[1] will be
     // hashed to map to shard0 and shard1 respectively, we can delete these documents and then use
     // them as values to change the shard key to.
     st.s.getDB(kDbName).foo.remove({"x": docsOnShard0[1].x});
     st.s.getDB(kDbName).foo.remove({"x": docsOnShard1[1].x});
     let queries = [{"x": docsOnShard0[0].x}, {"x": docsOnShard1[0].x}];
-    let updates = [{"$set": {"x": docsOnShard1[1].x}}, {"$set": {"x": docsOnShard0[1].x}}];
+    let updates = [{"$set": {"x": docsOnShard1[1].x}}, {"x": docsOnShard0[1].x}];
 
     // Non-upsert case. The first update will move a doc from shard0 to shard1 and the second will
     // move a doc from shard1 to shard 0.
     let upsert = false;
+
+    // Op-style modify
     assertUpdateSucceeds(st, session, sessionDB, inTxn, queries[0], updates[0], upsert);
     assertHashedShardKeyUpdateCorrect(st, kDbName, queries[0], updates[0], upsert, false);
 
+    // Replacement style modify
     assertUpdateSucceeds(st, session, sessionDB, inTxn, queries[1], updates[1], upsert);
     assertHashedShardKeyUpdateCorrect(st, kDbName, queries[1], updates[1], upsert, true);
 
@@ -714,8 +714,12 @@ function assertCanUpdatePrimitiveShardKeyHashedChangeShards(
     st.s.getDB(kDbName).foo.remove({"x": docsOnShard1[1].x});
 
     upsert = true;
+
+    // Op-style upsert
     assertUpdateSucceeds(st, session, sessionDB, inTxn, queries[0], updates[0], upsert);
     assertHashedShardKeyUpdateCorrect(st, kDbName, queries[0], updates[0], upsert, false);
+
+    // Modify style upsert
     assertUpdateSucceeds(st, session, sessionDB, inTxn, queries[1], updates[1], upsert);
     assertHashedShardKeyUpdateCorrect(st, kDbName, queries[1], updates[1], upsert, true);
 
diff --git a/jstests/sharding/regex_targeting.js b/jstests/sharding/regex_targeting.js
index e55e0f6cab9..e2300a3e896 100644
--- a/jstests/sharding/regex_targeting.js
+++ b/jstests/sharding/regex_targeting.js
@@ -152,16 +152,48 @@
     assert.writeOK(collHashed.update({hash: /abcde.*/}, {$set: {updated: true}}, {multi: true}));
     assert.eq(collHashed.find().itcount(), collHashed.find({updated: true}).itcount());
 
+    collSharded.remove({});
+    collCompound.remove({});
+    collNested.remove({});
+
+    //
+    //
+    // Op-style updates with regex should fail on sharded collections.
+    // Query clause is targeted, and regex in query clause is ambiguous.
+    assert.commandFailedWithCode(
+        collSharded.update({a: /abcde-1/}, {"$set": {b: 1}}, {upsert: false}),
+        ErrorCodes.InvalidOptions);
+    assert.commandFailedWithCode(
+        collSharded.update({a: /abcde-[1-2]/}, {"$set": {b: 1}}, {upsert: false}),
+        ErrorCodes.InvalidOptions);
+    assert.commandFailedWithCode(
+        collNested.update({a: {b: /abcde-1/}}, {"$set": {"a.b": /abcde-1/, b: 1}}, {upsert: false}),
+        ErrorCodes.InvalidOptions);
+    assert.commandFailedWithCode(
+        collNested.update({"a.b": /abcde.*/}, {"$set": {b: 1}}, {upsert: false}),
+        ErrorCodes.InvalidOptions);
+
+    //
+    //
+    // Replacement style updates with regex should work on sharded collections.
+    // If query clause is ambiguous, we fallback to using update clause for targeting.
+    assert.commandWorked(collSharded.update({a: /abcde.*/}, {a: /abcde.*/, b: 1}, {upsert: false}));
+    assert.commandWorked(collSharded.update({a: /abcde-1/}, {a: /abcde-1/, b: 1}, {upsert: false}));
+    assert.commandWorked(
+        collNested.update({a: {b: /abcde.*/}}, {a: {b: /abcde.*/}}, {upsert: false}));
+    assert.commandWorked(
+        collNested.update({'a.b': /abcde-1/}, {a: {b: /abcde.*/}}, {upsert: false}));
+
     //
     //
     // Upsert with op-style regex should fail on sharded collections
     // Query clause is targeted, and regex in query clause is ambiguous
-    collSharded.remove({});
-    collCompound.remove({});
-    collNested.remove({});
+
+    // The queries will also be interpreted as regex based prefix search and cannot target a single
+    // shard.
     assert.writeError(collSharded.update({a: /abcde.*/}, {$set: {a: /abcde.*/}}, {upsert: true}));
     assert.writeError(
-        collCompound.update({a: /abcde.*/}, {$set: {a: /abcde.*/, b: 1}}, {upsert: true}));
+        collCompound.update({a: /abcde-1/}, {$set: {a: /abcde.*/, b: 1}}, {upsert: true}));
     // Exact regex in query never equality
     assert.writeError(
         collNested.update({'a.b': /abcde.*/}, {$set: {'a.b': /abcde.*/}}, {upsert: true}));
@@ -172,16 +204,21 @@
 
     //
     //
-    // Upsert by replacement-style regex should succeed on sharded collections
-    // Replacement clause is targeted, and regex is unambiguously a value
-    collSharded.remove({});
-    collCompound.remove({});
-    collNested.remove({});
-    assert.writeOK(collSharded.update({a: /abcde.*/}, {a: /abcde.*/}, {upsert: true}));
-    assert.writeOK(collCompound.update({a: /abcde.*/}, {a: /abcde.*/, b: 1}, {upsert: true}));
-    assert.writeOK(collNested.update({'a.b': /abcde.*/}, {a: {b: /abcde.*/}}, {upsert: true}));
-    assert.writeOK(collNested.update({a: {b: /abcde.*/}}, {a: {b: /abcde.*/}}, {upsert: true}));
-    assert.writeOK(collNested.update({c: 1}, {a: {b: /abcde.*/}}, {upsert: true}));
+    // Upsert by replacement-style regex should fail on sharded collections
+    // Query clause is targeted, and regex in query clause is ambiguous
+    assert.commandFailedWithCode(collSharded.update({a: /abcde.*/}, {a: /abcde.*/}, {upsert: true}),
+                                 ErrorCodes.ShardKeyNotFound);
+    assert.commandFailedWithCode(
+        collCompound.update({a: /abcde.*/}, {a: /abcde.*/, b: 1}, {upsert: true}),
+        ErrorCodes.ShardKeyNotFound);
+    assert.commandFailedWithCode(
+        collNested.update({'a.b': /abcde-1/}, {a: {b: /abcde.*/}}, {upsert: true}),
+        ErrorCodes.ShardKeyNotFound);
+    assert.commandFailedWithCode(
+        collNested.update({a: {b: /abcde.*/}}, {a: {b: /abcde.*/}}, {upsert: true}),
+        ErrorCodes.ShardKeyNotFound);
+    assert.commandFailedWithCode(collNested.update({c: 1}, {a: {b: /abcde.*/}}, {upsert: true}),
+                                 ErrorCodes.ShardKeyNotFound);
 
     //
     //
diff --git a/jstests/sharding/shard6.js b/jstests/sharding/shard6.js
index 5fa07ea1301..e00833ac179 100644
--- a/jstests/sharding/shard6.js
+++ b/jstests/sharding/shard6.js
@@ -103,14 +103,11 @@
 
     poolStats("after checking itcount");
 
-    // --- test save support ---
+    // --- Verify that modify & save style updates doesn't work on sharded clusters ---
 
     var o = db.data.findOne();
     o.x = 16;
-    db.data.save(o);
-    o = db.data.findOne({_id: o._id});
-    assert.eq(16, o.x, "x1 - did save fail? " + tojson(o));
-
+    assert.commandFailedWithCode(db.data.save(o), ErrorCodes.ShardKeyNotFound);
     poolStats("at end");
 
     print(summary);
diff --git a/jstests/sharding/update_compound_shard_key.js b/jstests/sharding/update_compound_shard_key.js
new file mode 100644
index 00000000000..fca04410c52
--- /dev/null
+++ b/jstests/sharding/update_compound_shard_key.js
@@ -0,0 +1,418 @@
+/**
+ * Tests to validate the functionality of update command in the presence of a compound shard key.
+ * @tags: [uses_transactions, uses_multi_shard_transaction]
+ */
+(function() {
+    'use strict';
+
+    load("jstests/sharding/libs/update_shard_key_helpers.js");
+
+    const st = new ShardingTest({mongos: 1, shards: 3});
+    const kDbName = 'update_compound_sk';
+    const ns = kDbName + '.coll';
+    const session = st.s.startSession();
+    const sessionDB = session.getDatabase(kDbName);
+
+    assert.commandWorked(st.s0.adminCommand({enableSharding: kDbName}));
+    st.ensurePrimaryShard(kDbName, st.shard0.shardName);
+
+    assert.commandWorked(
+        st.s.getDB('config').adminCommand({shardCollection: ns, key: {x: 1, y: 1, z: 1}}));
+
+    let docsToInsert = [
+        {_id: 0, x: 4, y: 3, z: 3},
+        {_id: 1, x: 100, y: 50, z: 3, a: 5},
+        {_id: 2, x: 100, y: 500, z: 3, a: 5}
+    ];
+
+    // Make sure that shard0, shard1 and shard2 has _id 0,1 and 2 documents respectively.
+    assert.commandWorked(st.s.adminCommand({split: ns, middle: {x: 100, y: 0, z: 3}}));
+    assert.commandWorked(st.s.adminCommand({split: ns, middle: {x: 100, y: 100, z: 3}}));
+
+    for (let i = 0; i < docsToInsert.length; i++) {
+        assert.commandWorked(st.s.getDB(kDbName).coll.insert(docsToInsert[i]));
+    }
+
+    assert.commandWorked(
+        st.s.adminCommand({moveChunk: ns, find: {x: 100, y: 50, z: 3}, to: st.shard1.shardName}));
+    assert.commandWorked(
+        st.s.adminCommand({moveChunk: ns, find: {x: 100, y: 500, z: 3}, to: st.shard2.shardName}));
+    cleanupOrphanedDocs(st, ns);
+
+    function assertUpdateWorked(query, update, isUpsert, _id) {
+        const res = st.s.getDB(kDbName).coll.update(query, update, {upsert: isUpsert});
+        assert.commandWorked(res);
+        assert.eq(0, res.nUpserted);
+        assert.eq(1, res.nMatched);
+        assert.eq(1, res.nModified);
+
+        // Skip find based validation for pipleline update.
+        if (!Array.isArray(update)) {
+            if (update["$set"] != undefined) {
+                update = update["$set"];
+            }
+            update["_id"] = _id;
+            // Make sure that the update modified the document with the given _id.
+            assert.eq(1, st.s.getDB(kDbName).coll.find(update).itcount());
+        }
+    }
+
+    /**
+     * For upserts this will insert a new document, for non-upserts it will be a no-op.
+     */
+    function assertUpdateWorkedWithNoMatchingDoc(query, update, isUpsert, inTransaction) {
+        const res = sessionDB.coll.update(query, update, {upsert: isUpsert});
+
+        assert.commandWorked(res);
+        assert.eq(isUpsert ? 1 : 0, res.nUpserted);
+        assert.eq(0, res.nMatched);
+        assert.eq(0, res.nModified);
+
+        // Skip find based validation for pipleline update or when inside a transaction.
+        if (Array.isArray(update) || inTransaction)
+            return;
+
+        // Make sure that the upsert inserted the correct document or update did not insert
+        // anything.
+        assert.eq(
+            isUpsert ? 1 : 0,
+            st.s.getDB(kDbName).coll.find(update["$set"] ? update["$set"] : update).itcount());
+    }
+
+    //
+    // Update Type Replacement-style.
+    //
+
+    // Test behaviours common to update and upsert.
+    [false, true].forEach(function(isUpsert) {
+        // Full shard key in query matches the update document.
+        assertUpdateWorked({x: 4, y: 3, z: 3}, {x: 4, y: 3, z: 3, a: 0}, isUpsert, 0);
+        assertUpdateWorked({x: 4, _id: 0, z: 3, y: 3}, {x: 4, y: 3, z: 3, a: 0}, isUpsert, 0);
+
+        // Case when upsert needs to insert a new document and the new document should belong in the
+        // same shard as the targeted shard. For non-upserts, it will be a no-op.
+        assertUpdateWorkedWithNoMatchingDoc(
+            {x: 4, y: 0, z: 0}, {x: 1, z: 3, y: 110, a: 90}, isUpsert);
+    });
+
+    //
+    // Test behaviours specific to non-upsert updates.
+    //
+
+    // Partial shard key in query can target a single shard, and shard key of existing document is
+    // the same as the replacement's.
+    assertUpdateWorked({x: 4}, {x: 4, y: 3, z: 3, a: 1}, false, 0);
+    assertUpdateWorked({x: 4, _id: 0, z: 3}, {y: 3, x: 4, z: 3, a: 3}, false, 0);
+
+    // Parital shard key in the query, update succeeds with no op when there is no matching document
+    // for the query.
+    assertUpdateWorkedWithNoMatchingDoc({x: 10}, {x: 10, y: 3, z: 3, a: 5}, false);
+    assertUpdateWorkedWithNoMatchingDoc({x: 100, y: 55, a: 15}, {x: 100, y: 55, z: 3, a: 6}, false);
+    assertUpdateWorkedWithNoMatchingDoc({x: 11, _id: 3}, {x: 11, y: 3, z: 3, a: 7}, false);
+
+    // Partial shard key in query can target a single shard, but fails while attempting to
+    // modify shard key value.
+    assert.commandFailedWithCode(
+        st.s.getDB(kDbName).coll.update(
+            {x: 100, y: 50, a: 5}, {x: 100, y: 55, z: 3, a: 1}, {upsert: false}),
+        [31025]);
+    assert.commandFailedWithCode(
+        st.s.getDB(kDbName).coll.update({x: 4, z: 3}, {x: 4, y: 3, z: 4, a: 1}, {upsert: false}),
+        [31025]);
+
+    // Full shard key in query, matches no document.
+    assertUpdateWorkedWithNoMatchingDoc({x: 4, y: 0, z: 0}, {x: 1110, y: 55, z: 3, a: 111}, false);
+
+    // Partial shard key in query, but can still target a single shard.
+    assertUpdateWorkedWithNoMatchingDoc({x: 100, y: 51, a: 5}, {x: 110, y: 55, z: 3, a: 8}, false);
+
+    // Partial shard key in query cannot target a single shard, targeting happens using update
+    // document.
+
+    // When query doesn't match any doc.
+    assertUpdateWorkedWithNoMatchingDoc({x: 4, y: 0}, {x: 110, y: 55, z: 3, a: 110}, false);
+    assertUpdateWorkedWithNoMatchingDoc({_id: 1}, {x: 110, y: 55, z: 3, a: 110}, false);
+
+    // When query matches a doc and updates sucessfully.
+    assertUpdateWorked({_id: 0, y: 3}, {z: 3, x: 4, y: 3, a: 2}, false, 0);
+    assertUpdateWorked({_id: 0}, {z: 3, x: 4, y: 3, replStyle: 2}, false, 0);
+
+    // When query matches a doc and fails to update because shard key needs to be updated.
+    assert.commandFailedWithCode(
+        st.s.getDB(kDbName).coll.update({}, {x: 110, y: 55, z: 3, a: 110}, false), 31025);
+    assert.commandFailedWithCode(
+        st.s.getDB(kDbName).coll.update({_id: 2}, {x: 110, y: 55, z: 3, a: 110}, false), 31025);
+
+    //
+    // 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 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);
+
+    // The above upsert works with transactions.
+    session.startTransaction();
+    assertUpdateWorkedWithNoMatchingDoc({x: 4, y: 0, z: 0}, updateDoc, true, true);
+    session.commitTransaction();
+    assert.eq(1, st.s.getDB(kDbName).coll.find(updateDoc).itcount());
+
+    // Full shard key not specified in query.
+
+    // Query on partial shard key.
+    assert.commandFailedWithCode(
+        st.s.getDB(kDbName).coll.update(
+            {x: 100, y: 50, a: 5}, {x: 100, y: 55, z: 3, a: 1}, {upsert: true}),
+        ErrorCodes.ShardKeyNotFound);
+    assert.commandFailedWithCode(
+        st.s.getDB(kDbName).coll.update(
+            {x: 100, y: 50, nonExistingField: true}, {x: 100, y: 55, z: 3, a: 1}, {upsert: true}),
+        ErrorCodes.ShardKeyNotFound);
+
+    // Query on partial shard key with _id.
+    assert.commandFailedWithCode(
+        st.s.getDB(kDbName).coll.update(
+            {x: 100, y: 50, a: 5, _id: 0}, {x: 100, y: 55, z: 3, a: 1}, {upsert: true}),
+        ErrorCodes.ShardKeyNotFound);
+    assert.commandFailedWithCode(
+        st.s.getDB(kDbName).coll.update({x: 100, y: 50, a: 5, _id: 0, nonExistingField: true},
+                                        {x: 100, y: 55, z: 3, a: 1},
+                                        {upsert: true}),
+        ErrorCodes.ShardKeyNotFound);
+
+    // Query on only _id.
+    assert.commandFailedWithCode(
+        st.s.getDB(kDbName).coll.update({_id: 0}, {z: 3, x: 4, y: 3, a: 2}, {upsert: true}),
+        ErrorCodes.ShardKeyNotFound);
+    assert.commandFailedWithCode(
+        st.s.getDB(kDbName).coll.update(
+            {_id: "nonExisting"}, {z: 3, x: 4, y: 3, a: 2}, {upsert: true}),
+        ErrorCodes.ShardKeyNotFound);
+
+    //
+    // Update Type Op-style.
+    //
+
+    // Test behaviours common to update and upsert.
+    [false, true].forEach(function(isUpsert) {
+        // Full shard key in query.
+        assertUpdateWorked({x: 4, _id: 0, z: 3, y: 3}, {"$set": {opStyle: 1}}, isUpsert, 0);
+        assertUpdateWorked({x: 4, z: 3, y: 3}, {"$set": {opStyle: 2}}, isUpsert, 0);
+
+        // Case when upsert needs to insert a new document and the new document should belong in the
+        // same shard as the targetted shard. For non-upserts, it will be a no op.
+        assertUpdateWorkedWithNoMatchingDoc(
+            {x: 4, y: 0, z: 0}, {"$set": {x: 1, z: 3, y: 111, a: 90}}, isUpsert);
+    });
+
+    // Test behaviours specific to non-upsert updates.
+
+    // Full shard key in query, matches no document.
+    assertUpdateWorkedWithNoMatchingDoc(
+        {x: 4, y: 0, z: 0}, {"$set": {x: 2110, y: 55, z: 3, a: 111}}, false);
+
+    // Partial shard key in query, but can still target a single shard.
+    assertUpdateWorkedWithNoMatchingDoc(
+        {x: 100, y: 51, a: 112}, {"$set": {x: 110, y: 55, z: 3, a: 8}}, false);
+
+    // Query on _id works for update.
+    assertUpdateWorked({_id: 0}, {"$set": {opStyle: 6}}, false, 0);
+    assertUpdateWorked({_id: 0, y: 3}, {"$set": {opStyle: 8, y: 3, x: 4}}, false, 0);
+
+    // Parital shard key in the query targets single shard. Update succeeds with no op when there is
+    // no matching document for the query.
+    assertUpdateWorkedWithNoMatchingDoc({x: 14, _id: 0}, {"$set": {opStyle: 5}}, false);
+    assertUpdateWorkedWithNoMatchingDoc({x: 14}, {"$set": {opStyle: 5}}, false);
+
+    assertUpdateWorkedWithNoMatchingDoc({x: -1, y: 0}, {"$set": {z: 3, y: 110, a: 91}}, false);
+
+    // Partial shard key in query can target a single shard and doesn't try to update shard key
+    // value.
+    assertUpdateWorked({x: 4, z: 3}, {"$set": {opStyle: 3}}, false, 0);
+    assertUpdateWorked({x: 4, _id: 0, z: 3}, {"$set": {y: 3, x: 4, z: 3, opStyle: 4}}, false, 0);
+
+    // Partial shard key in query can target a single shard, but fails while attempting to modify
+    // shard key value.
+    assert.commandFailedWithCode(
+        st.s.getDB(kDbName).coll.update(
+            {_id: 1, x: 100, z: 3, a: 5}, {"$set": {y: 55, a: 11}}, {upsert: false}),
+        [31025]);
+    assert.commandFailedWithCode(
+        st.s.getDB(kDbName).coll.update(
+            {x: 4, z: 3}, {"$set": {x: 4, y: 3, z: 4, a: 1}}, {upsert: false}),
+        [31025]);
+
+    // 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}};
+    assert.commandFailedWithCode(
+        st.s.getDB(kDbName).coll.update({x: 4, y: 0, z: 0, opStyle: true}, update, {upsert: true}),
+        ErrorCodes.IllegalOperation);
+
+    // The above upsert works with transactions.
+    session.startTransaction();
+    assertUpdateWorkedWithNoMatchingDoc({x: 4, y: 0, z: 0, opStyle: true}, update, true, true);
+    session.commitTransaction();
+    assert.eq(1, st.s.getDB(kDbName).coll.find(update["$set"]).itcount());
+
+    // Full shard key not specified in query.
+
+    // Query on _id doesn't work for upserts.
+    assert.commandFailedWithCode(
+        st.s.getDB(kDbName).coll.update(
+            {_id: 0}, {"$set": {x: 2, y: 11, z: 10, opStyle: 7}}, {upsert: true}),
+        ErrorCodes.ShardKeyNotFound);
+
+    // Partial shard key can target single shard. This style of update can work if SERVER-41243 is
+    // implemented.
+    assert.commandFailedWithCode(
+        st.s.getDB(kDbName).coll.update({x: 14}, {"$set": {opStyle: 5}}, {upsert: true}),
+        ErrorCodes.ShardKeyNotFound);
+    assert.commandFailedWithCode(
+        st.s.getDB(kDbName).coll.update({x: 100, y: 51, nonExistingField: true},
+                                        {"$set": {x: 110, y: 55, z: 3, a: 8}},
+                                        {upsert: true}),
+        ErrorCodes.ShardKeyNotFound);
+
+    // Partial shard key cannot target single shard.
+    assert.commandFailedWithCode(
+        st.s.getDB(kDbName).coll.update(
+            {_id: 0, y: 3}, {"$set": {z: 3, x: 4, y: 3, a: 2}}, {upsert: true}),
+        ErrorCodes.ShardKeyNotFound);
+    assert.commandFailedWithCode(
+        st.s.getDB(kDbName).coll.update({y: 3}, {"$set": {z: 3, x: 4, y: 3, a: 2}}, {upsert: true}),
+        ErrorCodes.ShardKeyNotFound);
+    //
+    // Update with pipeline.
+    //
+
+    // Test behaviours common to update and upsert.
+    [false, true].forEach(function(isUpsert) {
+        // Full shard key in query.
+        assertUpdateWorked(
+            {_id: 0, x: 4, z: 3, y: 3}, [{$addFields: {pipelineUpdate: isUpsert}}], isUpsert, 0);
+        assert.eq(1,
+                  st.s.getDB(kDbName)
+                      .coll.find({_id: 0, x: 4, z: 3, y: 3, pipelineUpdate: isUpsert})
+                      .itcount());
+        assertUpdateWorkedWithNoMatchingDoc(
+            {_id: 15, x: 44, z: 3, y: 3}, [{$addFields: {pipelineUpdate: true}}], isUpsert);
+        assert.eq(isUpsert ? 1 : 0,
+                  st.s.getDB(kDbName)
+                      .coll.find({_id: 15, x: 44, z: 3, y: 3, pipelineUpdate: true})
+                      .itcount());
+
+        assertUpdateWorkedWithNoMatchingDoc(
+            {x: 45, z: 4, y: 3}, [{$addFields: {pipelineUpdate: true}}], isUpsert);
+        assert.eq(
+            isUpsert ? 1 : 0,
+            st.s.getDB(kDbName).coll.find({x: 45, z: 4, y: 3, pipelineUpdate: true}).itcount());
+
+        // Case when upsert needs to insert a new document and the new document should belong in the
+        // same shard as the targeted shard.
+        assertUpdateWorkedWithNoMatchingDoc({x: 4, y: 0, z: 0},
+                                            [{
+                                               "$project": {
+                                                   x: {$literal: 3},
+                                                   y: {$literal: 33},
+                                                   z: {$literal: 3},
+                                                   pipelineUpdate: {$literal: true}
+                                               }
+                                            }],
+                                            isUpsert);
+        assert.eq(
+            isUpsert ? 1 : 0,
+            st.s.getDB(kDbName).coll.find({x: 3, z: 3, y: 33, pipelineUpdate: true}).itcount());
+    });
+
+    // Test behaviours specific to non-upsert updates.
+
+    // Full shard key in query, matches no document.
+    assertUpdateWorkedWithNoMatchingDoc({x: 4, y: 0, z: 0},
+                                        [{
+                                           "$project": {
+                                               x: {$literal: 2111},
+                                               y: {$literal: 55},
+                                               z: {$literal: 3},
+                                               pipelineUpdate: {$literal: true}
+                                           }
+                                        }],
+                                        false);
+    assert.eq(
+        0, st.s.getDB(kDbName).coll.find({x: 2111, z: 3, y: 55, pipelineUpdate: true}).itcount());
+
+    // Partial shard key in query targets single shard but doesn't match any document on that shard.
+    assertUpdateWorkedWithNoMatchingDoc({_id: 14, z: 4, x: 3}, [{$addFields: {foo: 4}}], false);
+
+    // Partial shard key in query can target a single shard and doesn't try to update shard key
+    // value.
+    assertUpdateWorkedWithNoMatchingDoc(
+        {x: 46, z: 4}, [{$addFields: {y: 10, pipelineUpdateNoOp: false}}], false);
+    assertUpdateWorked({x: 4, z: 3}, [{$addFields: {pipelineUpdateDoc: false}}], false, 0);
+
+    // Partial shard key in query cannot target a single shard.
+    assert.commandFailedWithCode(
+        st.s.getDB(kDbName).coll.update({z: 3, y: 3}, [{$addFields: {foo: 4}}], {upsert: false}),
+        [72, ErrorCodes.InvalidOptions]);
+
+    // 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.
+    session.startTransaction();
+    assertUpdateWorkedWithNoMatchingDoc({x: 4, y: 0, z: 0, pipelineUpdate: true},
+                                        [{
+                                           "$project": {
+                                               x: {$literal: 2111},
+                                               y: {$literal: 55},
+                                               z: {$literal: 3},
+                                               pipelineUpdate: {$literal: true}
+                                           }
+                                        }],
+                                        true);
+    session.commitTransaction();
+    assert.eq(
+        1, st.s.getDB(kDbName).coll.find({x: 2111, y: 55, z: 3, pipelineUpdate: true}).itcount());
+
+    // Full shard key not specified in query.
+    assert.commandFailedWithCode(
+        st.s.getDB(kDbName).coll.update(
+            {_id: 18, z: 4, x: 3}, [{$addFields: {foo: 4}}], {upsert: true}),
+        ErrorCodes.ShardKeyNotFound);
+    assert.commandFailedWithCode(
+        st.s.getDB(kDbName).coll.update({_id: 0},
+                                        [{
+                                           "$project": {
+                                               x: {$literal: 2111},
+                                               y: {$literal: 55},
+                                               z: {$literal: 3},
+                                               pipelineUpdate: {$literal: true}
+                                           }
+                                        }],
+                                        {upsert: true}),
+        ErrorCodes.ShardKeyNotFound);
+
+    st.stop();
+})();
diff --git a/jstests/sharding/update_immutable_fields.js b/jstests/sharding/update_immutable_fields.js
index 5e81d075c37..287e750c176 100644
--- a/jstests/sharding/update_immutable_fields.js
+++ b/jstests/sharding/update_immutable_fields.js
@@ -58,12 +58,18 @@
     assert.writeOK(shard0Coll.update({_id: 1}, {a: 1}));
 
     // Update existing doc ($set), same shard key value
-    assert.writeOK(shard0Coll.update({_id: 1}, {$set: {a: 1}}));
+    assert.commandWorked(shard0Coll.update({_id: 1}, {$set: {a: 1}}));
 
-    // Error due to mutating the shard key (replacement)
-    assert.writeError(shard0Coll.update({_id: 1}, {b: 1}));
+    // Error when trying to update a shard key outside of a transaction.
+    assert.commandFailedWithCode(shard0Coll.update({_id: 1, a: 1}, {_id: 1, a: 2}),
+                                 ErrorCodes.IllegalOperation);
+    assert.commandFailedWithCode(shard0Coll.update({_id: 1, a: 1}, {"$set": {a: 2}}),
+                                 ErrorCodes.IllegalOperation);
 
-    // Error due to mutating the shard key ($set)
+    // Error when unsetting shard key.
+    assert.writeError(shard0Coll.update({_id: 1}, {b: 3}));
+
+    // Error when unsetting shard key ($set).
     assert.writeError(shard0Coll.update({_id: 1}, {$unset: {a: 1}}));
 
     // Error due to removing all the embedded fields.
diff --git a/jstests/sharding/update_replace_id.js b/jstests/sharding/update_replace_id.js
index 9fe3538e4f5..db8c878674f 100644
--- a/jstests/sharding/update_replace_id.js
+++ b/jstests/sharding/update_replace_id.js
@@ -37,8 +37,7 @@
         }
     }
 
-    // Run the set of tests relevant to the given shardKey.
-    function runReplacementUpdateTests(shardKey) {
+    function setUpData() {
         // Write a single document to shard0 and verify that it is present.
         mongosColl.insert({_id: -100, a: -100, msg: "not_updated"});
         assert.docEq(shard0DB.test.find({_id: -100}).toArray(),
@@ -50,6 +49,10 @@
 
         // Clear and restart the profiler on both shards.
         restartProfiling();
+    }
+
+    function runReplacementUpdateTestsForHashedShardKey() {
+        setUpData();
 
         // Perform a replacement update whose query is an exact match on _id and whose replacement
         // document contains the remainder of the shard key. Despite the fact that the replacement
@@ -104,12 +107,45 @@
                 errName: {$exists: false}
             }
         });
+    }
 
-        // The remainder of the tests are only relevant for a compound shard key. If the key is
-        // {_id: 'hashed'}, we stop at this point.
-        if (shardKey._id === "hashed") {
-            return;
-        }
+    function runReplacementUpdateTestsForCompoundShardKey() {
+        setUpData();
+
+        // Perform a replacement update whose query is an exact match on _id and whose replacement
+        // document contains the remainder of the shard key. Despite the fact that the replacement
+        // document does not contain the entire shard key, we expect that mongoS will extract the
+        // _id from the query and combine it with the replacement doc to target a single shard.
+        let writeRes = assert.commandWorked(
+            mongosColl.update({_id: -100}, {a: -100, msg: "update_extracted_id_from_query"}));
+
+        // Verify that the update did not modify the orphan document.
+        assert.docEq(shard1DB.test.find({_id: -100}).toArray(),
+                     [{_id: -100, a: -100, msg: "not_updated"}]);
+        assert.eq(writeRes.nMatched, 1);
+        assert.eq(writeRes.nModified, 1);
+
+        // Verify that the update only targeted shard0 and that the resulting document appears as
+        // expected.
+        assert.docEq(mongosColl.find({_id: -100}).toArray(),
+                     [{_id: -100, a: -100, msg: "update_extracted_id_from_query"}]);
+        profilerHasSingleMatchingEntryOrThrow({
+            profileDB: shard0DB,
+            filter: {op: "update", "command.u.msg": "update_extracted_id_from_query"}
+        });
+        profilerHasZeroMatchingEntriesOrThrow({
+            profileDB: shard1DB,
+            filter: {op: "update", "command.u.msg": "update_extracted_id_from_query"}
+        });
+
+        // An upsert whose query doesn't have full shard key will fail.
+        assert.commandFailedWithCode(
+            mongosColl.update(
+                {_id: 101}, {a: 101, msg: "upsert_extracted_id_from_query"}, {upsert: true}),
+            ErrorCodes.ShardKeyNotFound);
+
+        // Verify that the document did not perform any writes.
+        assert.docEq(mongosColl.find({_id: 101}).itcount(), 0);
 
         // Verify that an update whose query contains an exact match on _id but whose replacement
         // doc does not contain all other shard key fields will be rejected by mongoS.
@@ -153,14 +189,14 @@
         mongosColl, {_id: 1, a: 1}, {_id: 0, a: 0}, {_id: 1, a: 1}, mongosDB.getName(), true);
 
     // Run the replacement behaviour tests that are relevant to a compound key that includes _id.
-    runReplacementUpdateTests({_id: 1, a: 1});
+    runReplacementUpdateTestsForCompoundShardKey();
 
     // Drop and reshard the collection on {_id: "hashed"}, which will autosplit across both shards.
     assert(mongosColl.drop());
     mongosDB.adminCommand({shardCollection: mongosColl.getFullName(), key: {_id: "hashed"}});
 
     // Run the replacement behaviour tests relevant to a collection sharded on {_id: "hashed"}.
-    runReplacementUpdateTests({_id: "hashed"});
+    runReplacementUpdateTestsForHashedShardKey();
 
     st.stop();
 })();
\ No newline at end of file
diff --git a/jstests/sharding/update_shard_key_doc_moves_shards.js b/jstests/sharding/update_shard_key_doc_moves_shards.js
index 09117dddbd9..c381dd6809a 100644
--- a/jstests/sharding/update_shard_key_doc_moves_shards.js
+++ b/jstests/sharding/update_shard_key_doc_moves_shards.js
@@ -19,8 +19,6 @@
     assert.commandWorked(mongos.adminCommand({enableSharding: kDbName}));
     st.ensurePrimaryShard(kDbName, shard0);
 
-    // TODO SERVER-39158: Add tests that replaement style updates work as well.
-
     function changeShardKeyWhenFailpointsSet(session, sessionDB, runInTxn, isFindAndModify) {
         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});
@@ -101,13 +99,30 @@
         mongos.getDB(kDbName).foo.drop();
     }
 
+    //
     // Test that changing the shard key works correctly when either the update or findAndModify
     // command is used and when the command is run either as a retryable write or in a transaction.
-    // Pairs represent [shouldRunCommandInTxn, runUpdateAsFindAndModifyCmd]
-    let changeShardKeyOptions = [[false, false], [true, false], [false, true], [true, true]];
-    changeShardKeyOptions.forEach(function(updatePair) {
-        let runInTxn = updatePair[0];
-        let isFindAndModify = updatePair[1];
+    // Tuples represent [shouldRunCommandInTxn, runUpdateAsFindAndModifyCmd, isUpsert].
+    //
+
+    const changeShardKeyOptions = [
+        [false, false, false],
+        [true, false, false],
+        [true, true, false],
+        [false, true, false],
+        [false, false, true],
+        [true, false, true],
+        [false, true, true],
+        [true, true, true]
+    ];
+
+    //
+    //  Tests for op-style updates.
+    //
+    changeShardKeyOptions.forEach(function(updateConfig) {
+        let runInTxn = updateConfig[0];
+        let isFindAndModify = updateConfig[1];
+        let upsert = updateConfig[2];
 
         jsTestLog("Testing changing the shard key using " +
                   (isFindAndModify ? "findAndModify command " : "update command ") +
@@ -116,9 +131,6 @@
         let session = st.s.startSession({retryWrites: runInTxn ? false : true});
         let sessionDB = session.getDatabase(kDbName);
 
-        // Modify updates
-
-        // upsert : false
         assertCanUpdatePrimitiveShardKey(st,
                                          kDbName,
                                          ns,
@@ -128,7 +140,7 @@
                                          isFindAndModify,
                                          [{"x": 300}, {"x": 4}],
                                          [{"$set": {"x": 30}}, {"$set": {"x": 600}}],
-                                         false);
+                                         upsert);
         assertCanUpdateDottedPath(st,
                                   kDbName,
                                   ns,
@@ -138,7 +150,7 @@
                                   isFindAndModify,
                                   [{"x.a": 300}, {"x.a": 4}],
                                   [{"$set": {"x": {"a": 30}}}, {"$set": {"x": {"a": 600}}}],
-                                  false);
+                                  upsert);
         assertCanUpdatePartialShardKey(st,
                                        kDbName,
                                        ns,
@@ -148,9 +160,54 @@
                                        isFindAndModify,
                                        [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
                                        [{"$set": {"x": 30}}, {"$set": {"x": 600}}],
-                                       false);
+                                       upsert);
+
+        // Failure cases. These tests do not take 'upsert' as an option so we do not need to test
+        // them for both upsert true and false.
+        if (!upsert) {
+            assertCannotUpdate_id(
+                st, kDbName, ns, session, sessionDB, runInTxn, isFindAndModify, {"_id": 300}, {
+                    "$set": {"_id": 30}
+                });
+            assertCannotUpdate_idDottedPath(
+                st, kDbName, ns, session, sessionDB, runInTxn, isFindAndModify, {"_id.a": 300}, {
+                    "$set": {"_id": {"a": 30}}
+                });
+            assertCannotUpdateSKToArray(
+                st, kDbName, ns, session, sessionDB, runInTxn, isFindAndModify, {"x": 300}, {
+                    "$set": {"x": [30]}
+                });
+            assertCannotUnsetSKField(
+                st, kDbName, ns, session, sessionDB, runInTxn, isFindAndModify, {"x": 300}, {
+                    "$unset": {"x": 1}
+                });
+
+            if (!isFindAndModify) {
+                assertCannotUpdateWithMultiTrue(
+                    st, kDbName, ns, session, sessionDB, runInTxn, {"x": 300}, {"$set": {"x": 30}});
+            }
+            changeShardKeyWhenFailpointsSet(session, sessionDB, runInTxn, isFindAndModify);
+        }
+    });
+
+    //
+    // Tests for replacement style updates.
+    //
+
+    mongos.getDB(kDbName).foo.drop();
+
+    changeShardKeyOptions.forEach(function(updateConfig) {
+        let runInTxn = updateConfig[0];
+        let isFindAndModify = updateConfig[1];
+        let upsert = updateConfig[2];
+
+        jsTestLog("Testing changing the shard key using replacement style update and " +
+                  (isFindAndModify ? "findAndModify command " : "update command ") +
+                  (runInTxn ? "in transaction " : "as retryable write"));
+
+        let session = st.s.startSession({retryWrites: runInTxn ? false : true});
+        let sessionDB = session.getDatabase(kDbName);
 
-        // upsert : true
         assertCanUpdatePrimitiveShardKey(st,
                                          kDbName,
                                          ns,
@@ -159,8 +216,8 @@
                                          runInTxn,
                                          isFindAndModify,
                                          [{"x": 300}, {"x": 4}],
-                                         [{"$set": {"x": 30}}, {"$set": {"x": 600}}],
-                                         true);
+                                         [{"x": 30}, {"x": 600}],
+                                         upsert);
         assertCanUpdateDottedPath(st,
                                   kDbName,
                                   ns,
@@ -169,8 +226,8 @@
                                   runInTxn,
                                   isFindAndModify,
                                   [{"x.a": 300}, {"x.a": 4}],
-                                  [{"$set": {"x": {"a": 30}}}, {"$set": {"x": {"a": 600}}}],
-                                  true);
+                                  [{"x": {"a": 30}}, {"x": {"a": 600}}],
+                                  upsert);
         assertCanUpdatePartialShardKey(st,
                                        kDbName,
                                        ns,
@@ -179,33 +236,31 @@
                                        runInTxn,
                                        isFindAndModify,
                                        [{"x": 300, "y": 80}, {"x": 4, "y": 3}],
-                                       [{"$set": {"x": 30}}, {"$set": {"x": 600}}],
-                                       true);
-
-        // failing cases
-        assertCannotUpdate_id(
-            st, kDbName, ns, session, sessionDB, runInTxn, isFindAndModify, {"_id": 300}, {
-                "$set": {"_id": 30}
-            });
-        assertCannotUpdate_idDottedPath(
-            st, kDbName, ns, session, sessionDB, runInTxn, isFindAndModify, {"_id.a": 300}, {
-                "$set": {"_id": {"a": 30}}
-            });
-        assertCannotUpdateSKToArray(
-            st, kDbName, ns, session, sessionDB, runInTxn, isFindAndModify, {"x": 300}, {
-                "$set": {"x": [30]}
-            });
-        assertCannotUnsetSKField(
-            st, kDbName, ns, session, sessionDB, runInTxn, isFindAndModify, {"x": 300}, {
-                "$unset": {"x": 1}
-            });
-
-        if (!isFindAndModify) {
-            assertCannotUpdateWithMultiTrue(
-                st, kDbName, ns, session, sessionDB, runInTxn, {"x": 300}, {"$set": {"x": 30}});
+                                       [{"x": 30, "y": 80}, {"x": 600, "y": 3}],
+                                       upsert);
+
+        // Failure cases. These tests do not take 'upsert' as an option so we do not need to test
+        // them for both upsert true and false.
+        if (!upsert) {
+            assertCannotUpdate_id(
+                st, kDbName, ns, session, sessionDB, runInTxn, isFindAndModify, {"_id": 300}, {
+                    "_id": 30
+                });
+            assertCannotUpdate_idDottedPath(
+                st, kDbName, ns, session, sessionDB, runInTxn, isFindAndModify, {"_id.a": 300}, {
+                    "_id": {"a": 30}
+                });
+            if (!isFindAndModify) {
+                assertCannotUpdateWithMultiTrue(
+                    st, kDbName, ns, session, sessionDB, runInTxn, {"x": 300}, {"x": 30});
+            }
+            assertCannotUpdateSKToArray(
+                st, kDbName, ns, session, sessionDB, runInTxn, isFindAndModify, {"x": 300}, {
+                    "x": [30]
+                });
+            assertCannotUnsetSKField(
+                st, kDbName, ns, session, sessionDB, runInTxn, isFindAndModify, {"x": 300}, {});
         }
-
-        changeShardKeyWhenFailpointsSet(session, sessionDB, runInTxn, isFindAndModify);
     });
 
     let session = st.s.startSession({retryWrites: true});
@@ -353,10 +408,13 @@
 
     // ----Assert correct behavior when update is sent directly to a shard----
 
-    // TODO SERVER-39158: Add replacement tests as well
     shardCollectionMoveChunks(st, kDbName, ns, {"x": 1}, docsToInsert, {"x": 100}, {"x": 300});
     cleanupOrphanedDocs(st, ns);
 
+    //
+    // For Op-style updates.
+    //
+
     // An update sent directly to a shard cannot change the shard key.
     assert.commandFailedWithCode(
         st.rs1.getPrimary().getDB(kDbName).foo.update({"x": 500}, {$set: {"x": 2}}),
@@ -376,6 +434,28 @@
     assert.eq(0, res.nModified);
     assert.eq(0, res.nUpserted);
 
+    //
+    // For Replacement style updates.
+    //
+
+    // An update sent directly to a shard cannot change the shard key.
+    assert.commandFailedWithCode(
+        st.rs1.getPrimary().getDB(kDbName).foo.update({"x": 500}, {"x": 2}),
+        ErrorCodes.ImmutableField);
+    assert.commandFailedWithCode(
+        st.rs1.getPrimary().getDB(kDbName).foo.update({"x": 1000}, {"x": 2}, {upsert: true}),
+        ErrorCodes.ImmutableField);
+    assert.commandFailedWithCode(
+        st.rs0.getPrimary().getDB(kDbName).foo.update({"x": 1000}, {"x": 2}, {upsert: true}),
+        ErrorCodes.ImmutableField);
+
+    // The query will not match a doc and upsert is false, so this will not fail but will be a
+    // no-op.
+    res = assert.commandWorked(st.rs0.getPrimary().getDB(kDbName).foo.update({"x": 500}, {"x": 2}));
+    assert.eq(0, res.nMatched);
+    assert.eq(0, res.nModified);
+    assert.eq(0, res.nUpserted);
+
     mongos.getDB(kDbName).foo.drop();
 
     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 53e68480437..3c7874c670f 100644
--- a/jstests/sharding/update_shard_key_doc_on_same_shard.js
+++ b/jstests/sharding/update_shard_key_doc_on_same_shard.js
@@ -201,6 +201,9 @@
         st, kDbName, ns, session, sessionDB, false, {"x": 300}, {"x": 600});
     assertCannotDoReplacementUpdateWhereShardKeyMissingFields(
         st, kDbName, ns, session, sessionDB, false, false, {"x": 300, "y": 80}, {"x": 600});
+    // Shard key fields are missing in query.
+    assertCannotDoReplacementUpdateWhereShardKeyMissingFields(
+        st, kDbName, ns, session, sessionDB, false, false, {"x": 300}, {"x": 600, "y": 80, "a": 2});
     assertCannotUpdateSKToArray(
         st, kDbName, ns, session, sessionDB, false, false, {"x": 300}, {"x": [300]});
 
@@ -355,6 +358,9 @@
         st, kDbName, ns, session, sessionDB, false, true, {"_id.a": 300}, {"_id": {"a": 600}});
     assertCannotDoReplacementUpdateWhereShardKeyMissingFields(
         st, kDbName, ns, session, sessionDB, false, true, {"x": 300, "y": 80}, {"x": 600});
+    // Shard key fields are missing in query.
+    assertCannotDoReplacementUpdateWhereShardKeyMissingFields(
+        st, kDbName, ns, session, sessionDB, false, true, {"x": 300}, {"x": 600, "y": 80, "a": 2});
     assertCannotUpdateSKToArray(
         st, kDbName, ns, session, sessionDB, false, true, {"x": 300}, {"x": [300]});
 
@@ -532,6 +538,10 @@
         st, kDbName, ns, session, sessionDB, true, {"x": 300}, {"x": 600});
     assertCannotDoReplacementUpdateWhereShardKeyMissingFields(
         st, kDbName, ns, session, sessionDB, true, false, {"x": 300, "y": 80}, {"x": 600});
+    // Shard key fields are missing in query.
+    assertCannotDoReplacementUpdateWhereShardKeyMissingFields(
+        st, kDbName, ns, session, sessionDB, true, false, {"x": 300}, {"x": 600, "y": 80, "a": 2});
+
     assertCannotUpdateSKToArray(
         st, kDbName, ns, session, sessionDB, true, false, {"x": 300}, {"x": [300]});
 
@@ -685,7 +695,10 @@
     assertCannotUpdate_idDottedPath(
         st, kDbName, ns, session, sessionDB, true, true, {"_id.a": 300}, {"_id": {"a": 600}});
     assertCannotDoReplacementUpdateWhereShardKeyMissingFields(
-        st, kDbName, ns, session, sessionDB, true, false, {"x": 300, "y": 80}, {"x": 600});
+        st, kDbName, ns, session, sessionDB, true, true, {"x": 300, "y": 80}, {"x": 600});
+    // Shard key fields are missing in query.
+    assertCannotDoReplacementUpdateWhereShardKeyMissingFields(
+        st, kDbName, ns, session, sessionDB, true, true, {"x": 300}, {"x": 600, "y": 80, "a": 2});
     assertCannotUpdateSKToArray(
         st, kDbName, ns, session, sessionDB, true, true, {"x": 300}, {"x": [300]});
 
@@ -771,4 +784,4 @@
 
     st.stop();
 
-})();
\ No newline at end of file
+})();
diff --git a/jstests/sharding/update_sharded.js b/jstests/sharding/update_sharded.js
index 0873152c4e0..432329f7210 100644
--- a/jstests/sharding/update_sharded.js
+++ b/jstests/sharding/update_sharded.js
@@ -23,24 +23,22 @@
         coll = db.getCollection(collName);
         coll.insert({_id: 1, key: 1});
 
-        // these are both upserts
-        coll.save({_id: 2, key: 2});
-        coll.update({_id: 3, key: 3}, {$set: {foo: 'bar'}}, {upsert: true});
+        // Replacment and Opstyle upserts.
+        assert.commandWorked(coll.update({_id: 2, key: 2}, {key: 2, foo: 'bar'}, {upsert: true}));
+        assert.commandWorked(coll.update({_id: 3, key: 3}, {$set: {foo: 'bar'}}, {upsert: true}));
 
         assert.eq(coll.count(), 3, "count A");
         assert.eq(coll.findOne({_id: 3}).key, 3, "findOne 3 key A");
         assert.eq(coll.findOne({_id: 3}).foo, 'bar', "findOne 3 foo A");
 
-        // update existing using save()
-        coll.save({_id: 1, key: 1, other: 1});
-
         // update existing using update()
-        coll.update({_id: 2}, {key: 2, other: 2});
-        coll.update({_id: 3}, {key: 3, other: 3});
+        assert.commandWorked(coll.update({_id: 1}, {key: 1, other: 1}));
+        assert.commandWorked(coll.update({_id: 2}, {key: 2, other: 2}));
+        assert.commandWorked(coll.update({_id: 3}, {key: 3, other: 3}));
 
         // do a replacement-style update which queries the shard key and keeps it constant
-        coll.save({_id: 4, key: 4});
-        coll.update({key: 4}, {key: 4, other: 4});
+        assert.commandWorked(coll.update({key: 4}, {_id: 4, key: 4}, {upsert: true}));
+        assert.commandWorked(coll.update({key: 4}, {key: 4, other: 4}));
         assert.eq(coll.find({key: 4, other: 4}).count(), 1, 'replacement update error');
         coll.remove({_id: 4});
 
diff --git a/jstests/sharding/upsert_sharded.js b/jstests/sharding/upsert_sharded.js
index 5d3dde7a5cd..317b3107e09 100644
--- a/jstests/sharding/upsert_sharded.js
+++ b/jstests/sharding/upsert_sharded.js
@@ -51,16 +51,28 @@
     assert.eq(1, upsertedXVal({$and: [{x: {$eq: 1}}]}, {$set: {a: 1}}));
     assert.eq(1, upsertedXVal({$or: [{x: {$eq: 1}}]}, {$set: {a: 1}}));
 
-    // shard key not extracted
-    assert.writeError(upsertedResult({}, {$set: {a: 1, x: 1}}));
-    assert.writeError(upsertedResult({x: {$gt: 1}}, {$set: {a: 1, x: 1}}));
-
-    // Shard key type errors
-    assert.writeError(upsertedResult({x: undefined}, {$set: {a: 1}}));
-    assert.writeError(upsertedResult({x: [1, 2]}, {$set: {a: 1}}));
-    assert.writeError(upsertedResult({x: {$eq: {$gt: 5}}}, {$set: {a: 1}}));
-    // Regex shard key is not extracted from queries, even exact matches
-    assert.writeError(upsertedResult({x: {$eq: /abc/}}, {$set: {a: 1}}));
+    // Missing shard key in query.
+    assert.commandFailedWithCode(upsertedResult({}, {$set: {a: 1, x: 1}}),
+                                 ErrorCodes.ShardKeyNotFound);
+
+    // Missing equality match on shard key in query.
+    assert.commandFailedWithCode(upsertedResult({x: {$gt: 10}}, {$set: {a: 1, x: 5}}),
+                                 ErrorCodes.ShardKeyNotFound);
+
+    // Regex shard key value in query is ambigious and cannot be extracted for an equality match.
+    assert.commandFailedWithCode(
+        upsertedResult({x: {$eq: /abc*/}}, {$set: {a: 1, x: "regexValue"}}),
+        ErrorCodes.ShardKeyNotFound);
+    assert.commandFailedWithCode(upsertedResult({x: {$eq: /abc/}}, {$set: {a: 1, x: /abc/}}),
+                                 ErrorCodes.ShardKeyNotFound);
+
+    // Shard key in query not extractable.
+    assert.commandFailedWithCode(upsertedResult({x: undefined}, {$set: {a: 1}}),
+                                 ErrorCodes.BadValue);
+    assert.commandFailedWithCode(upsertedResult({x: [1, 2]}, {$set: {a: 1}}),
+                                 ErrorCodes.ShardKeyNotFound);
+    assert.commandFailedWithCode(upsertedResult({x: {$eq: {$gt: 5}}}, {$set: {a: 1}}),
+                                 ErrorCodes.ShardKeyNotFound);
 
     // nested field extraction always fails with non-nested key - like _id, we require setting the
     // elements directly