SERVER-47436 Make shards validate shardKey in dataSize command

buildscripts/resmokeconfig/suites/sharding_last_stable_mongos_and_mixed_shards.yml

@@ -32,6 +32,7 @@ selector:
   - jstests/sharding/verify_sessions_expiration_sharded.js
   # Enable when 4.6 becomes last stable
   - jstests/sharding/disable_resumable_range_deleter_flag.js
+  - jstests/sharding/mongos_dataSize.js
   # Enable when SERVER-44733 is backported
   - jstests/sharding/change_streams_update_lookup_shard_metadata_missing.js
   # Enable when SERVER-43310 is backported

jstests/sharding/mongos_dataSize.js

@@ -0,0 +1,113 @@
+/*
+ * Tests the dataSize command on mongos.
+ */
+(function() {
+'use strict';
+
+const kDbName = "foo";
+const kCollName = "bar";
+const kNs = kDbName + "." + kCollName;
+const kNumDocs = 100;
+
+/*
+ * Returns the global min and max key for the given key pattern.
+ */
+function getGlobalMinMaxKey(keyPattern) {
+    let globalMin = {};
+    let globalMax = {};
+    for (let field in keyPattern) {
+        globalMin[field] = MinKey;
+        globalMax[field] = MaxKey;
+    }
+    return {globalMin, globalMax};
+}
+
+/*
+ * Runs a dataSize command with the given key pattern on the given connection, and
+ * asserts that command works and that the returned numObjects is equal to expectedNumObjects.
+ */
+function assertDataSizeCmdWorked(conn, keyPattern, expectedNumObjects) {
+    let res = assert.commandWorked(conn.adminCommand({dataSize: kNs, keyPattern: keyPattern}));
+    assert.eq(res.numObjects, expectedNumObjects);
+
+    const {globalMin, globalMax} = getGlobalMinMaxKey(keyPattern);
+    res = assert.commandWorked(
+        conn.adminCommand({dataSize: kNs, keyPattern: keyPattern, min: globalMin, max: globalMax}));
+    assert.eq(res.numObjects, expectedNumObjects);
+}
+
+/*
+ * Runs a dataSize command with the given key pattern on the given connection, and
+ * asserts that command failed with BadValue error.
+ */
+function assertDataSizeCmdFailedWithBadValue(conn, keyPattern) {
+    assert.commandFailedWithCode(conn.adminCommand({dataSize: kNs, keyPattern: keyPattern}),
+                                 ErrorCodes.BadValue);
+}
+
+/*
+ * Runs dataSize commands on the given connection. Asserts the command fails if run with an invalid
+ * namespace or with the min and max as given by each range in invalidRanges. Asserts that the
+ * command succeeds if run with valid min and max and returns the expected numObjects.
+ */
+function testDataSizeCmd(conn, keyPattern, invalidRanges, numObjects) {
+    assert.commandFailedWithCode(conn.adminCommand({dataSize: kCollName}),
+                                 ErrorCodes.InvalidNamespace);
+
+    for (const {min, max, errorCode} of invalidRanges) {
+        const cmdObj = {dataSize: kNs, keyPattern: keyPattern, min: min, max: max};
+        assert.commandFailedWithCode(conn.adminCommand(cmdObj), errorCode);
+    }
+
+    assertDataSizeCmdWorked(conn, {}, numObjects);
+    assertDataSizeCmdWorked(conn, keyPattern, numObjects);
+}
+
+const st = new ShardingTest({mongos: 3, shards: 2});
+assert.commandWorked(st.s.adminCommand({enableSharding: kDbName}));
+st.ensurePrimaryShard(kDbName, st.shard0.shardName);
+
+const shardKey1 = {
+    x: 1
+};
+jsTest.log(`Sharding the collection with key ${tojson(shardKey1)}`);
+assert.commandWorked(st.s0.adminCommand({shardCollection: kNs, key: shardKey1}));
+
+let bulk = st.s0.getCollection(kNs).initializeUnorderedBulkOp();
+for (let i = 0; i < kNumDocs; ++i) {
+    bulk.insert({_id: i, x: i, y: -i});
+}
+assert.commandWorked(bulk.execute());
+
+jsTest.log("Verify that keyPattern and key range validation works");
+const invalidRanges1 = [
+    {min: {y: MinKey}, max: {y: MaxKey}, errorCode: ErrorCodes.BadValue},
+    {min: {x: MinKey, y: MinKey}, max: {x: MaxKey, y: MaxKey}, errorCode: ErrorCodes.BadValue},
+    // The command does not throw any particular error when only one of min or max is specified.
+    {min: {}, max: {x: MaxKey}, errorCode: ErrorCodes.UnknownError},
+    {min: {x: MinKey}, max: {}, errorCode: ErrorCodes.UnknownError},
+];
+testDataSizeCmd(st.s0, shardKey1, invalidRanges1, kNumDocs);
+testDataSizeCmd(st.s1, shardKey1, invalidRanges1, kNumDocs);
+testDataSizeCmd(st.s2, shardKey1, invalidRanges1, kNumDocs);
+
+jsTest.log("Dropping the collection");
+st.s0.getCollection(kNs).drop();
+
+const shardKey2 = {
+    y: 1
+};
+jsTest.log(`Resharding the collection with key ${tojson(shardKey2)}`);
+assert.commandWorked(st.s0.adminCommand({shardCollection: kNs, key: shardKey2}));
+
+jsTest.log("Verify that validation occurs on shards not on mongos");
+// If validation occurs on mongos, the command would fail with BadValue as st.s1 is stale so
+// it thinks that shardKey1 is the shard key.
+assertDataSizeCmdWorked(st.s1, shardKey2, 0);
+
+// If validation occurs on mongos, the command would succeed as st.s2 is stale so it thinks
+// that shardKey1 is the shard key.
+assertDataSizeCmdFailedWithBadValue(st.s2, shardKey1);
+
+st.stop();
+})();

