SERVER-36424 Enforce uniqueKey is unique enough

12 files changed, 433 insertions, 3 deletions

diff --git a/buildscripts/resmokeconfig/suites/aggregation_mongos_passthrough.yml b/buildscripts/resmokeconfig/suites/aggregation_mongos_passthrough.yml
index 14151b4ea4f..69271f5ab14 100644
--- a/buildscripts/resmokeconfig/suites/aggregation_mongos_passthrough.yml
+++ b/buildscripts/resmokeconfig/suites/aggregation_mongos_passthrough.yml
@@ -10,6 +10,8 @@ selector:
   - jstests/aggregation/bugs/cursor_timeout.js
   - jstests/aggregation/bugs/lookup_unwind_getmore.js
   - jstests/aggregation/bugs/lookup_unwind_killcursor.js
+  # TODO: Remove when SERVER-36047 is resolved.
+  - jstests/aggregation/sources/out/unique_key_requires_index.js
   # TODO: Remove when SERVER-23229 is fixed.
   - jstests/aggregation/bugs/groupMissing.js
   exclude_with_any_tags:
diff --git a/buildscripts/resmokeconfig/suites/aggregation_one_shard_sharded_collections.yml b/buildscripts/resmokeconfig/suites/aggregation_one_shard_sharded_collections.yml
index 3f31669bd9f..32df6505ee0 100644
--- a/buildscripts/resmokeconfig/suites/aggregation_one_shard_sharded_collections.yml
+++ b/buildscripts/resmokeconfig/suites/aggregation_one_shard_sharded_collections.yml
@@ -10,6 +10,8 @@ selector:
   - jstests/aggregation/bugs/cursor_timeout.js
   - jstests/aggregation/bugs/lookup_unwind_getmore.js
   - jstests/aggregation/bugs/lookup_unwind_killcursor.js
+  # TODO: Remove when SERVER-36047 is resolved.
+  - jstests/aggregation/sources/out/unique_key_requires_index.js
   # TODO: Remove when SERVER-23229 is fixed.
   - jstests/aggregation/bugs/groupMissing.js
   exclude_with_any_tags:
diff --git a/buildscripts/resmokeconfig/suites/aggregation_sharded_collections_passthrough.yml b/buildscripts/resmokeconfig/suites/aggregation_sharded_collections_passthrough.yml
index 83817f9039e..237b1344c35 100644
--- a/buildscripts/resmokeconfig/suites/aggregation_sharded_collections_passthrough.yml
+++ b/buildscripts/resmokeconfig/suites/aggregation_sharded_collections_passthrough.yml
@@ -21,6 +21,8 @@ selector:
   - jstests/aggregation/use_query_project_and_sort.js
   - jstests/aggregation/use_query_projection.js
   - jstests/aggregation/use_query_sort.js
+  # TODO: Remove when SERVER-36047 is resolved.
+  - jstests/aggregation/sources/out/unique_key_requires_index.js
   # TODO: Remove when SERVER-23229 is fixed.
   - jstests/aggregation/bugs/groupMissing.js
   exclude_with_any_tags:
diff --git a/jstests/aggregation/sources/out/mode_replace_documents.js b/jstests/aggregation/sources/out/mode_replace_documents.js
index a147688adda..a6c723c4c47 100644
--- a/jstests/aggregation/sources/out/mode_replace_documents.js
+++ b/jstests/aggregation/sources/out/mode_replace_documents.js
@@ -30,8 +30,8 @@
     // Test 'replaceDocuments' mode with a dotted path unique key.
     coll.drop();
     outColl.drop();
