SERVER-32308: Add the ability for a $lookup stage to execute on mongos against a sharded foreign collection

1 files changed, 35 insertions, 2 deletions

diff --git a/src/mongo/db/pipeline/document_source_lookup.cpp b/src/mongo/db/pipeline/document_source_lookup.cpp
index e4fd70920f6..ecc9d572bfe 100644
--- a/src/mongo/db/pipeline/document_source_lookup.cpp
+++ b/src/mongo/db/pipeline/document_source_lookup.cpp
@@ -283,8 +283,8 @@ std::unique_ptr<Pipeline, PipelineDeleter> DocumentSourceLookUp::buildPipeline(
 
     if (!_cache->isServing()) {
         // The cache has either been abandoned or has not yet been built. Attach a cursor.
-        uassertStatusOK(pExpCtx->mongoProcessInterface->attachCursorSourceToPipeline(
-            _fromExpCtx, pipeline.get()));
+        pipeline = uassertStatusOK(pExpCtx->mongoProcessInterface->attachCursorSourceToPipeline(
+            _fromExpCtx, pipeline.release()));
     }
 
     // If the cache has been abandoned, release it.
@@ -614,6 +614,39 @@ void DocumentSourceLookUp::initializeIntrospectionPipeline() {
             sources.empty() || !sources.front()->constraints().isChangeStreamStage());
 }
 
+DocumentSource::StageConstraints DocumentSourceLookUp::constraints(
+    Pipeline::SplitState pipeState) const {
+
+    const bool mayUseDisk = wasConstructedWithPipelineSyntax() &&
+        std::any_of(_parsedIntrospectionPipeline->getSources().begin(),
+                    _parsedIntrospectionPipeline->getSources().end(),
+                    [](const auto& source) {
+                        return source->constraints().diskRequirement ==
+                            DiskUseRequirement::kWritesTmpData;
+                    });
+
+    // If executing on mongos and the foreign collection is sharded, then this stage can run on
+    // mongos.
+    HostTypeRequirement hostReq;
+    if (pExpCtx->inMongos) {
+        hostReq = pExpCtx->mongoProcessInterface->isSharded(pExpCtx->opCtx, _fromNs)
+            ? HostTypeRequirement::kMongoS
+            : HostTypeRequirement::kPrimaryShard;
+    } else {
+        hostReq = HostTypeRequirement::kPrimaryShard;
+    }
+
+    StageConstraints constraints(StreamType::kStreaming,
+                                 PositionRequirement::kNone,
+                                 hostReq,
+                                 mayUseDisk ? DiskUseRequirement::kWritesTmpData
+                                            : DiskUseRequirement::kNoDiskUse,
+                                 FacetRequirement::kAllowed);
+
+    constraints.canSwapWithMatch = true;
+    return constraints;
+}
+
 void DocumentSourceLookUp::serializeToArray(
     std::vector<Value>& array, boost::optional<ExplainOptions::Verbosity> explain) const {
     Document doc;