SERVER-20641 Equality query on _id with sort now works on mongos

3 files changed, 51 insertions, 67 deletions

diff --git a/jstests/core/sortl.js b/jstests/core/sortl.js
new file mode 100644
index 00000000000..b7cf9b34958
--- /dev/null
+++ b/jstests/core/sortl.js
@@ -0,0 +1,32 @@
+// Tests equality query on _id with a sort, intended to be tested on both mongos and mongod. For
+// SERVER-20641.
+(function() {
+    'use strict';
+    var coll = db.sortl;
+    coll.drop();
+
+    assert.writeOK(coll.insert({_id: 1, a: 2}));
+    var res = coll.find({_id: 1}).sort({a: 1});
+    assert.eq(res.next(), {_id: 1, a: 2});
+    assert.eq(res.hasNext(), false);
+
+    res = coll.find({_id: 1}, {b: {$meta: "sortKey"}}).sort({a: 1});
+    assert.eq(res.next(), {_id: 1, a: 2, b: {"": 2}});
+    assert.eq(res.hasNext(), false);
+
+    res = db.runCommand({findAndModify: coll.getName(),
+                         query: {_id: 1},
+                         update: {$set: {b: 1}},
+                         sort: {a: 1},
+                         fields: {c: {$meta: "sortKey"}}});
+    assert.commandFailedWithCode(res, ErrorCodes.BadValue, "$meta sortKey update");
+
+    res = db.runCommand({findAndModify: coll.getName(),
+                         query: {_id: 1},
+                         remove: true,
+                         sort: {b: 1},
+                         fields: {c: {$meta: "sortKey"}}});
+    assert.commandFailedWithCode(res, ErrorCodes.BadValue, "$meta sortKey delete");
+
+    coll.drop();
+})();
diff --git a/src/mongo/db/query/get_executor.cpp b/src/mongo/db/query/get_executor.cpp
index f0aaf0cf6ad..3e14fc037ae 100644
--- a/src/mongo/db/query/get_executor.cpp
+++ b/src/mongo/db/query/get_executor.cpp
@@ -48,6 +48,7 @@
 #include "mongo/db/exec/oplogstart.h"
 #include "mongo/db/exec/projection.h"
 #include "mongo/db/exec/shard_filter.h"
+#include "mongo/db/exec/sort_key_generator.h"
 #include "mongo/db/exec/subplan.h"
 #include "mongo/db/exec/update.h"
 #include "mongo/db/index_names.h"
