SERVER-32666 Allow $graphLookup against a sharded collection to execute on mongos

9 files changed, 11 insertions, 39 deletions

diff --git a/jstests/aggregation/sources/graphLookup/airports.js b/jstests/aggregation/sources/graphLookup/airports.js
index 9254fd992fa..71a38e268b8 100644
--- a/jstests/aggregation/sources/graphLookup/airports.js
+++ b/jstests/aggregation/sources/graphLookup/airports.js
@@ -1,7 +1,3 @@
-// Cannot implicitly shard accessed collections because unsupported use of sharded collection
-// for target collection of $lookup and $graphLookup.
-// @tags: [assumes_unsharded_collection]
-
 // In MongoDB 3.4, $graphLookup was introduced. In this file, we test some complex graphs.
 
 (function() {
diff --git a/jstests/aggregation/sources/graphLookup/basic.js b/jstests/aggregation/sources/graphLookup/basic.js
index c0bcb1a8a53..63b2a63f95a 100644
--- a/jstests/aggregation/sources/graphLookup/basic.js
+++ b/jstests/aggregation/sources/graphLookup/basic.js
@@ -1,7 +1,3 @@
-// Cannot implicitly shard accessed collections because unsupported use of sharded collection
-// for target collection of $lookup and $graphLookup.
-// @tags: [assumes_unsharded_collection]
-
 // In MongoDB 3.4, $graphLookup was introduced. In this file, we test basic behavior and correctness
 // of the stage.
 
diff --git a/jstests/aggregation/sources/graphLookup/collation_graphlookup.js b/jstests/aggregation/sources/graphLookup/collation_graphlookup.js
index 7b457289cc6..246163b1eed 100644
--- a/jstests/aggregation/sources/graphLookup/collation_graphlookup.js
+++ b/jstests/aggregation/sources/graphLookup/collation_graphlookup.js
@@ -1,7 +1,3 @@
-// Cannot implicitly shard accessed collections because unsupported use of sharded collection
-// for target collection of $lookup and $graphLookup.
-// @tags: [assumes_unsharded_collection]
-
 /**
  * Tests that the $graphLookup stage respects the collation when matching between the
  * 'connectFromField' and the 'connectToField'.  $graphLookup should use the collation
@@ -60,10 +56,11 @@
     assert.eq("erica", res[0].username);
     assert.eq(0, res[0].friendUsers.length);
 
-    coll.drop();
+    assert.commandWorked(db.runCommand({drop: coll.getName()}));
     assert.commandWorked(db.createCollection(coll.getName(), caseInsensitiveUS));
     assert.writeOK(coll.insert({username: "erica", friends: ["jeremy", "jimmy"]}));
-    foreignColl.drop();
+    assert.commandWorkedOrFailedWithCode(db.runCommand({drop: foreignColl.getName()}),
+                                         ErrorCodes.NamespaceNotFound);
     assert.commandWorked(db.createCollection(foreignColl.getName(), caseSensitiveUS));
     assert.writeOK(foreignColl.insert([{username: "JEREMY"}, {username: "JIMMY"}]));
 
diff --git a/jstests/aggregation/sources/graphLookup/error.js b/jstests/aggregation/sources/graphLookup/error.js
index 42d1203238c..1f58dcb32fb 100644
--- a/jstests/aggregation/sources/graphLookup/error.js
+++ b/jstests/aggregation/sources/graphLookup/error.js
@@ -1,7 +1,3 @@
-// Cannot implicitly shard accessed collections because unsupported use of sharded collection
-// for target collection of $lookup and $graphLookup.
-// @tags: [assumes_unsharded_collection]
-
 // In MongoDB 3.4, $graphLookup was introduced. In this file, we test the error cases.
 load("jstests/aggregation/extras/utils.js");  // For "assertErrorCode".
 
diff --git a/jstests/aggregation/sources/graphLookup/filter.js b/jstests/aggregation/sources/graphLookup/filter.js
index 69027500aae..7b65599eb74 100644
--- a/jstests/aggregation/sources/graphLookup/filter.js
+++ b/jstests/aggregation/sources/graphLookup/filter.js
@@ -1,7 +1,3 @@
-// Cannot implicitly shard accessed collections because unsupported use of sharded collection
-// for target collection of $lookup and $graphLookup.
-// @tags: [assumes_unsharded_collection]
-
 // In SERVER-24714, the 'restrictSearchWithMatch' option was added to $graphLookup. In this file,
 // we test the functionality and correctness of the option.
 
diff --git a/jstests/aggregation/sources/graphLookup/nested_objects.js b/jstests/aggregation/sources/graphLookup/nested_objects.js
index d40cced2ac4..b2e3f5ac59a 100644
--- a/jstests/aggregation/sources/graphLookup/nested_objects.js
+++ b/jstests/aggregation/sources/graphLookup/nested_objects.js
@@ -1,7 +1,3 @@
-// Cannot implicitly shard accessed collections because unsupported use of sharded collection
-// for target collection of $lookup and $graphLookup.
-// @tags: [assumes_unsharded_collection]
-
 // In MongoDB 3.4, $graphLookup was introduced. In this file, we test the behavior of graphLookup
 // when the 'connectToField' is a nested array, or when the 'connectFromField' is a nested array.
 
diff --git a/jstests/aggregation/sources/graphLookup/socialite.js b/jstests/aggregation/sources/graphLookup/socialite.js
index 228c0f56c0e..6efe2672540 100644
--- a/jstests/aggregation/sources/graphLookup/socialite.js
+++ b/jstests/aggregation/sources/graphLookup/socialite.js
@@ -1,7 +1,3 @@
-// Cannot implicitly shard accessed collections because unsupported use of sharded collection
-// for target collection of $lookup and $graphLookup.
-// @tags: [assumes_unsharded_collection]
-
 // In MongoDB 3.4, $graphLookup was introduced. In this file, we test $graphLookup as applied to the
 // Socialite schema example available here: https://github.com/mongodb-labs/socialite
 
diff --git a/src/mongo/db/pipeline/document_source_graph_lookup.h b/src/mongo/db/pipeline/document_source_graph_lookup.h
index 5f1af60d2ce..1c7cdd55924 100644
--- a/src/mongo/db/pipeline/document_source_graph_lookup.h
+++ b/src/mongo/db/pipeline/document_source_graph_lookup.h
@@ -56,9 +56,16 @@ public:
     GetModPathsReturn getModifiedPaths() const final;
 
     StageConstraints constraints(Pipeline::SplitState pipeState) const final {
+        // TODO SERVER-27533 Until we remove the restriction of only performing lookups from mongos,
+        // this stage must run on mongos if the output collection is sharded.
+        HostTypeRequirement hostRequirement =
+            (pExpCtx->inMongos && pExpCtx->mongoProcessInterface->isSharded(pExpCtx->opCtx, _from))
+            ? HostTypeRequirement::kMongoS
+            : HostTypeRequirement::kPrimaryShard;
+
         StageConstraints constraints(StreamType::kStreaming,
                                      PositionRequirement::kNone,
-                                     HostTypeRequirement::kPrimaryShard,
+                                     hostRequirement,
                                      DiskUseRequirement::kNoDiskUse,
                                      FacetRequirement::kAllowed,
                                      TransactionRequirement::kAllowed);
diff --git a/src/mongo/db/pipeline/lite_parsed_document_source.h b/src/mongo/db/pipeline/lite_parsed_document_source.h
index 8f4b85168cc..fc3b88163f3 100644
--- a/src/mongo/db/pipeline/lite_parsed_document_source.h
+++ b/src/mongo/db/pipeline/lite_parsed_document_source.h
@@ -189,14 +189,6 @@ public:
         return _requiredPrivileges;
     }
 
-    /**
-     * Returns true if 'nss' is in the list of foreign namespaces for this DocumentSource. By
-     * default, no involved namespace is allowed to be sharded.
-     */
-    bool allowShardedForeignCollection(NamespaceString nss) const {
-        return (_foreignNssSet.find(nss) == _foreignNssSet.end());
-    }
-
 protected:
     stdx::unordered_set<NamespaceString> _foreignNssSet;