SERVER-31447 Use correct collation for update lookup

7 files changed, 78 insertions, 30 deletions

diff --git a/src/mongo/db/pipeline/document_source.h b/src/mongo/db/pipeline/document_source.h
index 8e4278f040b..c576fd5bda8 100644
--- a/src/mongo/db/pipeline/document_source.h
+++ b/src/mongo/db/pipeline/document_source.h
@@ -840,13 +840,15 @@ public:
         virtual std::vector<FieldPath> collectDocumentKeyFields(UUID) const = 0;
 
         /**
-         * Returns zero or one documents matching the input filter, or throws if more than one match
-         * was found. The passed ExpressionContext may use a different namespace than the
-         * ExpressionContext used to construct the MongoProcessInterface. Returns boost::none if no
-         * matching documents were found, including cases where the given namespace does not exist.
+         * Returns zero or one documents with the document key 'documentKey'. 'documentKey' is
+         * treated as a unique identifier of a document, and may include an _id or all fields from
+         * the shard key and an _id. Throws if more than one match was found. Returns boost::none if
+         * no matching documents were found, including cases where the given namespace does not
+         * exist.
          */
-        virtual boost::optional<Document> lookupSingleDocument(
-            const boost::intrusive_ptr<ExpressionContext>& expCtx, const Document& filter) = 0;
+        virtual boost::optional<Document> lookupSingleDocument(const NamespaceString& nss,
+                                                               UUID collectionUUID,
+                                                               const Document& documentKey) = 0;
 
         // Add new methods as needed.
     };
diff --git a/src/mongo/db/pipeline/document_source_lookup_change_post_image.cpp b/src/mongo/db/pipeline/document_source_lookup_change_post_image.cpp
index 058a959f588..7752fbfd751 100644
--- a/src/mongo/db/pipeline/document_source_lookup_change_post_image.cpp
+++ b/src/mongo/db/pipeline/document_source_lookup_change_post_image.cpp
@@ -105,10 +105,9 @@ Value DocumentSourceLookupChangePostImage::lookupPostImage(const Document& updat
     auto resumeToken =
         ResumeToken::parse(updateOp[DocumentSourceChangeStream::kIdField].getDocument());
 
-    // TODO SERVER-29134 we need to extract the namespace from the document and set them on the new
-    // ExpressionContext if we're getting notifications from an entire database.
-    auto foreignExpCtx = pExpCtx->copyWith(nss, resumeToken.getData().uuid);
-    auto lookedUpDoc = _mongoProcessInterface->lookupSingleDocument(foreignExpCtx, documentKey);
+    invariant(resumeToken.getData().uuid);
+    auto lookedUpDoc =
+        _mongoProcessInterface->lookupSingleDocument(nss, *resumeToken.getData().uuid, documentKey);
 
     // Check whether the lookup returned any documents. Even if the lookup itself succeeded, it may
     // not have returned any results if the document was deleted in the time since the update op.
diff --git a/src/mongo/db/pipeline/document_source_lookup_change_post_image_test.cpp b/src/mongo/db/pipeline/document_source_lookup_change_post_image_test.cpp
index 05b92297359..3d35f875619 100644
--- a/src/mongo/db/pipeline/document_source_lookup_change_post_image_test.cpp
+++ b/src/mongo/db/pipeline/document_source_lookup_change_post_image_test.cpp
@@ -113,9 +113,11 @@ public:
         return Status::OK();
     }
 
