SERVER-44382 Allow upserting documents with missing shard key fields

1 files changed, 115 insertions, 16 deletions

diff --git a/jstests/sharding/update_sharded.js b/jstests/sharding/update_sharded.js
index c0466216647..74a034c0899 100644
--- a/jstests/sharding/update_sharded.js
+++ b/jstests/sharding/update_sharded.js
@@ -1,5 +1,8 @@
 // Test simple updates issued through mongos. Updates have different constraints through mongos,
 // since shard key is immutable.
+//
+// Updating a shard key in a single statement may use a multi shard transaction.
+// @tags: [uses_transactions, uses_multi_shard_transaction]
 (function() {
 
 const s = new ShardingTest({name: "auto1", shards: 2, mongos: 1});
@@ -7,7 +10,9 @@ const s = new ShardingTest({name: "auto1", shards: 2, mongos: 1});
 s.adminCommand({enablesharding: "test"});
 s.ensurePrimaryShard('test', s.shard1.shardName);
 
-db = s.getDB("test");
+const db = s.getDB("test");
+const sessionDb = s.s.startSession({retryWrites: true}).getDatabase("test");
+let coll, sessionColl;
 
 // Repeat same tests with hashed shard key, to ensure identical behavior.
 s.shardColl("update0", {key: 1}, {key: 0}, {key: 1}, db.getName(), true);
@@ -21,6 +26,7 @@ for (let i = 0; i < 2; i++) {
     const hashedKey = (collName == "update1");
 
     coll = db.getCollection(collName);
+    sessionColl = sessionDb.getCollection(collName);  // Used for updates of the shard key.
     coll.insert({_id: 1, key: 1});
 
     // Replacment and Opstyle upserts.
@@ -48,8 +54,10 @@ for (let i = 0; i < 2; i++) {
         assert.eq(x._id, x.other, "_id == other");
     });
 
-    assert.writeError(coll.update({_id: 1, key: 1}, {$set: {key: 2}}));
-    assert.eq(coll.findOne({_id: 1}).key, 1, 'key unchanged');
+    assert.commandWorked(sessionColl.update({_id: 1, key: 1}, {$set: {key: 2}}));
+    assert.eq(coll.findOne({_id: 1}).key, 2, 'key changed');
+    assert.commandWorked(
+        sessionColl.update({_id: 1, key: 2}, {$set: {key: 1}}));  // Reset the key value.
 
     assert.commandWorked(coll.update({_id: 1, key: 1}, {$set: {foo: 2}}));
 
@@ -68,13 +76,19 @@ for (let i = 0; i < 2; i++) {
     assert.commandWorked(coll.update({_id: {$in: [ObjectId()]}}, {$set: {x: 1}}, {multi: false}));
 
     // Invalid extraction of exact _id from query
-    assert.writeError(coll.update({}, {$set: {x: 1}}, {multi: false}));
-    assert.writeError(coll.update({_id: {$gt: ObjectId()}}, {$set: {x: 1}}, {multi: false}));
-    assert.writeError(
-        coll.update({$or: [{_id: ObjectId()}, {_id: ObjectId()}]}, {$set: {x: 1}}, {multi: false}));
-    assert.writeError(coll.update(
-        {$and: [{_id: ObjectId()}, {_id: ObjectId()}]}, {$set: {x: 1}}, {multi: false}));
-    assert.writeError(coll.update({'_id.x': ObjectId()}, {$set: {x: 1}}, {multi: false}));
+    assert.commandFailedWithCode(coll.update({}, {$set: {x: 1}}, {multi: false}),
+                                 ErrorCodes.InvalidOptions);
+    assert.commandFailedWithCode(
+        coll.update({_id: {$gt: ObjectId()}}, {$set: {x: 1}}, {multi: false}),
+        ErrorCodes.InvalidOptions);
+    assert.commandFailedWithCode(
+        coll.update({$or: [{_id: ObjectId()}, {_id: ObjectId()}]}, {$set: {x: 1}}, {multi: false}),
+        ErrorCodes.InvalidOptions);
+    assert.commandFailedWithCode(
+        coll.update({$and: [{_id: ObjectId()}, {_id: ObjectId()}]}, {$set: {x: 1}}, {multi: false}),
+        ErrorCodes.InvalidOptions);
+    assert.commandFailedWithCode(coll.update({'_id.x': ObjectId()}, {$set: {x: 1}}, {multi: false}),
+                                 ErrorCodes.InvalidOptions);
 
     // Make sure we can extract exact shard key from certain queries
     assert.commandWorked(coll.update({key: ObjectId()}, {$set: {x: 1}}, {multi: false}));
@@ -85,13 +99,19 @@ for (let i = 0; i < 2; i++) {
     assert.commandWorked(coll.update({$and: [{key: ObjectId()}]}, {$set: {x: 1}}, {multi: false}));
 
     // Invalid extraction of exact key from query
-    assert.writeError(coll.update({}, {$set: {x: 1}}, {multi: false}));
-    assert.writeError(coll.update({'key.x': ObjectId()}, {$set: {x: 1}}, {multi: false}));
+    assert.commandFailedWithCode(coll.update({}, {$set: {x: 1}}, {multi: false}),
+                                 ErrorCodes.InvalidOptions);
+    assert.commandFailedWithCode(coll.update({'key.x': ObjectId()}, {$set: {x: 1}}, {multi: false}),
+                                 ErrorCodes.InvalidOptions);
 
     // Inexact queries may target a single shard. Range queries may target a single shard as
     // long as the collection is not hashed.
-    assert[hashedKey ? "writeError" : "writeOK"](
-        coll.update({key: {$gt: 0}}, {$set: {x: 1}}, {multi: false}));
+    if (hashedKey) {
+        assert.commandFailedWithCode(coll.update({key: {$gt: 0}}, {$set: {x: 1}}, {multi: false}),
+                                     ErrorCodes.InvalidOptions);
+    } else {
+        assert.commandWorked(coll.update({key: {$gt: 0}}, {$set: {x: 1}}, {multi: false}));
+    }
     // Note: {key:-1} and {key:-2} fall on shard0 for both hashed and ascending shardkeys.
     assert.commandWorked(
         coll.update({$or: [{key: -1}, {key: -2}]}, {$set: {x: 1}}, {multi: false}));
@@ -99,15 +119,94 @@ for (let i = 0; i < 2; i++) {
         coll.update({$and: [{key: -1}, {key: -2}]}, {$set: {x: 1}}, {multi: false}));
 
     // In cases where an inexact query does target multiple shards, single update is rejected.
-    assert.writeError(coll.update({key: {$gt: MinKey}}, {$set: {x: 1}}, {multi: false}));
-    assert.writeError(coll.update({$or: [{key: -10}, {key: 10}]}, {$set: {x: 1}}, {multi: false}));
+    assert.commandFailedWithCode(coll.update({key: {$gt: MinKey}}, {$set: {x: 1}}, {multi: false}),
+                                 ErrorCodes.InvalidOptions);
+    assert.commandFailedWithCode(
+        coll.update({$or: [{key: -10}, {key: 10}]}, {$set: {x: 1}}, {multi: false}),
+        ErrorCodes.InvalidOptions);
 
     // Make sure failed shard key or _id extraction doesn't affect the other
     assert.commandWorked(
         coll.update({'_id.x': ObjectId(), key: 1}, {$set: {x: 1}}, {multi: false}));
     assert.commandWorked(
         coll.update({_id: ObjectId(), 'key.x': 1}, {$set: {x: 1}}, {multi: false}));
+
+    // Can unset shard key with op style update.
+    assert.commandWorked(coll.insert({_id: 11, key: 1}));
+    assert.commandWorked(sessionColl.update({_id: 11, key: 1}, {$unset: {key: 1}}));
+    assert.docEq(coll.findOne({_id: 11}), {_id: 11});
+
+    // Can unset shard key with replacement style update.
+    assert.commandWorked(coll.insert({_id: 12, key: 1}));
+    assert.commandWorked(sessionColl.update({_id: 12, key: 1}, {_id: 12}));
+    assert.docEq(coll.findOne({_id: 12}), {_id: 12});
+
+    // Can unset shard key with pipeline style update.
+    assert.commandWorked(coll.insert({_id: 13, key: 1}));
+    assert.commandWorked(sessionColl.update({_id: 13, key: 1}, [{$unset: "key"}, {$set: {x: 1}}]));
+    assert.docEq(coll.findOne({_id: 13}), {_id: 13, x: 1});
+
+    // Can unset nested fields in the shard key.
+    assert.commandWorked(coll.insert({_id: 14, key: {a: 1, b: 1}}));
+    assert.commandWorked(sessionColl.update({_id: 14, key: {a: 1, b: 1}}, {$unset: {"key.a": 1}}));
+    assert.docEq(coll.findOne({_id: 14}), {_id: 14, key: {b: 1}});
+}
+
+// Tests for nested shard keys.
+
+function testNestedShardKeys(collName, keyPattern) {
+    s.adminCommand({shardCollection: db.getName() + "." + collName, key: keyPattern});
+    coll = db.getCollection(collName);
+    sessionColl = sessionDb.getCollection(collName);
+
+    //
+    // Verify full shard key path can be unset.
+    //
+
+    // Can unset shard key with op style update.
+    assert.commandWorked(coll.insert({_id: 11, skey: {skey: 1}}));
+    assert.commandWorked(sessionColl.update({_id: 11, "skey.skey": 1}, {$unset: {skey: 1}}));
+    assert.docEq(coll.findOne({_id: 11}), {_id: 11});
+
+    // Can unset shard key with replacement style update.
+    assert.commandWorked(coll.insert({_id: 12, skey: {skey: 1}}));
+    assert.commandWorked(sessionColl.update({_id: 12, "skey.skey": 1}, {_id: 12}));
+    assert.docEq(coll.findOne({_id: 12}), {_id: 12});
+
+    // Can unset shard key with pipeline style update.
+    assert.commandWorked(coll.insert({_id: 13, skey: {skey: 1}}));
+    assert.commandWorked(
+        sessionColl.update({_id: 13, "skey.skey": 1}, [{$unset: "skey"}, {$set: {x: 1}}]));
+    assert.docEq(coll.findOne({_id: 13}), {_id: 13, x: 1});
+
+    //
+    // Verify each field in a nested shard key can be unset.
+    //
+
+    // For op-style.
+    assert.commandWorked(coll.insert({_id: 14, skey: {skey: 1}}));
+    assert.commandWorked(sessionColl.update({_id: 14, "skey.skey": 1}, {$unset: {"skey.skey": 1}}));
+    assert.docEq(coll.findOne({_id: 14}), {_id: 14, skey: {}});
+    assert.commandWorked(sessionColl.update({_id: 14, skey: {}}, {$unset: {skey: 1}}));
+    assert.docEq(coll.findOne({_id: 14}), {_id: 14});
+
+    // For replacement style.
+    assert.commandWorked(coll.insert({_id: 15, skey: {skey: 1}}));
+    assert.commandWorked(sessionColl.update({_id: 15, "skey.skey": 1}, {skey: 1}));
+    assert.docEq(coll.findOne({_id: 15}), {_id: 15, skey: 1});
+    assert.commandWorked(sessionColl.update({_id: 15, skey: 1}, {$unset: {skey: 1}}));
+    assert.docEq(coll.findOne({_id: 15}), {_id: 15});
+
+    // For pipeline style.
+    assert.commandWorked(coll.insert({_id: 16, skey: {skey: 1}}));
+    assert.commandWorked(sessionColl.update({_id: 16, "skey.skey": 1}, [{$unset: "skey.skey"}]));
+    assert.docEq(coll.findOne({_id: 16}), {_id: 16, skey: {}});
+    assert.commandWorked(sessionColl.update({_id: 16, skey: {}}, [{$unset: "skey"}]));
+    assert.docEq(coll.findOne({_id: 16}), {_id: 16});
 }
 
+testNestedShardKeys("update_nested", {"skey.skey": 1});
+testNestedShardKeys("update_nested_hashed", {"skey.skey": "hashed"});
+
 s.stop();
 })();