SERVER-36667: Add ability to have a mixture of sharded and unsharded foreign namespaces in an aggregation

8 files changed, 58 insertions, 31 deletions

diff --git a/src/mongo/db/pipeline/document_source_facet.cpp b/src/mongo/db/pipeline/document_source_facet.cpp
index e3c1f8480c1..4f2a2dd8547 100644
--- a/src/mongo/db/pipeline/document_source_facet.cpp
+++ b/src/mongo/db/pipeline/document_source_facet.cpp
@@ -141,6 +141,13 @@ stdx::unordered_set<NamespaceString> DocumentSourceFacet::LiteParsed::getInvolve
     return involvedNamespaces;
 }
 
+bool DocumentSourceFacet::LiteParsed::allowShardedForeignCollection(NamespaceString nss) const {
+    return std::all_of(
+        _liteParsedPipelines.begin(), _liteParsedPipelines.end(), [&nss](auto&& pipeline) {
+            return pipeline.allowShardedForeignCollection(nss);
+        });
+}
+
 REGISTER_DOCUMENT_SOURCE(facet,
                          DocumentSourceFacet::LiteParsed::parse,
                          DocumentSourceFacet::createFromBson);
diff --git a/src/mongo/db/pipeline/document_source_facet.h b/src/mongo/db/pipeline/document_source_facet.h
index 92e9f2f05d2..3ef1d716f01 100644
--- a/src/mongo/db/pipeline/document_source_facet.h
+++ b/src/mongo/db/pipeline/document_source_facet.h
@@ -80,6 +80,8 @@ public:
 
         stdx::unordered_set<NamespaceString> getInvolvedNamespaces() const final;
 
+        bool allowShardedForeignCollection(NamespaceString nss) const final;
+
     private:
         const std::vector<LiteParsedPipeline> _liteParsedPipelines;
         const PrivilegeVector _requiredPrivileges;
diff --git a/src/mongo/db/pipeline/document_source_graph_lookup.cpp b/src/mongo/db/pipeline/document_source_graph_lookup.cpp
index cabdd753cae..f99153aa6cc 100644
--- a/src/mongo/db/pipeline/document_source_graph_lookup.cpp
+++ b/src/mongo/db/pipeline/document_source_graph_lookup.cpp
@@ -74,9 +74,8 @@ std::unique_ptr<LiteParsedDocumentSourceForeignCollections> DocumentSourceGraphL
     PrivilegeVector privileges{
         Privilege(ResourcePattern::forExactNamespace(nss), ActionType::find)};
 
