SERVER-37982 Distinguish use cases for collecting document key fields

14 files changed, 165 insertions, 137 deletions

diff --git a/src/mongo/db/pipeline/document_source_change_stream_test.cpp b/src/mongo/db/pipeline/document_source_change_stream_test.cpp
index b1fd9c1491b..2e303b58f27 100644
--- a/src/mongo/db/pipeline/document_source_change_stream_test.cpp
+++ b/src/mongo/db/pipeline/document_source_change_stream_test.cpp
@@ -90,8 +90,8 @@ struct MockMongoInterface final : public StubMongoProcessInterface {
 
     MockMongoInterface(std::vector<FieldPath> fields) : _fields(std::move(fields)) {}
 
-    std::pair<std::vector<FieldPath>, bool> collectDocumentKeyFields(
-        OperationContext*, NamespaceStringOrUUID) const final {
+    std::pair<std::vector<FieldPath>, bool> collectDocumentKeyFieldsForHostedCollection(
+        OperationContext*, const NamespaceString&, UUID) const final {
         return {_fields, false};
     }
 
@@ -1091,16 +1091,18 @@ TEST_F(ChangeStreamStageTest, DocumentKeyShouldIncludeShardKeyFromResumeToken) {
     };
     // Although the chunk manager and sharding catalog are not aware of the shard key in this test,
     // the expectation is for the $changeStream stage to infer the shard key from the resume token.
-    checkTransformation(insertEntry,
-                        expectedInsert,
-                        {{"_id"}},  // Mock the 'collectDocumentKeyFields' response.
-                        BSON("$changeStream" << BSON("resumeAfter" << resumeToken)));
+    checkTransformation(
+        insertEntry,
+        expectedInsert,
+        {{"_id"}},  // Mock the 'collectDocumentKeyFieldsForHostedCollection' response.
+        BSON("$changeStream" << BSON("resumeAfter" << resumeToken)));
 
     // Verify the same behavior with resuming using 'startAfter'.
-    checkTransformation(insertEntry,
-                        expectedInsert,
-                        {{"_id"}},  // Mock the 'collectDocumentKeyFields' response.
-                        BSON("$changeStream" << BSON("startAfter" << resumeToken)));
+    checkTransformation(
+        insertEntry,
+        expectedInsert,
+        {{"_id"}},  // Mock the 'collectDocumentKeyFieldsForHostedCollection' response.
+        BSON("$changeStream" << BSON("startAfter" << resumeToken)));
 }
 
 TEST_F(ChangeStreamStageTest, DocumentKeyShouldNotIncludeShardKeyFieldsIfNotPresentInOplogEntry) {
@@ -1133,16 +1135,18 @@ TEST_F(ChangeStreamStageTest, DocumentKeyShouldNotIncludeShardKeyFieldsIfNotPres
         {DSChangeStream::kNamespaceField, D{{"db", nss.db()}, {"coll", nss.coll()}}},
         {DSChangeStream::kDocumentKeyField, D{{"_id", 2}}},
     };
-    checkTransformation(insertEntry,
-                        expectedInsert,
-                        {{"_id"}},  // Mock the 'collectDocumentKeyFields' response.
-                        BSON("$changeStream" << BSON("resumeAfter" << resumeToken)));
+    checkTransformation(
+        insertEntry,
+        expectedInsert,
+        {{"_id"}},  // Mock the 'collectDocumentKeyFieldsForHostedCollection' response.
+        BSON("$changeStream" << BSON("resumeAfter" << resumeToken)));
 
     // Verify the same behavior with resuming using 'startAfter'.
-    checkTransformation(insertEntry,
-                        expectedInsert,
-                        {{"_id"}},  // Mock the 'collectDocumentKeyFields' response.
-                        BSON("$changeStream" << BSON("startAfter" << resumeToken)));
+    checkTransformation(
+        insertEntry,
+        expectedInsert,
+        {{"_id"}},  // Mock the 'collectDocumentKeyFieldsForHostedCollection' response.
+        BSON("$changeStream" << BSON("startAfter" << resumeToken)));
 }
 
 TEST_F(ChangeStreamStageTest, ResumeAfterFailsIfResumeTokenDoesNotContainUUID) {
@@ -1562,10 +1566,11 @@ TEST_F(ChangeStreamStageDBTest, DocumentKeyShouldIncludeShardKeyFromResumeToken)
         {DSChangeStream::kNamespaceField, D{{"db", nss.db()}, {"coll", nss.coll()}}},
         {DSChangeStream::kDocumentKeyField, D{{"_id", 2}, {"shardKey", 3}}},
     };
-    checkTransformation(insertEntry,
-                        expectedInsert,
-                        {{"_id"}},  // Mock the 'collectDocumentKeyFields' response.
-                        BSON("$changeStream" << BSON("resumeAfter" << resumeToken)));
+    checkTransformation(
+        insertEntry,
+        expectedInsert,
+        {{"_id"}},  // Mock the 'collectDocumentKeyFieldsForHostedCollection' response.
+        BSON("$changeStream" << BSON("resumeAfter" << resumeToken)));
 }
 
 TEST_F(ChangeStreamStageDBTest, DocumentKeyShouldNotIncludeShardKeyFieldsIfNotPresentInOplogEntry) {
@@ -1598,10 +1603,11 @@ TEST_F(ChangeStreamStageDBTest, DocumentKeyShouldNotIncludeShardKeyFieldsIfNotPr
         {DSChangeStream::kNamespaceField, D{{"db", nss.db()}, {"coll", nss.coll()}}},
         {DSChangeStream::kDocumentKeyField, D{{"_id", 2}}},
     };
-    checkTransformation(insertEntry,
-                        expectedInsert,
-                        {{"_id"}},  // Mock the 'collectDocumentKeyFields' response.
-                        BSON("$changeStream" << BSON("resumeAfter" << resumeToken)));
+    checkTransformation(
+        insertEntry,
+        expectedInsert,
+        {{"_id"}},  // Mock the 'collectDocumentKeyFieldsForHostedCollection' response.
+        BSON("$changeStream" << BSON("resumeAfter" << resumeToken)));
 }
 
 TEST_F(ChangeStreamStageDBTest, DocumentKeyShouldNotIncludeShardKeyIfResumeTokenDoesntContainUUID) {
@@ -1635,10 +1641,11 @@ TEST_F(ChangeStreamStageDBTest, DocumentKeyShouldNotIncludeShardKeyIfResumeToken
         {DSChangeStream::kNamespaceField, D{{"db", nss.db()}, {"coll", nss.coll()}}},
         {DSChangeStream::kDocumentKeyField, D{{"_id", 2}}},
     };
-    checkTransformation(insertEntry,
-                        expectedInsert,
-                        {{"_id"}},  // Mock the 'collectDocumentKeyFields' response.
-                        BSON("$changeStream" << BSON("resumeAfter" << resumeToken)));
+    checkTransformation(
+        insertEntry,
+        expectedInsert,
+        {{"_id"}},  // Mock the 'collectDocumentKeyFieldsForHostedCollection' response.
+        BSON("$changeStream" << BSON("resumeAfter" << resumeToken)));
 }
 
 TEST_F(ChangeStreamStageDBTest, ResumeAfterWithTokenFromInvalidateShouldFail) {
@@ -1684,10 +1691,11 @@ TEST_F(ChangeStreamStageDBTest, ResumeAfterWithTokenFromDropDatabase) {
         {DSChangeStream::kNamespaceField, D{{"db", nss.db()}, {"coll", nss.coll()}}},
         {DSChangeStream::kDocumentKeyField, D{{"_id", 2}}},
     };
-    checkTransformation(insertEntry,
-                        expectedInsert,
-                        {{"_id"}},  // Mock the 'collectDocumentKeyFields' response.
-                        BSON("$changeStream" << BSON("resumeAfter" << resumeToken)));
+    checkTransformation(
+        insertEntry,
+        expectedInsert,
+        {{"_id"}},  // Mock the 'collectDocumentKeyFieldsForHostedCollection' response.
+        BSON("$changeStream" << BSON("resumeAfter" << resumeToken)));
 }
 
 
