SERVER-22760 Sharded aggregation pipelines which involve taking a simple union should merge on mongos

1 files changed, 12 insertions, 4 deletions

diff --git a/src/mongo/db/pipeline/document_source.h b/src/mongo/db/pipeline/document_source.h
index efb695c8f59..8f1c2fef4a0 100644
--- a/src/mongo/db/pipeline/document_source.h
+++ b/src/mongo/db/pipeline/document_source.h
@@ -128,9 +128,18 @@ public:
          */
         enum class PositionRequirement { kNone, kFirst, kLast };
 
+        /**
+         * A HostTypeRequirement defines where this stage is permitted to be executed when the
+         * pipeline is run on a sharded cluster.
+         */
+        enum class HostTypeRequirement { kPrimaryShard, kAnyShard, kAnyShardOrMongoS };
+
         // Set if this stage needs to be in a particular position of the pipeline.
         PositionRequirement requiredPosition = PositionRequirement::kNone;
 
+        // Set if this stage can only be executed on specific components of a sharded cluster.
+        HostTypeRequirement hostRequirement = HostTypeRequirement::kAnyShard;
+
         bool isAllowedInsideFacetStage = true;
 
         // True if this stage does not generate results itself, and instead pulls inputs from an
@@ -147,12 +156,11 @@ public:
         // must also override getModifiedPaths() to provide information about which particular
         // $match predicates be swapped before itself.
         bool canSwapWithMatch = false;
-
-        // True if this stage must run on the primary shard when the collection being aggregated is
-        // sharded.
-        bool mustRunOnPrimaryShardIfSharded = false;
     };
 
+    using HostTypeRequirement = StageConstraints::HostTypeRequirement;
+    using PositionRequirement = StageConstraints::PositionRequirement;
+
     /**
      * This is what is returned from the main DocumentSource API: getNext(). It is essentially a
      * (ReturnStatus, Document) pair, with the first entry being used to communicate information