-    constexpr bool allowedSharded = false;
-    return std::make_unique<LiteParsedDocumentSourceForeignCollections>(
-        std::move(nss), std::move(privileges), allowedSharded);
+    return std::make_unique<LiteParsedDocumentSourceForeignCollections>(std::move(nss),
+                                                                        std::move(privileges));
 }
 
 REGISTER_DOCUMENT_SOURCE(graphLookup,
diff --git a/src/mongo/db/pipeline/document_source_lookup.h b/src/mongo/db/pipeline/document_source_lookup.h
index d059cbc3a9e..5276aa917c4 100644
--- a/src/mongo/db/pipeline/document_source_lookup.h
+++ b/src/mongo/db/pipeline/document_source_lookup.h
@@ -82,8 +82,8 @@ public:
         /**
          * Lookup from a sharded collection is not allowed.
          */
-        bool allowShardedForeignCollections() const final {
-            return false;
+        bool allowShardedForeignCollection(NamespaceString nss) const final {
+            return (_foreignNssSet.find(nss) == _foreignNssSet.end());
         }
 
     private:
diff --git a/src/mongo/db/pipeline/document_source_out.cpp b/src/mongo/db/pipeline/document_source_out.cpp
index 8d8d3db49e8..181233a6aed 100644
--- a/src/mongo/db/pipeline/document_source_out.cpp
+++ b/src/mongo/db/pipeline/document_source_out.cpp
@@ -40,7 +40,7 @@ namespace mongo {
 using boost::intrusive_ptr;
 using std::vector;
 
-std::unique_ptr<LiteParsedDocumentSourceForeignCollections> DocumentSourceOut::liteParse(
+std::unique_ptr<DocumentSourceOut::LiteParsed> DocumentSourceOut::LiteParsed::parse(
     const AggregationRequest& request, const BSONElement& spec) {
 
     uassert(ErrorCodes::TypeMismatch,
@@ -79,11 +79,13 @@ std::unique_ptr<LiteParsedDocumentSourceForeignCollections> DocumentSourceOut::l
 
     PrivilegeVector privileges{Privilege(ResourcePattern::forExactNamespace(targetNss), actions)};
 
-    return stdx::make_unique<LiteParsedDocumentSourceForeignCollections>(
+    return stdx::make_unique<DocumentSourceOut::LiteParsed>(
         std::move(targetNss), std::move(privileges), allowSharded);
 }
 
-REGISTER_DOCUMENT_SOURCE(out, DocumentSourceOut::liteParse, DocumentSourceOut::createFromBson);
+REGISTER_DOCUMENT_SOURCE(out,
+                         DocumentSourceOut::LiteParsed::parse,
+                         DocumentSourceOut::createFromBson);
 
 const char* DocumentSourceOut::getSourceName() const {
     return "$out";
diff --git a/src/mongo/db/pipeline/document_source_out.h b/src/mongo/db/pipeline/document_source_out.h
index f8ee8829038..4ae2c9af5e8 100644
--- a/src/mongo/db/pipeline/document_source_out.h
+++ b/src/mongo/db/pipeline/document_source_out.h
@@ -38,8 +38,26 @@ namespace mongo {
  */
 class DocumentSourceOut : public DocumentSource, public NeedsMergerDocumentSource {
 public:
-    static std::unique_ptr<LiteParsedDocumentSourceForeignCollections> liteParse(
-        const AggregationRequest& request, const BSONElement& spec);
+    /**
+     * A "lite parsed" $out stage is similar to other stages involving foreign collections except in
+     * some cases the foreign collection is allowed to be sharded.
+     */
+    class LiteParsed final : public LiteParsedDocumentSourceForeignCollections {
+    public:
+        static std::unique_ptr<LiteParsed> parse(const AggregationRequest& request,
+                                                 const BSONElement& spec);
+
+        LiteParsed(NamespaceString outNss, PrivilegeVector privileges, bool allowShardedOutNss)
+            : LiteParsedDocumentSourceForeignCollections(outNss, privileges),
+              _allowShardedOutNss(allowShardedOutNss) {}
+
+        bool allowShardedForeignCollection(NamespaceString nss) const final {
+            return _allowShardedOutNss ? true : (_foreignNssSet.find(nss) == _foreignNssSet.end());
+        }
+
+    private:
+        bool _allowShardedOutNss;
+    };
 
     DocumentSourceOut(const NamespaceString& outputNs,
                       const boost::intrusive_ptr<ExpressionContext>& expCtx,
diff --git a/src/mongo/db/pipeline/lite_parsed_document_source.h b/src/mongo/db/pipeline/lite_parsed_document_source.h
index 4d439ad2f29..b5eafe7718f 100644
--- a/src/mongo/db/pipeline/lite_parsed_document_source.h
+++ b/src/mongo/db/pipeline/lite_parsed_document_source.h
@@ -128,10 +128,10 @@ public:
     }
 
     /**
-     * Returns true if the involved namespaces for this aggregation are allowed to be sharded.
+     * Returns true if the involved namespace 'nss' is allowed to be sharded.
      */
-    virtual bool allowShardedForeignCollections() const {
-        return false;
+    virtual bool allowShardedForeignCollection(NamespaceString nss) const {
+        return true;
     }
 
     /**
@@ -170,18 +170,12 @@ public:
 class LiteParsedDocumentSourceForeignCollections : public LiteParsedDocumentSource {
 public:
     LiteParsedDocumentSourceForeignCollections(NamespaceString foreignNss,
-                                               PrivilegeVector privileges,
-                                               bool allowSharded)
-        : _foreignNssSet{std::move(foreignNss)},
-          _requiredPrivileges(std::move(privileges)),
-          _allowSharded(allowSharded) {}
+                                               PrivilegeVector privileges)
+        : _foreignNssSet{std::move(foreignNss)}, _requiredPrivileges(std::move(privileges)) {}
 
     LiteParsedDocumentSourceForeignCollections(stdx::unordered_set<NamespaceString> foreignNssSet,
-                                               PrivilegeVector privileges,
-                                               bool allowSharded)
-        : _foreignNssSet(std::move(foreignNssSet)),
-          _requiredPrivileges(std::move(privileges)),
-          _allowSharded(allowSharded) {}
+                                               PrivilegeVector privileges)
+        : _foreignNssSet(std::move(foreignNssSet)), _requiredPrivileges(std::move(privileges)) {}
 
     stdx::unordered_set<NamespaceString> getInvolvedNamespaces() const final {
         return {_foreignNssSet};
@@ -191,13 +185,18 @@ public:
         return _requiredPrivileges;
     }
 
-    bool allowShardedForeignCollections() const final {
-        return _allowSharded;
+    /**
+     * 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());
     }
 
-private:
+protected:
     stdx::unordered_set<NamespaceString> _foreignNssSet;
+
+private:
     PrivilegeVector _requiredPrivileges;
-    bool _allowSharded;
 };
 }  // namespace mongo
diff --git a/src/mongo/db/pipeline/lite_parsed_pipeline.h b/src/mongo/db/pipeline/lite_parsed_pipeline.h
index 303754df4b0..82bcf9e8c7c 100644
--- a/src/mongo/db/pipeline/lite_parsed_pipeline.h
+++ b/src/mongo/db/pipeline/lite_parsed_pipeline.h
@@ -119,12 +119,12 @@ public:
     }
 
     /**
-     * Returns false if at least one of the stages does not allow an involved namespace to be
+     * Returns false if at least one of the stages does not allow the involved namespace 'nss' to be
      * sharded.
      */
-    bool allowShardedForeignCollections() const {
-        return std::all_of(_stageSpecs.begin(), _stageSpecs.end(), [](auto&& spec) {
-            return spec->allowShardedForeignCollections();
+    bool allowShardedForeignCollection(NamespaceString nss) const {
+        return std::all_of(_stageSpecs.begin(), _stageSpecs.end(), [&nss](auto&& spec) {
+            return spec->allowShardedForeignCollection(nss);
         });
     }