SERVER-43914 Permit sharding commands to accept compound hashed shard key specs

3 files changed, 77 insertions, 4 deletions

diff --git a/src/mongo/s/shard_key_pattern.cpp b/src/mongo/s/shard_key_pattern.cpp
index 5e1773ce391..1a6093825ca 100644
--- a/src/mongo/s/shard_key_pattern.cpp
+++ b/src/mongo/s/shard_key_pattern.cpp
@@ -66,6 +66,7 @@ std::vector<std::unique_ptr<FieldRef>> parseShardKeyPattern(const BSONObj& keyPa
 
     std::vector<std::unique_ptr<FieldRef>> parsedPaths;
 
+    auto numHashedFields = 0;
     for (const auto& patternEl : keyPattern) {
         auto newFieldRef(std::make_unique<FieldRef>(patternEl.fieldNameStringData()));
 
@@ -88,7 +89,9 @@ std::vector<std::unique_ptr<FieldRef>> parseShardKeyPattern(const BSONObj& keyPa
                     !newFieldRef->getPart(i).empty());
         }
 
-        // Numeric and ascending (1.0), or "hashed" and single field
+        // Numeric and ascending (1.0), or "hashed" with exactly hashed field.
+        auto isHashedPattern = ShardKeyPattern::isHashedPatternEl(patternEl);
+        numHashedFields += isHashedPattern ? 1 : 0;
         uassert(ErrorCodes::BadValue,
                 str::stream()
                     << "Shard key " << keyPattern.toString()
@@ -96,7 +99,7 @@ std::vector<std::unique_ptr<FieldRef>> parseShardKeyPattern(const BSONObj& keyPa
                     << " or multiple numerical fields set to a value of 1. Failed to parse field "
                     << patternEl.fieldNameStringData(),
                 (patternEl.isNumber() && patternEl.numberInt() == 1) ||
-                    (keyPattern.nFields() == 1 && ShardKeyPattern::isHashedPatternEl(patternEl)));
+                    (isHashedPattern && numHashedFields == 1));
 
         parsedPaths.emplace_back(std::move(newFieldRef));
     }
@@ -173,10 +176,19 @@ Status ShardKeyPattern::checkShardKeyIsValidForMetadataStorage(const BSONObj& sh
     return Status::OK();
 }
 
+BSONElement ShardKeyPattern::extractHashedField(BSONObj keyPattern) {
+    for (auto&& element : keyPattern) {
+        if (isHashedPatternEl(element)) {
+            return element;
+        }
+    }
+    return BSONElement();
+}
 ShardKeyPattern::ShardKeyPattern(const BSONObj& keyPattern)
     : _keyPattern(keyPattern),
       _keyPatternPaths(parseShardKeyPattern(keyPattern)),
-      _hasId(keyPattern.hasField("_id"_sd)) {}
+      _hasId(keyPattern.hasField("_id"_sd)),
+      _hashedField(extractHashedField(keyPattern)) {}
 
 ShardKeyPattern::ShardKeyPattern(const KeyPattern& keyPattern)
     : ShardKeyPattern(keyPattern.toBSON()) {}
@@ -186,9 +198,30 @@ bool ShardKeyPattern::isHashedPatternEl(const BSONElement& el) {
 }
 
 bool ShardKeyPattern::isHashedPattern() const {
+    return !_hashedField.eoo();
+}
+
+bool ShardKeyPattern::isValidHashedValue(const BSONElement& el) {
+    switch (el.type()) {
+        case MinKey:
+        case MaxKey:
+        case NumberLong:
+            return true;
+        default:
+            return false;
+    }
+    MONGO_UNREACHABLE;
+}
+
+
+bool ShardKeyPattern::hasHashedPrefix() const {
     return isHashedPatternEl(_keyPattern.toBSON().firstElement());
 }
 
+BSONElement ShardKeyPattern::getHashedField() const {
+    return _hashedField;
+}
+
 const KeyPattern& ShardKeyPattern::getKeyPattern() const {
     return _keyPattern;
 }
@@ -345,7 +378,7 @@ BSONObj ShardKeyPattern::extractShardKeyFromQuery(const CanonicalQuery& query) c
         if (!isValidShardKeyElementForStorage(equalEl))
             return BSONObj();
 
