SERVER-21510 trip invariant on attempt to create PlanStages with a no-op extension context

11 files changed, 65 insertions, 17 deletions

diff --git a/src/mongo/db/exec/cached_plan.cpp b/src/mongo/db/exec/cached_plan.cpp
index 43ba65a729f..0f6d2b9f9be 100644
--- a/src/mongo/db/exec/cached_plan.cpp
+++ b/src/mongo/db/exec/cached_plan.cpp
@@ -230,7 +230,8 @@ Status CachedPlanStage::replan(PlanYieldPolicy* yieldPolicy, bool shouldCache) {
 
         PlanStage* newRoot;
         // Only one possible plan. Build the stages from the solution.
-        verify(StageBuilder::build(getOpCtx(), _collection, *solutions[0], _ws, &newRoot));
+        verify(StageBuilder::build(
+            getOpCtx(), _collection, *_canonicalQuery, *solutions[0], _ws, &newRoot));
         _children.emplace_back(newRoot);
         _replannedQs.reset(solutions.popAndReleaseBack());
         return Status::OK();
@@ -250,7 +251,8 @@ Status CachedPlanStage::replan(PlanYieldPolicy* yieldPolicy, bool shouldCache) {
         }
 
         PlanStage* nextPlanRoot;
-        verify(StageBuilder::build(getOpCtx(), _collection, *solutions[ix], _ws, &nextPlanRoot));
+        verify(StageBuilder::build(
+            getOpCtx(), _collection, *_canonicalQuery, *solutions[ix], _ws, &nextPlanRoot));
 
         // Takes ownership of 'solutions[ix]' and 'nextPlanRoot'.
         multiPlanStage->addPlan(solutions.releaseAt(ix), nextPlanRoot, _ws);
diff --git a/src/mongo/db/exec/subplan.cpp b/src/mongo/db/exec/subplan.cpp
index c77b7d05454..df514861011 100644
--- a/src/mongo/db/exec/subplan.cpp
+++ b/src/mongo/db/exec/subplan.cpp
@@ -327,8 +327,12 @@ Status SubplanStage::choosePlanForSubqueries(PlanYieldPolicy* yieldPolicy) {
             // Dump all the solutions into the MPS.
             for (size_t ix = 0; ix < branchResult->solutions.size(); ++ix) {
                 PlanStage* nextPlanRoot;
-                invariant(StageBuilder::build(
-                    getOpCtx(), _collection, *branchResult->solutions[ix], _ws, &nextPlanRoot));
+                invariant(StageBuilder::build(getOpCtx(),
+                                              _collection,
+                                              *branchResult->canonicalQuery,
+                                              *branchResult->solutions[ix],
+                                              _ws,
+                                              &nextPlanRoot));
 
                 // Takes ownership of solution with index 'ix' and 'nextPlanRoot'.
                 multiPlanStage.addPlan(branchResult->solutions.releaseAt(ix), nextPlanRoot, _ws);
@@ -407,7 +411,8 @@ Status SubplanStage::choosePlanForSubqueries(PlanYieldPolicy* yieldPolicy) {
     // and set that solution as our child stage.
     _ws->clear();
     PlanStage* root;
-    invariant(StageBuilder::build(getOpCtx(), _collection, *_compositeSolution.get(), _ws, &root));
+    invariant(StageBuilder::build(
+        getOpCtx(), _collection, *_query, *_compositeSolution.get(), _ws, &root));
     invariant(_children.empty());
     _children.emplace_back(root);
 
@@ -440,7 +445,7 @@ Status SubplanStage::choosePlanWholeQuery(PlanYieldPolicy* yieldPolicy) {
     if (1 == solutions.size()) {
         PlanStage* root;
         // Only one possible plan.  Run it.  Build the stages from the solution.
-        verify(StageBuilder::build(getOpCtx(), _collection, *solutions[0], _ws, &root));
+        verify(StageBuilder::build(getOpCtx(), _collection, *_query, *solutions[0], _ws, &root));
         invariant(_children.empty());
         _children.emplace_back(root);
 
@@ -462,8 +467,8 @@ Status SubplanStage::choosePlanWholeQuery(PlanYieldPolicy* yieldPolicy) {
 
             // version of StageBuild::build when WorkingSet is shared
             PlanStage* nextPlanRoot;
-            verify(
-                StageBuilder::build(getOpCtx(), _collection, *solutions[ix], _ws, &nextPlanRoot));
+            verify(StageBuilder::build(
+                getOpCtx(), _collection, *_query, *solutions[ix], _ws, &nextPlanRoot));
 
             // Takes ownership of 'solutions[ix]' and 'nextPlanRoot'.
             multiPlanStage->addPlan(solutions.releaseAt(ix), nextPlanRoot, _ws);
diff --git a/src/mongo/db/matcher/extensions_callback.h b/src/mongo/db/matcher/extensions_callback.h
index a0cab74b652..237e882ca71 100644
--- a/src/mongo/db/matcher/extensions_callback.h
+++ b/src/mongo/db/matcher/extensions_callback.h
@@ -40,11 +40,21 @@ namespace mongo {
  */
 class ExtensionsCallback {
 public:
+    virtual ~ExtensionsCallback() {}
+
     virtual StatusWithMatchExpression parseText(BSONElement text) const = 0;
 
     virtual StatusWithMatchExpression parseWhere(BSONElement where) const = 0;
 
-    virtual ~ExtensionsCallback() {}
+    /**
+     * Returns true if extensions (e.g. $text and $where) are allowed but are converted into no-ops.
+     *
+     * Queries with a no-op extension context are special because they can be parsed and planned,
+     * but they cannot be executed.
+     */
+    virtual bool hasNoopExtensions() const {
+        return false;
+    }
 
 protected:
     /**
diff --git a/src/mongo/db/matcher/extensions_callback_noop.h b/src/mongo/db/matcher/extensions_callback_noop.h
index 127e1d422a6..e68ec03e432 100644
--- a/src/mongo/db/matcher/extensions_callback_noop.h
+++ b/src/mongo/db/matcher/extensions_callback_noop.h
@@ -48,6 +48,10 @@ public:
      * Returns a WhereNoOpMatchExpression, or an error Status if parsing fails.
      */
     StatusWithMatchExpression parseWhere(BSONElement where) const final;
+
+    bool hasNoopExtensions() const final {
+        return true;
+    }
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/query/canonical_query.cpp b/src/mongo/db/query/canonical_query.cpp
index ed22674e0db..a48b846c044 100644
--- a/src/mongo/db/query/canonical_query.cpp
+++ b/src/mongo/db/query/canonical_query.cpp
@@ -314,6 +314,8 @@ Status CanonicalQuery::init(LiteParsedQuery* lpq,
                             MatchExpression* root) {
     _pq.reset(lpq);
 
+    _hasNoopExtensions = extensionsCallback.hasNoopExtensions();
+
     // Normalize, sort and validate tree.
     root = normalizeTree(root);
 
diff --git a/src/mongo/db/query/canonical_query.h b/src/mongo/db/query/canonical_query.h
index 136c6c1d8f8..3936809ad05 100644
--- a/src/mongo/db/query/canonical_query.h
+++ b/src/mongo/db/query/canonical_query.h
@@ -194,6 +194,17 @@ public:
      */
     static size_t countNodes(const MatchExpression* root, MatchExpression::MatchType type);
 
+    /**
+     * Returns true if this canonical query converted extensions such as $where and $text into
+     * no-ops during parsing.
+     *
+     * Queries with a no-op extension context are special because they can be parsed and planned,
+     * but they cannot be executed.
+     */
+    bool hasNoopExtensions() const {
+        return _hasNoopExtensions;
+    }
+
 private:
     // You must go through canonicalize to create a CanonicalQuery.
     CanonicalQuery() {}
@@ -211,6 +222,8 @@ private:
     std::unique_ptr<MatchExpression> _root;
 
     std::unique_ptr<ParsedProjection> _proj;
+
+    bool _hasNoopExtensions = false;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/query/get_executor.cpp b/src/mongo/db/query/get_executor.cpp
index dd973bee36f..f0290933c0d 100644
--- a/src/mongo/db/query/get_executor.cpp
+++ b/src/mongo/db/query/get_executor.cpp
@@ -315,7 +315,7 @@ Status prepareExecution(OperationContext* opCtx,
         Status status = QueryPlanner::planFromCache(*canonicalQuery, plannerParams, *cs, &qs);
 
         if (status.isOK()) {
-            verify(StageBuilder::build(opCtx, collection, *qs, ws, rootOut));
+            verify(StageBuilder::build(opCtx, collection, *canonicalQuery, *qs, ws, rootOut));
             if ((plannerParams.options & QueryPlannerParams::PRIVATE_IS_COUNT) &&
                 turnIxscanIntoCount(qs)) {
                 LOG(2) << "Using fast count: " << canonicalQuery->toStringShort()
@@ -371,7 +371,8 @@ Status prepareExecution(OperationContext* opCtx,
                 }
 
                 // We're not going to cache anything that's fast count.
-                verify(StageBuilder::build(opCtx, collection, *solutions[i], ws, rootOut));
+                verify(StageBuilder::build(
+                    opCtx, collection, *canonicalQuery, *solutions[i], ws, rootOut));
 
                 LOG(2) << "Using fast count: " << canonicalQuery->toStringShort()
                        << ", planSummary: " << Explain::getPlanSummary(*rootOut);
@@ -384,7 +385,7 @@ Status prepareExecution(OperationContext* opCtx,
 
     if (1 == solutions.size()) {
         // Only one possible plan.  Run it.  Build the stages from the solution.
-        verify(StageBuilder::build(opCtx, collection, *solutions[0], ws, rootOut));
+        verify(StageBuilder::build(opCtx, collection, *canonicalQuery, *solutions[0], ws, rootOut));
 
         LOG(2) << "Only one plan is available; it will be run but will not be cached. "
                << canonicalQuery->toStringShort()
@@ -404,7 +405,8 @@ Status prepareExecution(OperationContext* opCtx,
 
             // version of StageBuild::build when WorkingSet is shared
             PlanStage* nextPlanRoot;
-            verify(StageBuilder::build(opCtx, collection, *solutions[ix], ws, &nextPlanRoot));
+            verify(StageBuilder::build(
+                opCtx, collection, *canonicalQuery, *solutions[ix], ws, &nextPlanRoot));
 
             // Owns none of the arguments
             multiPlanStage->addPlan(solutions[ix], nextPlanRoot, ws);
@@ -1378,7 +1380,7 @@ StatusWith<unique_ptr<PlanExecutor>> getExecutorDistinct(OperationContext* txn,
 
         unique_ptr<WorkingSet> ws = make_unique<WorkingSet>();
         PlanStage* rawRoot;
-        verify(StageBuilder::build(txn, collection, *soln, ws.get(), &rawRoot));
+        verify(StageBuilder::build(txn, collection, *cq, *soln, ws.get(), &rawRoot));
         unique_ptr<PlanStage> root(rawRoot);
 
         LOG(2) << "Using fast distinct: " << cq->toStringShort()
@@ -1414,7 +1416,7 @@ StatusWith<unique_ptr<PlanExecutor>> getExecutorDistinct(OperationContext* txn,
             unique_ptr<WorkingSet> ws = make_unique<WorkingSet>();
             unique_ptr<QuerySolution> currentSolution(solutions[i]);
             PlanStage* rawRoot;
-            verify(StageBuilder::build(txn, collection, *currentSolution, ws.get(), &rawRoot));
+            verify(StageBuilder::build(txn, collection, *cq, *currentSolution, ws.get(), &rawRoot));
             unique_ptr<PlanStage> root(rawRoot);
 
             LOG(2) << "Using fast distinct: " << cq->toStringShort()
diff --git a/src/mongo/db/query/stage_builder.cpp b/src/mongo/db/query/stage_builder.cpp
index 429651668a2..f26e9d82187 100644
--- a/src/mongo/db/query/stage_builder.cpp
+++ b/src/mongo/db/query/stage_builder.cpp
@@ -345,9 +345,16 @@ PlanStage* buildStages(OperationContext* txn,
 // static (this one is used for Cached and MultiPlanStage)
 bool StageBuilder::build(OperationContext* txn,
                          Collection* collection,
+                         const CanonicalQuery& cq,
                          const QuerySolution& solution,
                          WorkingSet* wsIn,
                          PlanStage** rootOut) {
+    // Only QuerySolutions derived from queries parsed with context, or QuerySolutions derived from
+    // queries that disallow extensions, can be properly executed. If the query does not have
+    // $text/$where context (and $text/$where are allowed), then no attempt should be made to
+    // execute the query.
+    invariant(!cq.hasNoopExtensions());
+
     if (NULL == wsIn || NULL == rootOut) {
         return false;
     }
diff --git a/src/mongo/db/query/stage_builder.h b/src/mongo/db/query/stage_builder.h
index c490b9a974d..ee5ae17bb75 100644
--- a/src/mongo/db/query/stage_builder.h
+++ b/src/mongo/db/query/stage_builder.h
@@ -44,6 +44,8 @@ public:
     /**
      * Turns 'solution' into an executable tree of PlanStage(s).
      *
+     * 'cq' must be the CanonicalQuery from which 'solution' is derived.
+     *
      * Returns true if the PlanStage tree was built successfully.  The root of the tree is in
      * *rootOut and the WorkingSet that the tree uses is in wsIn.
      *
@@ -51,6 +53,7 @@ public:
      */
     static bool build(OperationContext* txn,
                       Collection* collection,
+                      const CanonicalQuery& cq,
                       const QuerySolution& solution,
                       WorkingSet* wsIn,
                       PlanStage** rootOut);
diff --git a/src/mongo/dbtests/plan_ranking.cpp b/src/mongo/dbtests/plan_ranking.cpp
index 33b47fb257e..5012816caa4 100644
--- a/src/mongo/dbtests/plan_ranking.cpp
+++ b/src/mongo/dbtests/plan_ranking.cpp
@@ -122,7 +122,7 @@ public:
         // Put each solution from the planner into the MPR.
         for (size_t i = 0; i < solutions.size(); ++i) {
             PlanStage* root;
-            ASSERT(StageBuilder::build(&_txn, collection, *solutions[i], ws.get(), &root));
+            ASSERT(StageBuilder::build(&_txn, collection, *cq, *solutions[i], ws.get(), &root));
             // Takes ownership of all (actually some) arguments.
             _mps->addPlan(solutions[i], root, ws.get());
         }
diff --git a/src/mongo/dbtests/query_stage_multiplan.cpp b/src/mongo/dbtests/query_stage_multiplan.cpp
index 2de6a2bb126..816bae1b203 100644
--- a/src/mongo/dbtests/query_stage_multiplan.cpp
+++ b/src/mongo/dbtests/query_stage_multiplan.cpp
@@ -243,7 +243,7 @@ public:
         // Put each solution from the planner into the MPR.
         for (size_t i = 0; i < solutions.size(); ++i) {
             PlanStage* root;
-            ASSERT(StageBuilder::build(&_txn, collection, *solutions[i], ws.get(), &root));
+            ASSERT(StageBuilder::build(&_txn, collection, *cq, *solutions[i], ws.get(), &root));
             // Takes ownership of 'solutions[i]' and 'root'.
             mps->addPlan(solutions[i], root, ws.get());
         }