-    assert.commandWorked(coll.insert({_id: 0, a: {b: 1}}));
-    assert.commandWorked(coll.insert({_id: 1, a: {b: 1}, c: 1}));
+    assert.commandWorked(coll.insert([{_id: 0, a: {b: 1}}, {_id: 1, a: {b: 1}, c: 1}]));
+    assert.commandWorked(outColl.createIndex({"a.b": 1, _id: 1}, {unique: true}));
     coll.aggregate([
         {$addFields: {_id: 0}},
         {$out: {to: outColl.getName(), mode: "replaceDocuments", uniqueKey: {_id: 1, "a.b": 1}}}
@@ -52,10 +52,12 @@
                     50905);
 
     // Test that 'replaceDocuments' mode with a missing non-id unique key fails.
+    assert.commandWorked(outColl.createIndex({missing: 1}, {unique: true}));
     assertErrorCode(
         coll,
         [{$out: {to: outColl.getName(), mode: "replaceDocuments", uniqueKey: {missing: 1}}}],
-        50905);
+        50905  // This attempt should fail because there's no field 'missing' in the document.
+        );
 
     // Test that a replace fails to insert a document if it violates a unique index constraint. In
     // this example, $out will attempt to insert multiple documents with {a: 0} which is not allowed
@@ -73,6 +75,7 @@
     // Test that $out fails if the unique key contains an array.
     coll.drop();
     assert.commandWorked(coll.insert({_id: 0, a: [1, 2]}));
