SERVER-69429 Check shard key restriction for indexes with `prepareUnique` option

(cherry picked from commit bef1987c88f68d932da97291dd95868b9dd54d2c)

7 files changed, 81 insertions, 29 deletions

diff --git a/jstests/core/index_prepareUnique.js b/jstests/core/index_prepareUnique.js
index ae889db8f6e..196e80a1864 100644
--- a/jstests/core/index_prepareUnique.js
+++ b/jstests/core/index_prepareUnique.js
@@ -1,8 +1,7 @@
 /**
- * Tests that the createIndex command accepts a prepareUnique field and works accordingly
- * then.
+ * Tests that the createIndex command accepts a prepareUnique field and works accordingly.
  *
- * @tags: [requires_fcv_53]
+ * @tags: [assumes_no_implicit_collection_creation_after_drop]
  */
 (function() {
 "use strict";
diff --git a/jstests/noPassthrough/sharding_index_uniqueness.js b/jstests/noPassthrough/sharding_index_uniqueness.js
new file mode 100644
index 00000000000..23cfcbd9a98
--- /dev/null
+++ b/jstests/noPassthrough/sharding_index_uniqueness.js
@@ -0,0 +1,49 @@
+/**
+ * Tests that an index with unique or prepareUnique option cannot be built on a sharded collection
+ * with an incompatible shard key, or the collection cannot be sharded with an incompatible
+ * unique/prepareUnique index.
+ *
+ * @tags: [
+ *   requires_sharding,
+ * ]
+ */
+
+(function() {
+"use strict";
+
+const st = new ShardingTest({shards: 1, mongos: 1});
+const dbName = "db";
+const db = st.getDB(dbName);
+const collNamePrefix = jsTestName();
+let count = 0;
+
+assert.commandWorked(st.s.adminCommand({enableSharding: dbName}));
+
+// Cannot shard the collection with a conflicting prepareUnique index.
+let coll = db[collNamePrefix + count++];
+assert.commandWorked(coll.createIndex({a: 1}, {prepareUnique: 1}));
+assert.commandFailedWithCode(st.s.adminCommand({shardCollection: coll.getFullName(), key: {b: 1}}),
+                             ErrorCodes.InvalidOptions);
+// Can shard the collection with a compatible prepareUnique index.
+assert.commandWorked(st.s.adminCommand({shardCollection: coll.getFullName(), key: {a: 1}}));
+
+// Cannot shard the collection with a conflicting unique index.
+coll = db[collNamePrefix + count++];
+assert.commandWorked(coll.createIndex({a: 1}, {unique: 1}));
+assert.commandFailedWithCode(st.s.adminCommand({shardCollection: coll.getFullName(), key: {b: 1}}),
+                             ErrorCodes.InvalidOptions);
+// Can shard the collection with a compatible unique index.
+assert.commandWorked(st.s.adminCommand({shardCollection: coll.getFullName(), key: {a: 1}}));
+
+// Cannot create the index with a conflicting shard key.
+coll = db[collNamePrefix + count++];
+assert.commandWorked(st.s.adminCommand({shardCollection: coll.getFullName(), key: {b: 1}}));
+assert.commandFailedWithCode(coll.createIndex({a: 1}, {prepareUnique: 1}),
+                             ErrorCodes.CannotCreateIndex);
+assert.commandFailedWithCode(coll.createIndex({a: 1}, {unique: 1}), ErrorCodes.CannotCreateIndex);
+// Can create the index with a compatible shard key.
+assert.commandWorked(coll.createIndex({b: 1, a: 1}, {prepareUnique: 1}));
+assert.commandWorked(coll.createIndex({b: 1, c: 1}, {unique: 1}));
+
+st.stop();
+})();
diff --git a/src/mongo/db/index_builds_coordinator.cpp b/src/mongo/db/index_builds_coordinator.cpp
index 8eb4ebb267f..42201897bb8 100644
--- a/src/mongo/db/index_builds_coordinator.cpp
+++ b/src/mongo/db/index_builds_coordinator.cpp
@@ -106,6 +106,7 @@ constexpr StringData kAbortIndexBuildFieldName = "abortIndexBuild"_sd;
 constexpr StringData kIndexesFieldName = "indexes"_sd;
 constexpr StringData kKeyFieldName = "key"_sd;
 constexpr StringData kUniqueFieldName = "unique"_sd;
+constexpr StringData kPrepareUniqueFieldName = "prepareUnique"_sd;
 
 /**
  * Checks if unique index specification is compatible with sharding configuration.
@@ -121,9 +122,9 @@ void checkShardKeyRestrictions(OperationContext* opCtx,
 
     const ShardKeyPattern shardKeyPattern(collDesc.getKeyPattern());
     uassert(ErrorCodes::CannotCreateIndex,
-            str::stream() << "cannot create unique index over " << newIdxKey
-                          << " with shard key pattern " << shardKeyPattern.toBSON(),
-            shardKeyPattern.isUniqueIndexCompatible(newIdxKey));
+            str::stream() << "cannot create index with 'unique' or 'prepareUnique' option over "
+                          << newIdxKey << " with shard key pattern " << shardKeyPattern.toBSON(),
+            shardKeyPattern.isIndexUniquenessCompatible(newIdxKey));
 }
 
 /**
@@ -2964,7 +2965,7 @@ std::vector<BSONObj> IndexBuildsCoordinator::prepareSpecListForCreate(
 
     // Verify that each spec is compatible with the collection's sharding state.
     for (const BSONObj& spec : resultSpecs) {
-        if (spec[kUniqueFieldName].trueValue()) {
+        if (spec[kUniqueFieldName].trueValue() || spec[kPrepareUniqueFieldName].trueValue()) {
             checkShardKeyRestrictions(opCtx, nss, spec[kKeyFieldName].Obj());
         }
     }
diff --git a/src/mongo/db/s/shard_key_util.cpp b/src/mongo/db/s/shard_key_util.cpp
index 5a0acaeb2a4..4998154ae26 100644
--- a/src/mongo/db/s/shard_key_util.cpp
+++ b/src/mongo/db/s/shard_key_util.cpp
@@ -114,12 +114,14 @@ bool validShardKeyIndexExists(OperationContext* opCtx,
     for (const auto& idx : indexes) {
         BSONObj currentKey = idx["key"].embeddedObject();
         bool isUnique = idx["unique"].trueValue();
+        bool isPrepareUnique = idx["prepareUnique"].trueValue();
         uassert(ErrorCodes::InvalidOptions,
                 str::stream() << "can't shard collection '" << nss.ns() << "' with unique index on "
                               << currentKey << " and proposed shard key "
                               << shardKeyPattern.toBSON()
                               << ". Uniqueness can't be maintained unless shard key is a prefix",
-                !isUnique || shardKeyPattern.isUniqueIndexCompatible(currentKey));
+                (!isUnique && !isPrepareUnique) ||
+                    shardKeyPattern.isIndexUniquenessCompatible(currentKey));
     }
 
     // 2. Check for a useful index
diff --git a/src/mongo/s/shard_key_pattern.cpp b/src/mongo/s/shard_key_pattern.cpp
index 7bfb9c88237..c693528de08 100644
--- a/src/mongo/s/shard_key_pattern.cpp
+++ b/src/mongo/s/shard_key_pattern.cpp
@@ -545,12 +545,12 @@ BSONObj ShardKeyPattern::extractShardKeyFromQuery(const CanonicalQuery& query) c
     return keyBuilder.obj();
 }
 
-bool ShardKeyPattern::isUniqueIndexCompatible(const BSONObj& uniqueIndexPattern) const {
-    if (!uniqueIndexPattern.isEmpty() && uniqueIndexPattern.firstElementFieldName() == kIdField) {
+bool ShardKeyPattern::isIndexUniquenessCompatible(const BSONObj& indexPattern) const {
+    if (!indexPattern.isEmpty() && indexPattern.firstElementFieldName() == kIdField) {
         return true;
     }
 
-    return _keyPattern.toBSON().isFieldNamePrefixOf(uniqueIndexPattern);
+    return _keyPattern.toBSON().isFieldNamePrefixOf(indexPattern);
 }
 
 BoundList ShardKeyPattern::flattenBounds(const IndexBounds& indexBounds) const {
diff --git a/src/mongo/s/shard_key_pattern.h b/src/mongo/s/shard_key_pattern.h
index 545d1906e31..70f5aff336d 100644
--- a/src/mongo/s/shard_key_pattern.h
+++ b/src/mongo/s/shard_key_pattern.h
@@ -274,36 +274,37 @@ public:
     BSONObj extractShardKeyFromQuery(const CanonicalQuery& query) const;
 
     /**
-     * Returns true if the shard key pattern can ensure that the unique index pattern is
-     * respected across all shards.
+     * Returns true if the shard key pattern can ensure that the index uniqueness is respected
+     * across all shards.
      *
      * Primarily this just checks whether the shard key pattern field names are equal to or a
-     * prefix of the unique index pattern field names.  Since documents with the same fields in
-     * the shard key pattern are guaranteed to go to the same shard, and all documents must
-     * contain the full shard key, a unique index with a shard key pattern prefix can be sure
-     * when resolving duplicates that documents on other shards will have different shard keys,
-     * and so are not duplicates.
+     * prefix of the 'unique' or 'prepareUnique' index pattern field names. Since documents with the
+     * same fields in the shard key pattern are guaranteed to go to the same shard, and all
+     * documents must contain the full shard key, an index with {unique: true} or {prepareUnique:
+     * true} and a shard key pattern prefix can be sure when resolving duplicates that documents on
+     * other shards will have different shard keys, and so are not duplicates.
      *
      * Hashed shard key patterns are similar to ordinary patterns in that they guarantee similar
      * shard keys go to the same shard.
      *
      * Examples:
-     *     shard key {a : 1} is compatible with a unique index on {_id : 1}
-     *     shard key {a : 1} is compatible with a unique index on {a : 1 , b : 1}
-     *     shard key {a : 1} is compatible with a unique index on {a : -1 , b : 1 }
-     *     shard key {a : "hashed"} is compatible with a unique index on {a : 1}
-     *     shard key {a : 1} is not compatible with a unique index on {b : 1}
-     *     shard key {a : "hashed" , b : 1 } is not compatible with unique index on { b : 1 }
+     *     shard key {a : 1} is compatible with a unique/prepareUnique index on {_id : 1}
+     *     shard key {a : 1} is compatible with a unique/prepareUnique index on {a : 1, b : 1}
+     *     shard key {a : 1} is compatible with a unique/prepareUnique index on {a : -1, b : 1}
+     *     shard key {a : "hashed"} is compatible with a unique/prepareUnique index on {a : 1}
+     *     shard key {a : 1} is not compatible with a unique/prepareUnique index on {b : 1}
+     *     shard key {a : "hashed", b : 1} is not compatible with unique/prepareUnique index on
+     *     {b : 1}
      *
      * All unique index patterns starting with _id are assumed to be enforceable by the fact
      * that _ids must be unique, and so all unique _id prefixed indexes are compatible with
      * any shard key pattern.
      *
-     * NOTE: We assume 'uniqueIndexPattern' is a valid unique index pattern - a pattern like
-     * { k : "hashed" } is not capable of being a unique index and is an invalid argument to
-     * this method.
+     * NOTE: We assume 'indexPattern' is a valid unique/prepareUnique index pattern - a pattern like
+     * { k : "hashed" } is not capable of being a unique/prepareUnique index and is an invalid
+     * argument to this method.
      */
-    bool isUniqueIndexCompatible(const BSONObj& uniqueIndexPattern) const;
+    bool isIndexUniquenessCompatible(const BSONObj& indexPattern) const;
 
     /**
      * Return an ordered list of bounds generated using this KeyPattern and the
diff --git a/src/mongo/s/shard_key_pattern_test.cpp b/src/mongo/s/shard_key_pattern_test.cpp
index b7c2a3e24b1..67ed42a41ad 100644
--- a/src/mongo/s/shard_key_pattern_test.cpp
+++ b/src/mongo/s/shard_key_pattern_test.cpp
@@ -608,7 +608,7 @@ TEST_F(ShardKeyPatternTest, ExtractQueryShardKeyHashed) {
 }
 
 static bool indexComp(const ShardKeyPattern& pattern, const BSONObj& indexPattern) {
-    return pattern.isUniqueIndexCompatible(indexPattern);
+    return pattern.isIndexUniquenessCompatible(indexPattern);
 }
 
 TEST_F(ShardKeyPatternTest, UniqueIndexCompatibleSingle) {