@@ -1711,10 +1719,11 @@ TEST_F(ChangeStreamStageDBTest, StartAfterSucceedsEvenIfResumeTokenDoesNotContai
         {DSChangeStream::kNamespaceField, D{{"db", nss.db()}, {"coll", nss.coll()}}},
         {DSChangeStream::kDocumentKeyField, D{{"_id", 2}}},
     };
-    checkTransformation(insertEntry,
-                        expectedInsert,
-                        {{"_id"}},  // Mock the 'collectDocumentKeyFields' response.
-                        BSON("$changeStream" << BSON("startAfter" << resumeToken)));
+    checkTransformation(
+        insertEntry,
+        expectedInsert,
+        {{"_id"}},  // Mock the 'collectDocumentKeyFieldsForHostedCollection' response.
+        BSON("$changeStream" << BSON("startAfter" << resumeToken)));
 }
 
 }  // namespace
diff --git a/src/mongo/db/pipeline/document_source_change_stream_transform.cpp b/src/mongo/db/pipeline/document_source_change_stream_transform.cpp
index 8a94978ef15..db600fb9e5b 100644
--- a/src/mongo/db/pipeline/document_source_change_stream_transform.cpp
+++ b/src/mongo/db/pipeline/document_source_change_stream_transform.cpp
@@ -229,8 +229,9 @@ Document DocumentSourceChangeStreamTransform::applyTransformation(const Document
         // unsharded when the entry was last populated.
         auto it = _documentKeyCache.find(uuid.getUuid());
         if (it == _documentKeyCache.end() || !it->second.isFinal) {
-            auto docKeyFields = pExpCtx->mongoProcessInterface->collectDocumentKeyFields(
-                pExpCtx->opCtx, NamespaceStringOrUUID(nss.db().toString(), uuid.getUuid()));
+            auto docKeyFields =
+                pExpCtx->mongoProcessInterface->collectDocumentKeyFieldsForHostedCollection(
+                    pExpCtx->opCtx, nss, uuid.getUuid());
             if (it == _documentKeyCache.end() || docKeyFields.second) {
                 _documentKeyCache[uuid.getUuid()] = DocumentKeyCacheEntry(docKeyFields);
             }
diff --git a/src/mongo/db/pipeline/document_source_out.cpp b/src/mongo/db/pipeline/document_source_out.cpp
index 48df0c2b7b8..9d8e31da131 100644
--- a/src/mongo/db/pipeline/document_source_out.cpp
+++ b/src/mongo/db/pipeline/document_source_out.cpp
@@ -365,6 +365,7 @@ intrusive_ptr<DocumentSource> DocumentSourceOut::createFromBson(
                         expCtx->mongoProcessInterface->uniqueKeyIsSupportedByIndex(
                             expCtx, outputNs, uniqueKey));
         } else {
+            uassert(51009, "Expected uniqueKey to be provided from mongos", !expCtx->fromMongos);
             if (expCtx->inMongos && mode != WriteModeEnum::kModeReplaceCollection) {
                 // In case there are multiple shards which will perform this $out in parallel, we
                 // need to figure out and attach the collection's epoch to ensure each shard is
@@ -382,8 +383,11 @@ intrusive_ptr<DocumentSource> DocumentSourceOut::createFromBson(
                 // waiting for that request to return instead of forcing another refresh.
                 targetEpoch = expCtx->mongoProcessInterface->refreshAndGetEpoch(expCtx, outputNs);
             }
-            std::vector<FieldPath> docKeyPaths = std::get<0>(
-                expCtx->mongoProcessInterface->collectDocumentKeyFields(expCtx->opCtx, outputNs));
+            // Even if we're not on mongos, we're still acting as a router here - the targeted
+            // collection may not be completely on our shard.
+            auto docKeyPaths =
+                expCtx->mongoProcessInterface->collectDocumentKeyFieldsActingAsRouter(expCtx->opCtx,
+                                                                                      outputNs);
             uniqueKey = std::set<FieldPath>(std::make_move_iterator(docKeyPaths.begin()),
                                             std::make_move_iterator(docKeyPaths.end()));
         }
diff --git a/src/mongo/db/pipeline/document_source_out_test.cpp b/src/mongo/db/pipeline/document_source_out_test.cpp
index 8c45e663a01..44256d1fb4b 100644
--- a/src/mongo/db/pipeline/document_source_out_test.cpp
+++ b/src/mongo/db/pipeline/document_source_out_test.cpp
@@ -63,9 +63,9 @@ public:
      * For the purposes of these tests, pretend each collection is unsharded and has a document key
      * of just "_id".
      */
-    std::pair<std::vector<FieldPath>, bool> collectDocumentKeyFields(
-        OperationContext* opCtx, NamespaceStringOrUUID nssOrUUID) const override {
-        return {{"_id"}, false};
+    std::vector<FieldPath> collectDocumentKeyFieldsActingAsRouter(
+        OperationContext* opCtx, const NamespaceString& nss) const override {
+        return {"_id"};
     }
 };
 
diff --git a/src/mongo/db/pipeline/mongo_process_common.cpp b/src/mongo/db/pipeline/mongo_process_common.cpp
index 6b990f5373f..3ab1d1a2bfc 100644
--- a/src/mongo/db/pipeline/mongo_process_common.cpp
+++ b/src/mongo/db/pipeline/mongo_process_common.cpp
@@ -122,6 +122,18 @@ std::vector<BSONObj> MongoProcessCommon::getCurrentOps(
     return ops;
 }
 
+std::vector<FieldPath> MongoProcessCommon::collectDocumentKeyFieldsActingAsRouter(
+    OperationContext* opCtx, const NamespaceString& nss) const {
+    if (auto chunkManager =
+            uassertStatusOK(Grid::get(opCtx)->catalogCache()->getCollectionRoutingInfo(opCtx, nss))
+                .cm()) {
+        return _shardKeyToDocumentKeyFields(
+            chunkManager->getShardKeyPattern().getKeyPatternFields());
+    }
+    // We have no evidence this collection is sharded, so the document key is just _id.
+    return {"_id"};
+}
+
 bool MongoProcessCommon::keyPatternNamesExactPaths(const BSONObj& keyPattern,
                                                    const std::set<FieldPath>& uniqueKeyPaths) {
     size_t nFieldsMatched = 0;
@@ -149,4 +161,18 @@ boost::optional<OID> MongoProcessCommon::refreshAndGetEpoch(
     }
     return boost::none;
 }
+
+std::vector<FieldPath> MongoProcessCommon::_shardKeyToDocumentKeyFields(
+    const std::vector<std::unique_ptr<FieldRef>>& keyPatternFields) const {
+    std::vector<FieldPath> result;
+    bool gotId = false;
+    for (auto& field : keyPatternFields) {
+        result.emplace_back(field->dottedField());
+        gotId |= (result.back().fullPath() == "_id");
+    }
+    if (!gotId) {  // If not part of the shard key, "_id" comes last.
+        result.emplace_back("_id");
+    }
+    return result;
+}
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/mongo_process_common.h b/src/mongo/db/pipeline/mongo_process_common.h
index 53ce182e93a..ff2a7a5595a 100644
--- a/src/mongo/db/pipeline/mongo_process_common.h
+++ b/src/mongo/db/pipeline/mongo_process_common.h
@@ -60,12 +60,22 @@ public:
                                        CurrentOpTruncateMode truncateMode,
                                        CurrentOpCursorMode cursorMode) const final;
 
+    virtual std::vector<FieldPath> collectDocumentKeyFieldsActingAsRouter(
+        OperationContext*, const NamespaceString&) const override;
+
     virtual boost::optional<OID> refreshAndGetEpoch(
         const boost::intrusive_ptr<ExpressionContext>& expCtx,
         const NamespaceString& nss) const final;
 
 protected:
     /**
+     * Converts the fields from a ShardKeyPattern to a vector of FieldPaths, including the _id if
+     * it's not already in 'keyPatternFields'.
+     */
+    std::vector<FieldPath> _shardKeyToDocumentKeyFields(
+        const std::vector<std::unique_ptr<FieldRef>>& keyPatternFields) const;
+
+    /**
      * Returns a BSONObj representing a report of the operation which is currently being
      * executed by the supplied client. This method is called by the getCurrentOps method of
      * MongoProcessCommon to delegate to the mongoS- or mongoD- specific implementation.
diff --git a/src/mongo/db/pipeline/mongo_process_interface.h b/src/mongo/db/pipeline/mongo_process_interface.h
index c5f11d566cf..8071043c1b3 100644
--- a/src/mongo/db/pipeline/mongo_process_interface.h
+++ b/src/mongo/db/pipeline/mongo_process_interface.h
@@ -220,12 +220,26 @@ public:
     /**
      * Returns the fields of the document key (in order) for the collection corresponding to 'uuid',
      * including the shard key and _id. If _id is not in the shard key, it is added last. If the
-     * collection is not sharded or no longer exists, returns only _id. Also retrurns a boolean that
+     * collection is not sharded or no longer exists, returns only _id. Also returns a boolean that
      * indicates whether the returned fields of the document key are final and will never change for
      * the given collection, either because the collection was dropped or has become sharded.
+     *
+     * This method is meant to be called from a mongod which owns at least one chunk for this
+     * collection. It will inspect the CollectionShardingState, not the CatalogCache. If asked about
+     * a collection not hosted on this shard, the answer will be incorrect.
+     */
+    virtual std::pair<std::vector<FieldPath>, bool> collectDocumentKeyFieldsForHostedCollection(
+        OperationContext* opCtx, const NamespaceString&, UUID) const = 0;
+
+    /**
+     * Returns the fields of the document key (in order) for the collection 'nss', according to the
+     * CatalogCache. The document key fields are the shard key (if sharded) and the _id (if not
+     * already in the shard key). If _id is not in the shard key, it is added last. If the
+     * collection is not sharded or is not known to exist, returns only _id. Does not refresh the
+     * CatalogCache.
      */
-    virtual std::pair<std::vector<FieldPath>, bool> collectDocumentKeyFields(
-        OperationContext* opCtx, NamespaceStringOrUUID nssOrUUID) const = 0;
+    virtual std::vector<FieldPath> collectDocumentKeyFieldsActingAsRouter(
+        OperationContext* opCtx, const NamespaceString&) const = 0;
 
     /**
      * Returns zero or one documents with the document key 'documentKey'. 'documentKey' is treated
diff --git a/src/mongo/db/pipeline/mongos_process_interface.cpp b/src/mongo/db/pipeline/mongos_process_interface.cpp
index 75d69161eb0..85f141c9277 100644
--- a/src/mongo/db/pipeline/mongos_process_interface.cpp
+++ b/src/mongo/db/pipeline/mongos_process_interface.cpp
@@ -210,36 +210,6 @@ boost::optional<Document> MongoSInterface::lookupSingleDocument(
     return (!batch.empty() ? Document(batch.front()) : boost::optional<Document>{});
 }
 
-std::pair<std::vector<FieldPath>, bool> MongoSInterface::collectDocumentKeyFields(
-    OperationContext* opCtx, NamespaceStringOrUUID nssOrUUID) const {
-
-    invariant(!nssOrUUID.uuid(), "Did not expect to use this method with a UUID on mongos");
-    const NamespaceString& nss = *nssOrUUID.nss();
-
-    auto collRoutInfo = Grid::get(opCtx)->catalogCache()->getCollectionRoutingInfo(opCtx, nss);
-    if (collRoutInfo == ErrorCodes::NamespaceNotFound) {
-        return {{"_id"}, false};
-    }
-    uassertStatusOKWithContext(collRoutInfo, "Collection Routing Info is unavailable");
-
-    auto cm = collRoutInfo.getValue().cm();
-    if (!cm)
-        return {{"_id"}, false};
-
-    // Unpack the shard key.
-    std::vector<FieldPath> result;
-    bool gotId = false;
-    for (auto& field : cm->getShardKeyPattern().getKeyPatternFields()) {
-        result.emplace_back(field->dottedField());
-        gotId |= (result.back().fullPath() == "_id");
-    }
-    if (!gotId) {  // If not part of the shard key, "_id" comes last.
-        result.emplace_back("_id");
-    }
-    // Collection is sharded so the document key fields will never change, mark as final.
-    return {result, true};
-}
-
 BSONObj MongoSInterface::_reportCurrentOpForClient(OperationContext* opCtx,
                                                    Client* client,
                                                    CurrentOpTruncateMode truncateOps) const {
diff --git a/src/mongo/db/pipeline/mongos_process_interface.h b/src/mongo/db/pipeline/mongos_process_interface.h
index d9e19b22f11..6f7c1e35505 100644
--- a/src/mongo/db/pipeline/mongos_process_interface.h
+++ b/src/mongo/db/pipeline/mongos_process_interface.h
@@ -128,8 +128,10 @@ public:
         MONGO_UNREACHABLE;
     }
 
-    std::pair<std::vector<FieldPath>, bool> collectDocumentKeyFields(
-        OperationContext* opCtx, NamespaceStringOrUUID nssOrUUID) const final;
+    std::pair<std::vector<FieldPath>, bool> collectDocumentKeyFieldsForHostedCollection(
+        OperationContext* opCtx, const NamespaceString&, UUID) const final {
+        MONGO_UNREACHABLE;
+    }
 
     StatusWith<std::unique_ptr<Pipeline, PipelineDeleter>> makePipeline(
         const std::vector<BSONObj>& rawPipeline,
diff --git a/src/mongo/db/pipeline/process_interface_shardsvr.cpp b/src/mongo/db/pipeline/process_interface_shardsvr.cpp
index c8c9aa66e28..47bd71bd753 100644
--- a/src/mongo/db/pipeline/process_interface_shardsvr.cpp
+++ b/src/mongo/db/pipeline/process_interface_shardsvr.cpp
@@ -73,62 +73,29 @@ void attachWriteConcern(BatchedCommandRequest* request, const WriteConcernOption
 
 }  // namespace
 
-std::pair<std::vector<FieldPath>, bool> MongoInterfaceShardServer::collectDocumentKeyFields(
-    OperationContext* opCtx, NamespaceStringOrUUID nssOrUUID) const {
+std::pair<std::vector<FieldPath>, bool>
+MongoInterfaceShardServer::collectDocumentKeyFieldsForHostedCollection(OperationContext* opCtx,
+                                                                       const NamespaceString& nss,
+                                                                       UUID uuid) const {
     invariant(serverGlobalParams.clusterRole == ClusterRole::ShardServer);
 
-    boost::optional<UUID> uuid;
-    NamespaceString nss;
-    if (nssOrUUID.uuid()) {
-        uuid = *(nssOrUUID.uuid());
-        nss = UUIDCatalog::get(opCtx).lookupNSSByUUID(*uuid);
-        // An empty namespace indicates that the collection has been dropped. Treat it as unsharded
-        // and mark the fields as final.
-        if (nss.isEmpty()) {
-            return {{"_id"}, true};
-        }
-    } else if (nssOrUUID.nss()) {
-        nss = *(nssOrUUID.nss());
-    }
-
-    // Before taking a collection lock to retrieve the shard key fields, consult the catalog cache
-    // to determine whether the collection is sharded in the first place.
-    auto catalogCache = Grid::get(opCtx)->catalogCache();
-
-    const bool collectionIsSharded = catalogCache && [&]() {
-        auto routingInfo = catalogCache->getCollectionRoutingInfo(opCtx, nss);
-        return routingInfo.isOK() && routingInfo.getValue().cm();
-    }();
-
-    // Collection exists and is not sharded, mark as not final.
-    if (!collectionIsSharded) {
-        return {{"_id"}, false};
-    }
-
-    const auto metadata = [opCtx, &nss]() {
-        AutoGetCollection autoColl(opCtx, nss, MODE_IS);
+    const auto metadata = [opCtx, &nss]() -> ScopedCollectionMetadata {
+        Lock::DBLock dbLock(opCtx, nss.db(), MODE_IS);
+        Lock::CollectionLock collLock(opCtx->lockState(), nss.ns(), MODE_IS);
         return CollectionShardingState::get(opCtx, nss)->getCurrentMetadata();
     }();
 
-    // If the UUID is set in 'nssOrUuid', check that the UUID in the ScopedCollectionMetadata
-    // matches. Otherwise, this implies that the collection has been dropped and recreated as
-    // sharded.
-    if (!metadata->isSharded() || (uuid && !metadata->uuidMatches(*uuid))) {
+    if (!metadata->isSharded() || !metadata->uuidMatches(uuid)) {
+        // An unsharded collection can still become sharded so is not final. If the uuid doesn't
+        // match the one stored in the ScopedCollectionMetadata, this implies that the collection
+        // has been dropped and recreated as sharded. We don't know what the old document key fields
+        // might have been in this case so we return just _id.
         return {{"_id"}, false};
     }
 
-    // Unpack the shard key.
-    std::vector<FieldPath> result;
-    bool gotId = false;
-    for (auto& field : metadata->getKeyPatternFields()) {
-        result.emplace_back(field->dottedField());
-        gotId |= (result.back().fullPath() == "_id");
-    }
-    if (!gotId) {  // If not part of the shard key, "_id" comes last.
-        result.emplace_back("_id");
-    }
-    // Collection is now sharded so the document key fields will never change, mark as final.
-    return {result, true};
+    // Unpack the shard key. Collection is now sharded so the document key fields will never change,
+    // mark as final.
+    return {_shardKeyToDocumentKeyFields(metadata->getKeyPatternFields()), true};
 }
 
 void MongoInterfaceShardServer::insert(const boost::intrusive_ptr<ExpressionContext>& expCtx,
diff --git a/src/mongo/db/pipeline/process_interface_shardsvr.h b/src/mongo/db/pipeline/process_interface_shardsvr.h
index ddafac24cdb..fba537a2e1a 100644
--- a/src/mongo/db/pipeline/process_interface_shardsvr.h
+++ b/src/mongo/db/pipeline/process_interface_shardsvr.h
@@ -43,8 +43,19 @@ class MongoInterfaceShardServer final : public MongoInterfaceStandalone {
 public:
     using MongoInterfaceStandalone::MongoInterfaceStandalone;
 
-    std::pair<std::vector<FieldPath>, bool> collectDocumentKeyFields(
-        OperationContext* opCtx, NamespaceStringOrUUID nssOrUUID) const final;
+    std::pair<std::vector<FieldPath>, bool> collectDocumentKeyFieldsForHostedCollection(
+        OperationContext* opCtx, const NamespaceString&, UUID) const final;
+
+    std::vector<FieldPath> collectDocumentKeyFieldsActingAsRouter(
+        OperationContext*, const NamespaceString&) const final {
+        // We don't expect anyone to use this method on the shard itself (yet). This is currently
+        // only used for $out. For $out in a sharded cluster, the mongos is responsible for
+        // collecting the document key fields before serializing them and sending them to the
+        // shards. This is logically a MONGO_UNREACHABLE, but a malicious user could construct a
+        // request to send directly to the shards which does not include the uniqueKey, so we must
+        // be prepared to gracefully error.
+        uasserted(50997, "Unexpected attempt to consult catalog cache on a shard server");
+    }
 
     /**
      * Inserts the documents 'objs' into the namespace 'ns' using the ClusterWriter for locking,
diff --git a/src/mongo/db/pipeline/process_interface_standalone.cpp b/src/mongo/db/pipeline/process_interface_standalone.cpp
index a071a848a5a..86dfbd0f09a 100644
--- a/src/mongo/db/pipeline/process_interface_standalone.cpp
+++ b/src/mongo/db/pipeline/process_interface_standalone.cpp
@@ -247,7 +247,7 @@ void MongoInterfaceStandalone::renameIfOptionsAndIndexesHaveNotChanged(
     const NamespaceString& targetNs,
     const BSONObj& originalCollectionOptions,
     const std::list<BSONObj>& originalIndexes) {
-    Lock::GlobalWrite globalLock(opCtx);
+    Lock::DBLock(opCtx, targetNs.db(), MODE_X);
 
     uassert(ErrorCodes::CommandFailed,
             str::stream() << "collection options of target collection " << targetNs.ns()
@@ -341,11 +341,18 @@ std::string MongoInterfaceStandalone::getShardName(OperationContext* opCtx) cons
     return std::string();
 }
 
-std::pair<std::vector<FieldPath>, bool> MongoInterfaceStandalone::collectDocumentKeyFields(
-    OperationContext* opCtx, NamespaceStringOrUUID nssOrUUID) const {
+std::pair<std::vector<FieldPath>, bool>
+MongoInterfaceStandalone::collectDocumentKeyFieldsForHostedCollection(OperationContext* opCtx,
+                                                                      const NamespaceString& nss,
+                                                                      UUID uuid) const {
     return {{"_id"}, false};  // Nothing is sharded.
 }
 
+std::vector<FieldPath> MongoInterfaceStandalone::collectDocumentKeyFieldsActingAsRouter(
+    OperationContext* opCtx, const NamespaceString& nss) const {
+    return {"_id"};  // Nothing is sharded.
+}
+
 std::vector<GenericCursor> MongoInterfaceStandalone::getIdleCursors(
     const intrusive_ptr<ExpressionContext>& expCtx, CurrentOpUserMode userMode) const {
     return CursorManager::getIdleCursors(expCtx->opCtx, userMode);
diff --git a/src/mongo/db/pipeline/process_interface_standalone.h b/src/mongo/db/pipeline/process_interface_standalone.h
index 257a42440e1..cdbbcd93c9d 100644
--- a/src/mongo/db/pipeline/process_interface_standalone.h
+++ b/src/mongo/db/pipeline/process_interface_standalone.h
@@ -94,8 +94,10 @@ public:
     Status attachCursorSourceToPipeline(const boost::intrusive_ptr<ExpressionContext>& expCtx,
                                         Pipeline* pipeline) final;
     std::string getShardName(OperationContext* opCtx) const final;
-    std::pair<std::vector<FieldPath>, bool> collectDocumentKeyFields(
-        OperationContext* opCtx, NamespaceStringOrUUID nssOrUUID) const override;
+    std::pair<std::vector<FieldPath>, bool> collectDocumentKeyFieldsForHostedCollection(
+        OperationContext* opCtx, const NamespaceString&, UUID) const override;
+    std::vector<FieldPath> collectDocumentKeyFieldsActingAsRouter(
+        OperationContext* opCtx, const NamespaceString&) const override;
     boost::optional<Document> lookupSingleDocument(
         const boost::intrusive_ptr<ExpressionContext>& expCtx,
         const NamespaceString& nss,
diff --git a/src/mongo/db/pipeline/stub_mongo_process_interface.h b/src/mongo/db/pipeline/stub_mongo_process_interface.h
index 906a750c40a..bdcea8dcd48 100644
--- a/src/mongo/db/pipeline/stub_mongo_process_interface.h
+++ b/src/mongo/db/pipeline/stub_mongo_process_interface.h
@@ -141,8 +141,13 @@ public:
         MONGO_UNREACHABLE;
     }
 
-    std::pair<std::vector<FieldPath>, bool> collectDocumentKeyFields(
-        OperationContext* opCtx, NamespaceStringOrUUID nssOrUUID) const override {
+    std::pair<std::vector<FieldPath>, bool> collectDocumentKeyFieldsForHostedCollection(
+        OperationContext*, const NamespaceString&, UUID) const override {
+        MONGO_UNREACHABLE;
+    }
+
+    std::vector<FieldPath> collectDocumentKeyFieldsActingAsRouter(
+        OperationContext*, const NamespaceString&) const override {
         MONGO_UNREACHABLE;
     }