-    boost::optional<Document> lookupSingleDocument(
-        const boost::intrusive_ptr<ExpressionContext>& expCtx, const Document& filter) {
-        auto swPipeline = makePipeline({BSON("$match" << filter)}, expCtx);
+    boost::optional<Document> lookupSingleDocument(const NamespaceString& nss,
+                                                   UUID collectionUUID,
+                                                   const Document& documentKey) {
+        boost::intrusive_ptr<ExpressionContextForTest> expCtx(new ExpressionContextForTest(nss));
+        auto swPipeline = makePipeline({BSON("$match" << documentKey)}, expCtx);
         if (swPipeline == ErrorCodes::NamespaceNotFound) {
             return boost::none;
         }
@@ -124,11 +126,12 @@ public:
         auto lookedUpDocument = pipeline->getNext();
         if (auto next = pipeline->getNext()) {
             uasserted(ErrorCodes::TooManyMatchingDocuments,
-                      str::stream() << "found more than one document matching " << filter.toString()
+                      str::stream() << "found more than one document matching "
+                                    << documentKey.toString()
                                     << " ["
-                                    << (*lookedUpDocument).toString()
+                                    << lookedUpDocument->toString()
                                     << ", "
-                                    << (*next).toString()
+                                    << next->toString()
                                     << "]");
         }
         return lookedUpDocument;
diff --git a/src/mongo/db/pipeline/expression_context.cpp b/src/mongo/db/pipeline/expression_context.cpp
index 64c813ed5a4..e2b2e1e758d 100644
--- a/src/mongo/db/pipeline/expression_context.cpp
+++ b/src/mongo/db/pipeline/expression_context.cpp
@@ -29,6 +29,7 @@
 #include "mongo/platform/basic.h"
 
 #include "mongo/db/pipeline/expression_context.h"
+#include "mongo/db/query/collation/collation_spec.h"
 #include "mongo/db/query/collation/collator_factory_interface.h"
 
 namespace mongo {
@@ -121,8 +122,10 @@ void ExpressionContext::setCollator(const CollatorInterface* collator) {
     _valueComparator = ValueComparator(_collator);
 }
 
-intrusive_ptr<ExpressionContext> ExpressionContext::copyWith(NamespaceString ns,
-                                                             boost::optional<UUID> uuid) const {
+intrusive_ptr<ExpressionContext> ExpressionContext::copyWith(
+    NamespaceString ns,
+    boost::optional<UUID> uuid,
+    boost::optional<std::unique_ptr<CollatorInterface>> collator) const {
     intrusive_ptr<ExpressionContext> expCtx =
         new ExpressionContext(std::move(ns), timeZoneDatabase);
 
@@ -140,11 +143,17 @@ intrusive_ptr<ExpressionContext> ExpressionContext::copyWith(NamespaceString ns,
 
     expCtx->opCtx = opCtx;
 
-    expCtx->collation = collation;
-    if (_ownedCollator) {
-        expCtx->setCollator(_ownedCollator->clone());
-    } else if (_collator) {
-        expCtx->setCollator(_collator);
+    if (collator) {
+        expCtx->collation =
+            *collator ? (*collator)->getSpec().toBSON() : CollationSpec::kSimpleSpec;
+        expCtx->setCollator(std::move(*collator));
+    } else {
+        expCtx->collation = collation;
+        if (_ownedCollator) {
+            expCtx->setCollator(_ownedCollator->clone());
+        } else if (_collator) {
+            expCtx->setCollator(_collator);
+        }
     }
 
     expCtx->_resolvedNamespaces = _resolvedNamespaces;
diff --git a/src/mongo/db/pipeline/expression_context.h b/src/mongo/db/pipeline/expression_context.h
index 59e85edfe90..a2e16163a73 100644
--- a/src/mongo/db/pipeline/expression_context.h
+++ b/src/mongo/db/pipeline/expression_context.h
@@ -140,7 +140,9 @@ public:
      * separate aggregation pipeline on 'ns' with the optional 'uuid'.
      */
     boost::intrusive_ptr<ExpressionContext> copyWith(
-        NamespaceString ns, boost::optional<UUID> uuid = boost::none) const;
+        NamespaceString ns,
+        boost::optional<UUID> uuid = boost::none,
+        boost::optional<std::unique_ptr<CollatorInterface>> collator = boost::none) const;
 
     /**
      * Returns the ResolvedNamespace corresponding to 'nss'. It is an error to call this method on a
diff --git a/src/mongo/db/pipeline/pipeline_d.cpp b/src/mongo/db/pipeline/pipeline_d.cpp
index ee1ded37c99..65214a154bf 100644
--- a/src/mongo/db/pipeline/pipeline_d.cpp
+++ b/src/mongo/db/pipeline/pipeline_d.cpp
@@ -384,9 +384,14 @@ public:
         return result;
     }
 
-    boost::optional<Document> lookupSingleDocument(const intrusive_ptr<ExpressionContext>& expCtx,
-                                                   const Document& filter) final {
-        auto swPipeline = makePipeline({BSON("$match" << filter)}, expCtx);
+    boost::optional<Document> lookupSingleDocument(const NamespaceString& nss,
+                                                   UUID collectionUUID,
+                                                   const Document& documentKey) final {
+        // Be sure to do the lookup using the collection default collation.
+        auto foreignExpCtx =
+            _ctx->copyWith(nss, collectionUUID, _getCollectionDefaultCollator(nss, collectionUUID));
+        auto swPipeline = makePipeline({BSON("$match" << documentKey)}, foreignExpCtx);
+
         if (swPipeline == ErrorCodes::NamespaceNotFound) {
             return boost::none;
         }
@@ -395,7 +400,8 @@ public:
         auto lookedUpDocument = pipeline->getNext();
         if (auto next = pipeline->getNext()) {
             uasserted(ErrorCodes::TooManyMatchingDocuments,
-                      str::stream() << "found more than one document matching " << filter.toString()
+                      str::stream() << "found more than one document with document key "
+                                    << documentKey.toString()
                                     << " ["
                                     << lookedUpDocument->toString()
                                     << ", "
@@ -406,8 +412,34 @@ public:
     }
 
 private:
+    /**
+     * Looks up the collection default collator for the collection given by 'collectionUUID'. A
+     * collection's default collation is not allowed to change, so we cache the result to allow for
+     * quick lookups in the future. Looks up the collection by UUID, and returns 'nullptr' if the
+     * collection does not exist or if the collection's default collation is the simple collation.
+     */
+    std::unique_ptr<CollatorInterface> _getCollectionDefaultCollator(const NamespaceString& nss,
+                                                                     UUID collectionUUID) {
+        if (_collatorCache.find(collectionUUID) == _collatorCache.end()) {
+            AutoGetCollection autoColl(_ctx->opCtx, nss, collectionUUID, MODE_IS);
+            if (!autoColl.getCollection()) {
+                // This collection doesn't exist - since we looked up by UUID, it will never exist
+                // in the future, so we cache a null pointer as the default collation.
+                _collatorCache[collectionUUID] = nullptr;
+            } else {
+                auto defaultCollator = autoColl.getCollection()->getDefaultCollator();
+                // Clone the collator so that we can safely use the pointer if the collection
+                // disappears right after we release the lock.
+                _collatorCache[collectionUUID] =
+                    defaultCollator ? defaultCollator->clone() : nullptr;
+            }
+        }
+        return _collatorCache[collectionUUID] ? _collatorCache[collectionUUID]->clone() : nullptr;
+    }
+
     intrusive_ptr<ExpressionContext> _ctx;
     DBDirectClient _client;
+    std::map<UUID, std::unique_ptr<const CollatorInterface>> _collatorCache;
 };
 
 /**
diff --git a/src/mongo/db/pipeline/stub_mongo_process_interface.h b/src/mongo/db/pipeline/stub_mongo_process_interface.h
index 06e22b24bdd..bebf8f02a34 100644
--- a/src/mongo/db/pipeline/stub_mongo_process_interface.h
+++ b/src/mongo/db/pipeline/stub_mongo_process_interface.h
@@ -118,8 +118,9 @@ public:
         MONGO_UNREACHABLE;
     }
 
-    boost::optional<Document> lookupSingleDocument(
-        const boost::intrusive_ptr<ExpressionContext>& expCtx, const Document& filter) {
+    boost::optional<Document> lookupSingleDocument(const NamespaceString& nss,
+                                                   UUID collectionUUID,
+                                                   const Document& documentKey) {
         MONGO_UNREACHABLE;
     }
 };