@@ -267,9 +268,20 @@ Status prepareExecution(OperationContext* opCtx,
             ProjectionStageParams params(WhereCallbackReal(opCtx, collection->ns().db()));
             params.projObj = canonicalQuery->getProj()->getProjObj();
 
+            // Add a SortKeyGeneratorStage if there is a $meta sortKey projection.
+            if (canonicalQuery->getProj()->wantSortKey()) {
+                *rootOut = new SortKeyGeneratorStage(opCtx,
+                                                     *rootOut,
+                                                     ws,
+                                                     collection,
+                                                     canonicalQuery->getParsed().getSort(),
+                                                     canonicalQuery->getParsed().getFilter());
+            }
+
             // Stuff the right data into the params depending on what proj impl we use.
             if (canonicalQuery->getProj()->requiresDocument() ||
-                canonicalQuery->getProj()->wantIndexKey()) {
+                canonicalQuery->getProj()->wantIndexKey() ||
+                canonicalQuery->getProj()->wantSortKey()) {
                 params.fullExpression = canonicalQuery->root();
                 params.projImpl = ProjectionStageParams::NO_FAST_PATH;
             } else {
@@ -435,53 +447,6 @@ StatusWith<unique_ptr<PlanExecutor>> getExecutor(OperationContext* txn,
                               yieldPolicy);
 }
 
-StatusWith<unique_ptr<PlanExecutor>> getExecutor(OperationContext* txn,
-                                                 Collection* collection,
-                                                 const std::string& ns,
-                                                 const BSONObj& unparsedQuery,
-                                                 PlanExecutor::YieldPolicy yieldPolicy,
-                                                 size_t plannerOptions) {
-    if (!collection) {
-        LOG(2) << "Collection " << ns << " does not exist."
-               << " Using EOF stage: " << unparsedQuery.toString();
-        auto eofStage = make_unique<EOFStage>(txn);
-        auto ws = make_unique<WorkingSet>();
-        return PlanExecutor::make(txn, std::move(ws), std::move(eofStage), ns, yieldPolicy);
-    }
-
-    const IndexDescriptor* descriptor = collection->getIndexCatalog()->findIdIndex(txn);
-
-    if (!descriptor || !CanonicalQuery::isSimpleIdQuery(unparsedQuery)) {
-        const WhereCallbackReal whereCallback(txn, collection->ns().db());
-        auto statusWithCQ =
-            CanonicalQuery::canonicalize(collection->ns(), unparsedQuery, whereCallback);
-        if (!statusWithCQ.isOK()) {
-            return statusWithCQ.getStatus();
-        }
-        unique_ptr<CanonicalQuery> cq = std::move(statusWithCQ.getValue());
-
-        // Takes ownership of 'cq'.
-        return getExecutor(txn, collection, std::move(cq), yieldPolicy, plannerOptions);
-    }
-
-    LOG(2) << "Using idhack: " << unparsedQuery.toString();
-
-    unique_ptr<WorkingSet> ws = make_unique<WorkingSet>();
-    unique_ptr<PlanStage> root = make_unique<IDHackStage>(
-        txn, collection, unparsedQuery["_id"].wrap(), ws.get(), descriptor);
-
-    // Might have to filter out orphaned docs.
-    if (plannerOptions & QueryPlannerParams::INCLUDE_SHARD_FILTER) {
-        root = make_unique<ShardFilterStage>(
-            txn,
-            ShardingState::get(txn)->getCollectionMetadata(collection->ns().ns()),
-            ws.get(),
-            root.release());
-    }
-
-    return PlanExecutor::make(txn, std::move(ws), std::move(root), collection, yieldPolicy);
-}
-
 //
 // Find
 //
@@ -644,6 +609,12 @@ StatusWith<unique_ptr<PlanStage>> applyProjection(OperationContext* txn,
                 "cannot use a positional projection and return the new document"};
     }
 
+    // $meta sortKey is not allowed to be projected in findAndModify commands.
+    if (pp->wantSortKey()) {
+        return {ErrorCodes::BadValue,
+                "Cannot use a $meta sortKey projection in findAndModify commands."};
+    }
+
     ProjectionStageParams params(WhereCallbackReal(txn, nsString.db()));
     params.projObj = proj;
     params.fullExpression = cq->root();
diff --git a/src/mongo/db/query/get_executor.h b/src/mongo/db/query/get_executor.h
index cd90a56d487..49b83166ac5 100644
--- a/src/mongo/db/query/get_executor.h
+++ b/src/mongo/db/query/get_executor.h
@@ -78,25 +78,6 @@ StatusWith<std::unique_ptr<PlanExecutor>> getExecutor(
     size_t plannerOptions = 0);
 
 /**
- * Get a plan executor for query. This differs from the getExecutor(...) function
- * above in that the above requires a non-NULL canonical query, whereas this
- * function can retrieve a plan executor from the raw query object.
- *
- * Used to support idhack updates that do not create a canonical query.
- *
- * If the query is valid and an executor could be created, returns a StatusWith with the
- * PlanExecutor.
- *
- * If the query cannot be executed, returns a Status indicating why.
- */
-StatusWith<std::unique_ptr<PlanExecutor>> getExecutor(OperationContext* txn,
-                                                      Collection* collection,
-                                                      const std::string& ns,
-                                                      const BSONObj& unparsedQuery,
-                                                      PlanExecutor::YieldPolicy yieldPolicy,
-                                                      size_t plannerOptions = 0);
-
-/**
  * Get a plan executor for a .find() operation.
  *
  * If the query is valid and an executor could be created, returns a StatusWith with the