SERVER-39166 $graphLookup should force a pipeline to split in sharded cluster

3 files changed, 47 insertions, 1 deletions

diff --git a/buildscripts/resmokeconfig/suites/sharding_last_stable_mongos_and_mixed_shards.yml b/buildscripts/resmokeconfig/suites/sharding_last_stable_mongos_and_mixed_shards.yml
index c7ff1b24395..628b51f1a00 100644
--- a/buildscripts/resmokeconfig/suites/sharding_last_stable_mongos_and_mixed_shards.yml
+++ b/buildscripts/resmokeconfig/suites/sharding_last_stable_mongos_and_mixed_shards.yml
@@ -97,6 +97,8 @@ selector:
     - jstests/sharding/views.js
     # SERVER-20392: Reenable when backported to 3.4 and released as last-stable.
     - jstests/sharding/shard_existing_coll_chunk_count.js
+    # $graphLookup is new in 3.4
+    - jstests/sharding/graph_lookup.js
 
 executor:
   js_test:
diff --git a/jstests/sharding/graph_lookup.js b/jstests/sharding/graph_lookup.js
new file mode 100644
index 00000000000..421a2a4acf6
--- /dev/null
+++ b/jstests/sharding/graph_lookup.js
@@ -0,0 +1,35 @@
+// Test aggregating a sharded collection while using $graphLookup on an unsharded collection.
+(function() {
+    'use strict';
+
+    const st = new ShardingTest({shards: 2, rs: {nodes: 1}});
+
+    assert.commandWorked(st.s0.adminCommand({enablesharding: "test"}));
+    assert.commandWorked(st.s0.adminCommand({shardCollection: "test.foo", key: {_id: "hashed"}}));
+
+    let db = st.s0.getDB("test");
+
+    assert.writeOK(db.foo.insert([{}, {}, {}, {}]));
+    assert.writeOK(db.bar.insert({_id: 1, x: 1}));
+
+    const res = db.foo
+                    .aggregate([{
+                        $graphLookup: {
+                            from: "bar",
+                            startWith: {$literal: 1},
+                            connectFromField: "x",
+                            connectToField: "_id",
+                            as: "res"
+                        }
+                    }])
+                    .toArray();
+
+    assert.eq(res.length, 4);
+    res.forEach(function(c) {
+        assert.eq(c.res.length, 1);
+        assert.eq(c.res[0]._id, 1);
+        assert.eq(c.res[0].x, 1);
+    });
+
+    st.stop();
+})();
diff --git a/src/mongo/db/pipeline/document_source_graph_lookup.h b/src/mongo/db/pipeline/document_source_graph_lookup.h
index be12637d404..661593c3482 100644
--- a/src/mongo/db/pipeline/document_source_graph_lookup.h
+++ b/src/mongo/db/pipeline/document_source_graph_lookup.h
@@ -36,7 +36,8 @@
 
 namespace mongo {
 
-class DocumentSourceGraphLookUp final : public DocumentSourceNeedsMongod {
+class DocumentSourceGraphLookUp final : public DocumentSourceNeedsMongod,
+                                        public SplittableDocumentSource {
 public:
     static std::unique_ptr<LiteParsedDocumentSourceOneForeignCollection> liteParse(
         const AggregationRequest& request, const BSONElement& spec);
@@ -71,6 +72,14 @@ public:
         return true;
     }
 
+    boost::intrusive_ptr<DocumentSource> getShardSource() final {
+        return nullptr;
+    }
+
+    boost::intrusive_ptr<DocumentSource> getMergeSource() final {
+        return this;
+    }
+
     void addInvolvedCollections(std::vector<NamespaceString>* collections) const final {
         collections->push_back(_from);
     }