SERVER-37578 Assert that a zone is associated with a shard during shard collection

(cherry picked from commit e720f0b57c8ce2ecf73926b542eb2c29ab7c1bc3)

3 files changed, 54 insertions, 3 deletions

diff --git a/jstests/sharding/shard_collection_existing_zones.js b/jstests/sharding/shard_collection_existing_zones.js
index 6bb7d233710..d6cb10c304e 100644
--- a/jstests/sharding/shard_collection_existing_zones.js
+++ b/jstests/sharding/shard_collection_existing_zones.js
@@ -15,6 +15,29 @@
     assert.commandWorked(mongos.adminCommand({enableSharding: kDbName}));
 
     /**
+     * Test that shardCollection correctly validates that a zone is associated with a shard.
+     */
+    function testShardZoneAssociationValidation(proposedShardKey, numberLongMin, numberLongMax) {
+        var zoneMin = numberLongMin ? {x: NumberLong(0)} : {x: 0};
+        var zoneMax = numberLongMax ? {x: NumberLong(10)} : {x: 10};
+        assert.commandWorked(configDB.tags.insert(
+            {_id: {ns: ns, min: zoneMin}, ns: ns, min: zoneMin, max: zoneMax, tag: zoneName}));
+
+        var tagDoc = configDB.tags.findOne();
+        assert.eq(ns, tagDoc.ns);
+        assert.eq(zoneMin, tagDoc.min);
+        assert.eq(zoneMax, tagDoc.max);
+        assert.eq(zoneName, tagDoc.tag);
+
+        assert.commandFailed(mongos.adminCommand({shardCollection: ns, key: proposedShardKey}));
+
+        assert.commandWorked(st.s.adminCommand({addShardToZone: shardName, zone: zoneName}));
+        assert.commandWorked(mongos.adminCommand({shardCollection: ns, key: proposedShardKey}));
+
+        assert.commandWorked(testDB.runCommand({drop: kCollName}));
+    }
+
+    /**
      * Test that shardCollection correctly validates shard key against existing zones.
      */
     function testShardKeyValidation(proposedShardKey, numberLongMin, numberLongMax, success) {
@@ -27,6 +50,7 @@
             {updateZoneKeyRange: ns, min: zoneMin, max: zoneMax, zone: zoneName}));
 
         var tagDoc = configDB.tags.findOne();
+        jsTestLog("xxx tag doc " + tojson(tagDoc));
         assert.eq(ns, tagDoc.ns);
         assert.eq(zoneMin, tagDoc.min);
         assert.eq(zoneMax, tagDoc.max);
@@ -104,8 +128,10 @@
         assert.commandWorked(testDB.runCommand({drop: kCollName}));
     }
 
-    // test that shardCollection uses existing zones to validate shard key
+    // test that shardCollection checks that a zone is associated with a shard.
+    testShardZoneAssociationValidation({x: 1}, false, false);
 
+    // test that shardCollection uses existing zones to validate shard key
     testShardKeyValidation({x: 1}, false, false, true);
 
     // cannot use a completely different key from the zone shard key or a key
diff --git a/src/mongo/db/s/config/initial_split_policy.cpp b/src/mongo/db/s/config/initial_split_policy.cpp
index c98df4b1414..e73f90de3ac 100644
--- a/src/mongo/db/s/config/initial_split_policy.cpp
+++ b/src/mongo/db/s/config/initial_split_policy.cpp
@@ -236,13 +236,24 @@ InitialSplitPolicy::generateShardCollectionInitialZonedChunks(
             const ShardId shardId = allShardIds[indx++ % allShardIds.size()];
             appendChunk(nss, lastChunkMax, tag.getMinKey(), &version, validAfter, shardId, &chunks);
         }
-        // create a chunk for the zone
+
+        // check that this tag is associated with a shard and if so create a chunk for the zone.
+        const auto& shardIdsForChunk = tagToShards.find(tag.getTag())->second;
+        uassert(50973,
+                str::stream()
+                    << "cannot shard collection "
+                    << tag.getNS().ns()
+                    << " because it is associated with zone: "
+                    << tag.getTag()
+                    << " which is not associated with a shard. please add this zone to a shard.",
+                !shardIdsForChunk.empty());
+
         appendChunk(nss,
                     tag.getMinKey(),
                     tag.getMaxKey(),
                     &version,
                     validAfter,
-                    tagToShards.find(tag.getTag())->second[0],
+                    shardIdsForChunk[0],
                     &chunks);
         lastChunkMax = tag.getMaxKey();
     }
diff --git a/src/mongo/db/s/config/initial_split_policy_test.cpp b/src/mongo/db/s/config/initial_split_policy_test.cpp
index a7a039f58b0..2b0c9198eee 100644
--- a/src/mongo/db/s/config/initial_split_policy_test.cpp
+++ b/src/mongo/db/s/config/initial_split_policy_test.cpp
@@ -430,5 +430,19 @@ TEST_F(GenerateShardCollectionInitialZonedChunksTest, EmptyTagsShouldFail) {
         ErrorCodes::InvalidOptions);
 }
 
+TEST_F(GenerateShardCollectionInitialZonedChunksTest, ZoneNotAssociatedWithAnyShardShouldFail) {
+    const std::vector<TagsType> tags = {
+        makeTag(ChunkRange(keyPattern().globalMin(), BSON(shardKey() << 0)), zoneName("0")),
+        makeTag(ChunkRange(BSON(shardKey() << 0), keyPattern().globalMax()), zoneName("1"))};
+    StringMap<std::vector<ShardId>> tagToShards = {{zoneName("0"), {ShardId("Shard0")}},
+                                                   {zoneName("1"), {}}};
+
+    ASSERT_THROWS_CODE(
+        InitialSplitPolicy::generateShardCollectionInitialZonedChunks(
+            nss(), shardKeyPattern(), timeStamp(), tags, tagToShards, makeShardIds(1)),
+        AssertionException,
+        50973);
+}
+
 }  // namespace
 }  // namespace mongo