SERVER-36424 Enforce uniqueKey is unique enough

7 files changed, 102 insertions, 0 deletions

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