+    assert.commandWorked(outColl.createIndex({"a.b": 1, _id: 1}, {unique: true}));
     assertErrorCode(
         coll,
         [
diff --git a/jstests/aggregation/sources/out/unique_key_requires_index.js b/jstests/aggregation/sources/out/unique_key_requires_index.js
new file mode 100644
index 00000000000..4e183a2d0bc
--- /dev/null
+++ b/jstests/aggregation/sources/out/unique_key_requires_index.js
@@ -0,0 +1,319 @@
+// Tests that the $out stage enforces that the uniqueKey argument can be used to uniquely identify
+// documents by checking that there is a supporting unique, non-partial, collator-compatible index
+// in the index catalog.
+(function() {
+    "use strict";
+
+    const testDB = db.getSiblingDB("unique_key_requires_index");
+    const source = testDB.source;
+    source.drop();
+    assert.commandWorked(source.insert([{_id: 0, a: 0}, {_id: 1, a: 1}]));
+
+    function withEachOutMode(callback) {
+        callback("replaceCollection");
+        callback("insertDocuments");
+        callback("replaceDocuments");
+    }
+
+    // Test that using {_id: 1} or not providing a unique key does not require any special indexes.
+    (function simpleIdUniqueKeyOrDefaultShouldNotRequireIndexes() {
+        function assertDefaultUniqueKeySuceeds({setupCallback, collName}) {
+            // Legacy style $out - "replaceCollection".
+            setupCallback();
+            assert.doesNotThrow(() => source.aggregate([{$out: collName}]));
+
+            withEachOutMode((mode) => {
+                setupCallback();
+                assert.doesNotThrow(() => source.aggregate([{$out: {to: collName, mode: mode}}]));
+                setupCallback();
+                assert.doesNotThrow(() => source.aggregate(
+                                        [{$out: {to: collName, uniqueKey: {_id: 1}, mode: mode}}]));
+            });
+        }
+
+        // Test that using {_id: 1} or not specifying a uniqueKey works for a collection which does
+        // not exist.
+        const non_existent = testDB.non_existent;
+        assertDefaultUniqueKeySuceeds(
+            {setupCallback: () => non_existent.drop(), collName: non_existent.getName()});
+
+        const unindexed = testDB.unindexed;
+        assertDefaultUniqueKeySuceeds({
+            setupCallback: () => {
+                unindexed.drop();
+                assert.commandWorked(testDB.runCommand({create: unindexed.getName()}));
+            },
+            collName: unindexed.getName()
+        });
+    }());
+
+    function assertUniqueKeyIsInvalid({uniqueKey, targetColl, options}) {
+        let cmd = {
+            aggregate: source.getName(),
+            pipeline: [{$out: {mode: "replaceDocuments", uniqueKey: uniqueKey, to: targetColl}}],
+            cursor: {}
+        };
+        withEachOutMode((mode) => {
+            cmd.pipeline[0].$out.mode = mode;
+            assert.commandFailedWithCode(testDB.runCommand(Object.merge(cmd, options)), 50938);
+        });
+    }
+
+    // Test that a unique index on the unique key can be used to satisfy the requirement.
+    (function basicUniqueIndexWorks() {
+        const target = testDB.regular_unique;
+        target.drop();
+        assertUniqueKeyIsInvalid({uniqueKey: {_id: 1, a: 1}, targetColl: target.getName()});
+
+        assert.commandWorked(target.createIndex({a: 1, _id: 1}, {unique: true}));
+        assert.doesNotThrow(() => source.aggregate([
+            {$out: {to: target.getName(), mode: "replaceDocuments", uniqueKey: {_id: 1, a: 1}}}
+        ]));
+        assert.doesNotThrow(() => source.aggregate([
+            {$out: {to: target.getName(), mode: "replaceDocuments", uniqueKey: {a: 1, _id: 1}}}
+        ]));
+
+        assertUniqueKeyIsInvalid({uniqueKey: {_id: 1, a: 1, b: 1}, targetColl: target.getName()});
+        assertUniqueKeyIsInvalid({uniqueKey: {a: 1, b: 1}, targetColl: target.getName()});
+        assertUniqueKeyIsInvalid({uniqueKey: {b: 1}, targetColl: target.getName()});
+        assertUniqueKeyIsInvalid({uniqueKey: {a: 1}, targetColl: target.getName()});
+        assertUniqueKeyIsInvalid({uniqueKey: {}, targetColl: target.getName()});
+
+        assert.commandWorked(target.dropIndex({a: 1, _id: 1}));
+        assert.commandWorked(target.createIndex({a: 1}, {unique: true}));
+        assert.doesNotThrow(
+            () => source.aggregate(
+                [{$out: {to: target.getName(), mode: "replaceDocuments", uniqueKey: {a: 1}}}]));
+
+        // Create a non-unique index and make sure that doesn't work.
+        assert.commandWorked(target.dropIndex({a: 1}));
+        assert.commandWorked(target.createIndex({a: 1}));
+        assertUniqueKeyIsInvalid({uniqueKey: {a: 1}, targetColl: target.getName()});
+        assertUniqueKeyIsInvalid({uniqueKey: {a: 1, _id: 1}, targetColl: target.getName()});
+    }());
+
+    // Test that a unique index on the unique key cannot be used to satisfy the requirement if it is
+    // a partial index.
+    (function uniqueButPartialShouldNotWork() {
+        const target = testDB.unique_but_partial_indexes;
+        target.drop();
+        assertUniqueKeyIsInvalid({uniqueKey: {a: 1}, targetColl: target.getName()});
+
+        assert.commandWorked(
+            target.createIndex({a: 1}, {unique: true, partialFilterExpression: {a: {$gte: 2}}}));
+        assertUniqueKeyIsInvalid({uniqueKey: {a: 1}, targetColl: target.getName()});
+        assertUniqueKeyIsInvalid({uniqueKey: {a: 1, _id: 1}, targetColl: target.getName()});
+    }());
+
+    // Test that a unique index on the unique key cannot be used to satisfy the requirement if it
+    // has a different collation.
+    (function indexMustMatchCollationOfOperation() {
+        const target = testDB.collation_indexes;
+        target.drop();
+        assertUniqueKeyIsInvalid({uniqueKey: {a: 1}, targetColl: target.getName()});
+
+        assert.commandWorked(
+            target.createIndex({a: 1}, {unique: true, collation: {locale: "en_US"}}));
+        assertUniqueKeyIsInvalid({uniqueKey: {a: 1}, targetColl: target.getName()});
+        assertUniqueKeyIsInvalid({
+            uniqueKey: {a: 1},
+            targetColl: target.getName(),
+            options: {collation: {locale: "en"}}
+        });
+        assertUniqueKeyIsInvalid({
+            uniqueKey: {a: 1},
+            targetColl: target.getName(),
+            options: {collation: {locale: "simple"}}
+        });
+        assertUniqueKeyIsInvalid({
+            uniqueKey: {a: 1},
+            targetColl: target.getName(),
+            options: {collation: {locale: "en_US", strength: 1}}
+        });
+        assert.doesNotThrow(
+            () => source.aggregate(
+                [{$out: {to: target.getName(), mode: "replaceDocuments", uniqueKey: {a: 1}}}],
+                {collation: {locale: "en_US"}}));
+
+        // Test that a non-unique index with the same collation cannot be used.
+        assert.commandWorked(target.dropIndex({a: 1}));
+        assert.commandWorked(target.createIndex({a: 1}, {collation: {locale: "en_US"}}));
+        assertUniqueKeyIsInvalid({
+            uniqueKey: {a: 1},
+            targetColl: target.getName(),
+            options: {collation: {locale: "en_US"}}
+        });
+
+        // Test that a collection-default collation will be applied to the index, but not the $out's
+        // update or insert into that collection. The pipeline will inherit a collection-default
+        // collation, but from the source collection, not the $out's target collection.
+        target.drop();
+        assert.commandWorked(
+            testDB.runCommand({create: target.getName(), collation: {locale: "en_US"}}));
+        assert.commandWorked(target.createIndex({a: 1}, {unique: true}));
+        assertUniqueKeyIsInvalid({
+            uniqueKey: {a: 1},
+            targetColl: target.getName(),
+        });
+        assert.doesNotThrow(
+            () => source.aggregate(
+                [{$out: {to: target.getName(), mode: "replaceDocuments", uniqueKey: {a: 1}}}],
+                {collation: {locale: "en_US"}}));
+
+        // Test that when the source collection and foreign collection have the same default
+        // collation, a unique index on the foreign collection can be used.
+        const newSourceColl = testDB.new_source;
+        newSourceColl.drop();
+        assert.commandWorked(
+            testDB.runCommand({create: newSourceColl.getName(), collation: {locale: "en_US"}}));
+        assert.commandWorked(newSourceColl.insert([{_id: 1, a: 1}, {_id: 2, a: 2}]));
+        // This aggregate does not specify a collation, but it should inherit the default collation
+        // from 'newSourceColl', and therefor the index on 'target' should be eligible for use since
+        // it has the same collation.
+        assert.doesNotThrow(
+            () => newSourceColl.aggregate(
+                [{$out: {to: target.getName(), mode: "replaceDocuments", uniqueKey: {a: 1}}}]));
+
+        // Test that an explicit "simple" collation can be used with an index without a collation.
+        target.drop();
+        assert.commandWorked(target.createIndex({a: 1}, {unique: true}));
+        assert.doesNotThrow(
+            () => source.aggregate(
+                [{$out: {to: target.getName(), mode: "replaceDocuments", uniqueKey: {a: 1}}}],
+                {collation: {locale: "simple"}}));
+        assertUniqueKeyIsInvalid({
+            uniqueKey: {a: 1},
+            targetColl: target.getName(),
+            options: {collation: {locale: "en_US"}}
+        });
+    }());
+
+    // Test that a unique index which is not simply ascending/descending fields cannot be used for
+    // the uniqueKey
+    (function testSpecialIndexTypes() {
+        const target = testDB.special_index_types;
+        target.drop();
+
+        assert.commandWorked(target.createIndex({a: 1, text: "text"}, {unique: true}));
+        assertUniqueKeyIsInvalid({uniqueKey: {a: 1}, targetColl: target.getName()});
+        assertUniqueKeyIsInvalid({uniqueKey: {a: 1, text: 1}, targetColl: target.getName()});
+        assertUniqueKeyIsInvalid({uniqueKey: {text: 1}, targetColl: target.getName()});
+
+        target.drop();
+        assert.commandWorked(target.createIndex({a: 1, geo: "2dsphere"}, {unique: true}));
+        assertUniqueKeyIsInvalid({uniqueKey: {a: 1}, targetColl: target.getName()});
+        assertUniqueKeyIsInvalid({uniqueKey: {a: 1, geo: 1}, targetColl: target.getName()});
+        assertUniqueKeyIsInvalid({uniqueKey: {geo: 1, a: 1}, targetColl: target.getName()});
+
+        target.drop();
+        assert.commandWorked(target.createIndex({geo: "2d"}, {unique: true}));
+        assertUniqueKeyIsInvalid({uniqueKey: {a: 1, geo: 1}, targetColl: target.getName()});
+        assertUniqueKeyIsInvalid({uniqueKey: {geo: 1}, targetColl: target.getName()});
+
+        target.drop();
+        assert.commandWorked(
+            target.createIndex({geo: "geoHaystack", a: 1}, {unique: true, bucketSize: 5}));
+        assertUniqueKeyIsInvalid({uniqueKey: {a: 1, geo: 1}, targetColl: target.getName()});
+        assertUniqueKeyIsInvalid({uniqueKey: {geo: 1, a: 1}, targetColl: target.getName()});
+
+        target.drop();
+        // MongoDB does not support unique hashed indexes.
+        assert.commandFailedWithCode(target.createIndex({a: "hashed"}, {unique: true}), 16764);
+        assert.commandWorked(target.createIndex({a: "hashed"}));
+        assertUniqueKeyIsInvalid({uniqueKey: {a: 1}, targetColl: target.getName()});
+    }());
+
+    // Test that a unique index with dotted field names can be used.
+    (function testDottedFieldNames() {
+        const target = testDB.dotted_field_paths;
+        target.drop();
+
+        assert.commandWorked(target.createIndex({a: 1, "b.c.d": -1}, {unique: true}));
+        assertUniqueKeyIsInvalid({uniqueKey: {a: 1}, targetColl: target.getName()});
+        assert.doesNotThrow(() => source.aggregate([
+            {$project: {_id: 1, a: 1, b: {c: {d: "x"}}}},
+            {
+              $out: {
+                  to: target.getName(),
+                  mode: "replaceDocuments",
+                  uniqueKey: {a: 1, "b.c.d": 1}
+              }
+            }
+        ]));
+
+        target.drop();
+        assert.commandWorked(target.createIndex({"id.x": 1, "id.y": -1}, {unique: true}));
+        assert.doesNotThrow(() => source.aggregate([
+            {$group: {_id: {x: "$_id", y: "$a"}}},
+            {$project: {id: "$_id"}},
+            {
+              $out: {
+                  to: target.getName(),
+                  mode: "replaceDocuments",
+                  uniqueKey: {"id.x": 1, "id.y": 1}
+              }
+            }
+        ]));
+        assert.doesNotThrow(() => source.aggregate([
+            {$group: {_id: {x: "$_id", y: "$a"}}},
+            {$project: {id: "$_id"}},
+            {
+              $out: {
+                  to: target.getName(),
+                  mode: "replaceDocuments",
+                  uniqueKey: {"id.y": 1, "id.x": 1}
+              }
+            }
+        ]));
+
+        // Test that we cannot use arrays with a dotted path within an $out.
+        target.drop();
+        assert.commandWorked(target.createIndex({"b.c": 1}, {unique: true}));
+        withEachOutMode((mode) => {
+            assert.commandFailedWithCode(testDB.runCommand({
+                aggregate: source.getName(),
+                pipeline: [
+                    {$replaceRoot: {newRoot: {b: [{c: 1}, {c: 2}]}}},
+                    {$out: {to: target.getName(), mode: mode, uniqueKey: {"b.c": 1}}}
+                ],
+                cursor: {}
+            }),
+                                         50905);
+        });
+    }());
+
+    // Test that a unique index that is multikey can still be used.
+    (function testMultikeyIndex() {
+        const target = testDB.multikey_index;
+        target.drop();
+
+        assert.commandWorked(target.createIndex({"a.b": 1}, {unique: true}));
+        assert.doesNotThrow(() => source.aggregate([
+            {$project: {_id: 1, "a.b": "$a"}},
+            {$out: {to: target.getName(), mode: "replaceDocuments", uniqueKey: {"a.b": 1}}}
+        ]));
+        assert.commandWorked(target.insert({_id: "TARGET", a: [{b: "hi"}, {b: "hello"}]}));
+        assert.commandWorked(source.insert({a: "hi", proofOfUpdate: "PROOF"}));
+        assert.doesNotThrow(() => source.aggregate([
+            {$project: {_id: 0, proofOfUpdate: "PROOF", "a.b": "$a"}},
+            {$out: {to: target.getName(), mode: "replaceDocuments", uniqueKey: {"a.b": 1}}}
+        ]));
+        assert.docEq(target.findOne({"a.b": "hi", proofOfUpdate: "PROOF"}),
+                     {_id: "TARGET", a: {b: "hi"}, proofOfUpdate: "PROOF"});
+    }());
+
+    // Test that a unique index that is sparse can still be used.
+    (function testSparseIndex() {
+        const target = testDB.multikey_index;
+        target.drop();
+
+        assert.commandWorked(target.createIndex({a: 1}, {unique: true, sparse: true}));
+        assert.doesNotThrow(
+            () => source.aggregate(
+                [{$out: {to: target.getName(), mode: "replaceDocuments", uniqueKey: {a: 1}}}]));
+        assert.commandWorked(target.insert([{b: 1, c: 1}, {a: null}, {d: 4}]));
+        assert.doesNotThrow(
+            () => source.aggregate(
+                [{$out: {to: target.getName(), mode: "replaceDocuments", uniqueKey: {a: 1}}}]));
+    }());
+}());
diff --git a/src/mongo/db/pipeline/document_source_out.cpp b/src/mongo/db/pipeline/document_source_out.cpp
index dee594e7675..becbdf60214 100644
--- a/src/mongo/db/pipeline/document_source_out.cpp
+++ b/src/mongo/db/pipeline/document_source_out.cpp
@@ -242,6 +242,18 @@ intrusive_ptr<DocumentSource> DocumentSourceOut::createFromBson(
         // Convert unique key object to a vector of FieldPaths.
         if (auto uniqueKeyObj = spec.getUniqueKey()) {
             uniqueKey = uniqueKeyObj->getFieldNames<std::set<FieldPath>>();
+            // TODO SERVER-36047 Check if this is the document key and skip the check below.
+            const bool isDocumentKey = false;
+            // Make sure the uniqueKey has a supporting index. TODO SERVER-36047 Figure out where
+            // this assertion should take place in a sharded environment. For now we will skip the
+            // check on mongos and will not test this check against a sharded collection or another
+            // database - meaning the assertion should always happen on the primary shard where the
+            // collection must exist.
+            uassert(50938,
+                    "Cannot find index to verify that $out's unique key will be unique",
+                    expCtx->inMongos || isDocumentKey ||
+                        expCtx->mongoProcessInterface->uniqueKeyIsSupportedByIndex(
+                            expCtx, outputNs, uniqueKey));
         } else {
             std::vector<FieldPath> docKeyPaths = std::get<0>(
                 expCtx->mongoProcessInterface->collectDocumentKeyFields(expCtx->opCtx, outputNs));
diff --git a/src/mongo/db/pipeline/field_path.h b/src/mongo/db/pipeline/field_path.h
index 64fa7fe0611..25015a6f644 100644
--- a/src/mongo/db/pipeline/field_path.h
+++ b/src/mongo/db/pipeline/field_path.h
@@ -123,4 +123,8 @@ private:
 inline bool operator<(const FieldPath& lhs, const FieldPath& rhs) {
     return lhs.fullPath() < rhs.fullPath();
 }
+
+inline bool operator==(const FieldPath& lhs, const FieldPath& rhs) {
+    return lhs.fullPath() == rhs.fullPath();
+}
 }
diff --git a/src/mongo/db/pipeline/mongo_process_interface.h b/src/mongo/db/pipeline/mongo_process_interface.h
index cea2ee19449..1d7add701d7 100644
--- a/src/mongo/db/pipeline/mongo_process_interface.h
+++ b/src/mongo/db/pipeline/mongo_process_interface.h
@@ -245,6 +245,18 @@ public:
     virtual std::vector<BSONObj> getMatchingPlanCacheEntryStats(OperationContext*,
                                                                 const NamespaceString&,
                                                                 const MatchExpression*) const = 0;
+
+    /**
+     * Returns true if there is an index on 'nss' with properties that will guarantee that a
+     * document with non-array values for each of 'uniqueKeyPaths' will have at most one matching
+     * document in 'nss'.
+     *
+     * Specifically, such an index must include all the fields, be unique, not be a partial index,
+     * and match the operation's collation as given by 'expCtx'.
+     */
+    virtual bool uniqueKeyIsSupportedByIndex(const boost::intrusive_ptr<ExpressionContext>& expCtx,
+                                             const NamespaceString& nss,
+                                             const std::set<FieldPath>& uniqueKeyPaths) const = 0;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/mongod_process_interface.cpp b/src/mongo/db/pipeline/mongod_process_interface.cpp
index 884eae805af..14eb073ef6e 100644
--- a/src/mongo/db/pipeline/mongod_process_interface.cpp
+++ b/src/mongo/db/pipeline/mongod_process_interface.cpp
@@ -34,8 +34,10 @@
 
 #include "mongo/db/auth/authorization_session.h"
 #include "mongo/db/catalog/collection.h"
+#include "mongo/db/catalog/database_holder.h"
 #include "mongo/db/catalog/document_validation.h"
 #include "mongo/db/catalog/uuid_catalog.h"
+#include "mongo/db/concurrency/d_concurrency.h"
 #include "mongo/db/curop.h"
 #include "mongo/db/db_raii.h"
 #include "mongo/db/ops/write_ops_exec.h"
@@ -117,6 +119,32 @@ Update buildUpdateOp(const NamespaceString& nss,
     return updateOp;
 }
 
+// Returns true if the field names of 'keyPattern' are exactly those in 'uniqueKeyPaths', and each
+// of the elements of 'keyPattern' is numeric, i.e. not "text", "$**", or any other special type of
+// index.
+bool keyPatternNamesExactPaths(const BSONObj& keyPattern,
+                               const std::set<FieldPath>& uniqueKeyPaths) {
+    size_t nFieldsMatched = 0;
+    for (auto&& elem : keyPattern) {
+        if (!elem.isNumber()) {
+            return false;
+        }
+        if (uniqueKeyPaths.find(elem.fieldNameStringData()) == uniqueKeyPaths.end()) {
+            return false;
+        }
+        ++nFieldsMatched;
+    }
+    return nFieldsMatched == uniqueKeyPaths.size();
+}
+
+bool supportsUniqueKey(const boost::intrusive_ptr<ExpressionContext>& expCtx,
+                       const IndexCatalogEntry* index,
+                       const std::set<FieldPath>& uniqueKeyPaths) {
+    return (index->descriptor()->unique() && !index->descriptor()->isPartial() &&
+            keyPatternNamesExactPaths(index->descriptor()->keyPattern(), uniqueKeyPaths) &&
+            CollatorInterface::collatorsMatch(index->getCollator(), expCtx->getCollator()));
+}
+
 }  // namespace
 
 // static
@@ -450,6 +478,34 @@ std::vector<BSONObj> MongoDInterface::getMatchingPlanCacheEntryStats(
     return planCache->getMatchingStats(serializer, predicate);
 }
 
+bool MongoDInterface::uniqueKeyIsSupportedByIndex(
+    const boost::intrusive_ptr<ExpressionContext>& expCtx,
+    const NamespaceString& nss,
+    const std::set<FieldPath>& uniqueKeyPaths) const {
+    auto* opCtx = expCtx->opCtx;
+    // We purposefully avoid a helper like AutoGetCollection here because we don't want to check the
+    // db version or do anything else. We simply want to protect against concurrent modifications to
+    // the catalog.
+    Lock::DBLock dbLock(opCtx, nss.db(), MODE_IS);
+    Lock::CollectionLock collLock(opCtx->lockState(), nss.ns(), MODE_IS);
+    const auto* collection = [&]() -> Collection* {
+        auto db = DatabaseHolder::getDatabaseHolder().get(opCtx, nss.db());
+        return db ? db->getCollection(opCtx, nss) : nullptr;
+    }();
+    if (!collection) {
+        return uniqueKeyPaths == std::set<FieldPath>{"_id"};
+    }
+
+    auto indexIterator = collection->getIndexCatalog()->getIndexIterator(opCtx, false);
+    while (indexIterator.more()) {
+        IndexDescriptor* descriptor = indexIterator.next();
+        if (supportsUniqueKey(expCtx, indexIterator.catalogEntry(descriptor), uniqueKeyPaths)) {
+            return true;
+        }
+    }
+    return false;
+}
+
 BSONObj MongoDInterface::_reportCurrentOpForClient(OperationContext* opCtx,
                                                    Client* client,
                                                    CurrentOpTruncateMode truncateOps) const {
diff --git a/src/mongo/db/pipeline/mongod_process_interface.h b/src/mongo/db/pipeline/mongod_process_interface.h
index f419578d456..0b0ccb7e7e7 100644
--- a/src/mongo/db/pipeline/mongod_process_interface.h
+++ b/src/mongo/db/pipeline/mongod_process_interface.h
@@ -102,6 +102,10 @@ public:
                                                         const NamespaceString&,
                                                         const MatchExpression*) const final;
 
+    bool uniqueKeyIsSupportedByIndex(const boost::intrusive_ptr<ExpressionContext>& expCtx,
+                                     const NamespaceString& nss,
+                                     const std::set<FieldPath>& uniqueKeyPaths) const final;
+
 protected:
     BSONObj _reportCurrentOpForClient(OperationContext* opCtx,
                                       Client* client,
diff --git a/src/mongo/db/pipeline/mongos_process_interface.h b/src/mongo/db/pipeline/mongos_process_interface.h
index 55dc4a8d06c..df96e9f1043 100644
--- a/src/mongo/db/pipeline/mongos_process_interface.h
+++ b/src/mongo/db/pipeline/mongos_process_interface.h
@@ -163,6 +163,14 @@ public:
         MONGO_UNREACHABLE;
     }
 
+    bool uniqueKeyIsSupportedByIndex(const boost::intrusive_ptr<ExpressionContext>&,
+                                     const NamespaceString&,
+                                     const std::set<FieldPath>& uniqueKeyPaths) const final {
+        // TODO SERVER-36047 we'll have to contact the primary shard for the database to ask for the
+        // index specs.
+        return true;
+    }
+
 protected:
     BSONObj _reportCurrentOpForClient(OperationContext* opCtx,
                                       Client* client,
diff --git a/src/mongo/db/pipeline/stub_mongo_process_interface.h b/src/mongo/db/pipeline/stub_mongo_process_interface.h
index d1d33f7da4d..24376600f5e 100644
--- a/src/mongo/db/pipeline/stub_mongo_process_interface.h
+++ b/src/mongo/db/pipeline/stub_mongo_process_interface.h
@@ -174,5 +174,11 @@ public:
                                                         const MatchExpression*) const override {
         MONGO_UNREACHABLE;
     }
+
+    bool uniqueKeyIsSupportedByIndex(const boost::intrusive_ptr<ExpressionContext>& expCtx,
+                                     const NamespaceString& nss,
+                                     const std::set<FieldPath>& uniqueKeyPaths) const override {
+        return true;
+    }
 };
 }  // namespace mongo