-        if (isHashedPattern()) {
+        if (_hashedField && _hashedField.fieldNameStringData() == patternPath.dottedField()) {
             keyBuilder.append(
                 patternPath.dottedField(),
                 BSONElementHasher::hash64(equalEl, BSONElementHasher::DEFAULT_HASH_SEED));
diff --git a/src/mongo/s/shard_key_pattern.h b/src/mongo/s/shard_key_pattern.h
index 91eee10080a..c5baaabc407 100644
--- a/src/mongo/s/shard_key_pattern.h
+++ b/src/mongo/s/shard_key_pattern.h
@@ -88,8 +88,23 @@ public:
      */
     static bool isHashedPatternEl(const BSONElement& el);
 
+    /**
+     * Returns the BSONElement pointing to the hashed field. Returns empty BSONElement if not found.
+     */
+    static BSONElement extractHashedField(BSONObj keyPattern);
+
+    /**
+     * Check if the given BSONElement is of type 'MinKey', 'MaxKey' or 'NumberLong', which are the
+     * only acceptable values for hashed fields.
+     */
+    static bool isValidHashedValue(const BSONElement& el);
+
     bool isHashedPattern() const;
 
+    bool hasHashedPrefix() const;
+
+    BSONElement getHashedField() const;
+
     const KeyPattern& getKeyPattern() const;
 
     const std::vector<std::unique_ptr<FieldRef>>& getKeyPatternFields() const;
@@ -255,6 +270,7 @@ private:
     std::vector<std::unique_ptr<FieldRef>> _keyPatternPaths;
 
     bool _hasId;
+    BSONElement _hashedField;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/s/shard_key_pattern_test.cpp b/src/mongo/s/shard_key_pattern_test.cpp
index fd51f371f22..a1196e36e17 100644
--- a/src/mongo/s/shard_key_pattern_test.cpp
+++ b/src/mongo/s/shard_key_pattern_test.cpp
@@ -433,6 +433,19 @@ TEST(ShardKeyPattern, ExtractQueryShardKeyHashed) {
     ASSERT_BSONOBJ_EQ(queryKey(pattern, BSON("a" << BSON("c" << value))), BSONObj());
     ASSERT_BSONOBJ_EQ(queryKey(pattern, BSON("a" << BSON("b" << BSON_ARRAY(value)))), BSONObj());
     ASSERT_BSONOBJ_EQ(queryKey(pattern, BSON("a" << BSON_ARRAY(BSON("b" << value)))), BSONObj());
+
+    pattern = ShardKeyPattern(BSON("a.b"
+                                   << "hashed"
+                                   << "c.d" << 1));
+
+    ASSERT_BSONOBJ_EQ(queryKey(pattern, BSON("a.b" << value << "c.d" << value)),
+                      BSON("a.b" << hashValue << "c.d" << value));
+    ASSERT_BSONOBJ_EQ(
+        queryKey(pattern, fromjson("{a : {b: '12345', p : 1}, c : {d : '12345', q: 2}}")),
+        BSON("a.b" << hashValue << "c.d" << value));
+
+    ASSERT_BSONOBJ_EQ(queryKey(pattern, BSON("a.b" << value)), BSONObj());
+    ASSERT_BSONOBJ_EQ(queryKey(pattern, fromjson("{'a.b': [10], 'c.d': 1}")), BSONObj());
 }
 
 static bool indexComp(const ShardKeyPattern& pattern, const BSONObj& indexPattern) {
@@ -513,5 +526,16 @@ TEST(ShardKeyPattern, UniqueIndexCompatibleHashed) {
     ASSERT(!indexComp(pattern, BSON("c" << -1 << "a.b" << 1)));
 }
 
+TEST(ShardKeyPattern, IsHashedPattern) {
+    ASSERT(ShardKeyPattern(BSON("a.b"
+                                << "hashed"))
+               .isHashedPattern());
+    ASSERT(ShardKeyPattern(BSON("a.b" << 1 << "c"
+                                      << "hashed"
+                                      << "d" << 1))
+               .isHashedPattern());
+    ASSERT(!ShardKeyPattern(BSON("a.b" << 1 << "d" << 1)).isHashedPattern());
+}
+
 }  // namespace
 }  // namespace mongo