jstests/sharding/mongos_dataSize_test.js

@@ -1,14 +0,0 @@
-// This tests the command dataSize on sharded clusters to ensure that they can use the command.
-
-(function() {
-'use strict';
-
-let s = new ShardingTest({shards: 2, mongos: 1});
-let db = s.getDB("test");
-assert.commandWorked(s.s0.adminCommand({enableSharding: "test"}));
-assert.commandWorked(s.s0.adminCommand({shardcollection: "test.foo", key: {num: 1}}));
-assert.commandWorked(s.getPrimaryShard("test").getDB("admin").runCommand({datasize: "test.foo"}));
-assert.commandFailedWithCode(s.getPrimaryShard("test").getDB("admin").runCommand({datasize: "foo"}),
-                             ErrorCodes.InvalidNamespace);
-s.stop();
-})();

src/mongo/db/commands/dbcommands.cpp

@@ -452,9 +452,31 @@ public:
         BSONObj keyPattern = jsobj.getObjectField("keyPattern");
         bool estimate = jsobj["estimate"].trueValue();
 
-        AutoGetCollectionForReadCommand ctx(opCtx, NamespaceString(ns));
-
+        const NamespaceString nss(ns);
+        AutoGetCollectionForReadCommand ctx(opCtx, nss);
         Collection* collection = ctx.getCollection();
+
+        const auto collDesc = CollectionShardingState::get(opCtx, nss)->getCollectionDescription();
+
+        if (collDesc.isSharded()) {
+            const ShardKeyPattern shardKeyPattern(collDesc.getKeyPattern());
+            uassert(ErrorCodes::BadValue,
+                    "keyPattern must be empty or must be an object that equals the shard key",
+                    keyPattern.isEmpty() ||
+                        (SimpleBSONObjComparator::kInstance.evaluate(shardKeyPattern.toBSON() ==
+                                                                     keyPattern)));
+
+            uassert(ErrorCodes::BadValue,
+                    str::stream() << "min value " << min << " does not have shard key",
+                    min.isEmpty() || shardKeyPattern.isShardKey(min));
+            min = shardKeyPattern.normalizeShardKey(min);
+
+            uassert(ErrorCodes::BadValue,
+                    str::stream() << "max value " << max << " does not have shard key",
+                    max.isEmpty() || shardKeyPattern.isShardKey(max));
+            max = shardKeyPattern.normalizeShardKey(max);
+        }
+
         long long numRecords = 0;
         if (collection) {
             numRecords = collection->numRecords(opCtx);

src/mongo/s/commands/cluster_data_size_cmd.cpp

@@ -75,31 +75,6 @@ public:
 
         auto routingInfo =
             uassertStatusOK(Grid::get(opCtx)->catalogCache()->getCollectionRoutingInfo(opCtx, nss));
-        const auto cm = routingInfo.cm();
-
-        BSONObj min = cmdObj.getObjectField("min");
-        BSONObj max = cmdObj.getObjectField("max");
-
-        if (cm) {
-            auto keyPattern = cmdObj["keyPattern"];
-
-            uassert(ErrorCodes::BadValue,
-                    "keyPattern must be empty or must be an object that equals the shard key",
-                    !keyPattern ||
-                        (keyPattern.type() == Object &&
-                         SimpleBSONObjComparator::kInstance.evaluate(
-                             cm->getShardKeyPattern().toBSON() == keyPattern.Obj())));
-
-            uassert(ErrorCodes::BadValue,
-                    str::stream() << "min value " << min << " does not have shard key",
-                    cm->getShardKeyPattern().isShardKey(min));
-            min = cm->getShardKeyPattern().normalizeShardKey(min);
-
-            uassert(ErrorCodes::BadValue,
-                    str::stream() << "max value " << max << " does not have shard key",
-                    cm->getShardKeyPattern().isShardKey(max));
-            max = cm->getShardKeyPattern().normalizeShardKey(max);
-        }
 
         auto shardResults = scatterGatherVersionedTargetByRoutingTable(
             opCtx,