SERVER-10026 returnKey in new system

17 files changed, 108 insertions, 56 deletions

diff --git a/jstests/explainc.js b/jstests/explainc.js
index b9850e68911..b9cfef4ad79 100644
--- a/jstests/explainc.js
+++ b/jstests/explainc.js
@@ -64,7 +64,7 @@ assertUnhintedExplain( { n:1, nscanned:1, nscannedObjects:1, scanAndOrder:true }
                            .hint( { a:1, b:1 } ) );
 
 // In memory sort $returnKey query.
-assertUnhintedExplain( { n:1, nscanned:1, nscannedObjects:0, scanAndOrder:true },
+assertUnhintedExplain( { n:1, nscanned:1, scanAndOrder:true },
                        t.find( { a:{ $gt:0 } } )._addSpecial( "$returnKey", true ).sort( { b:1 } )
                            .hint( { a:1, b:1 } ) );
 
@@ -155,27 +155,14 @@ assertUnhintedExplain( { // cursor:'BtreeCursor b_1',
                         },
                        t.find( { b:1 }, { _id:0, b:1 } ).sort( { a:1 } ).skip( 1 ) );
 
-// Ordered plan chosen, $returnKey specified.
-assertUnhintedExplain( { cursor:'BtreeCursor a_1_b_1', n:30, nscanned:30, nscannedObjects:0,
-                         scanAndOrder:false },
-                       t.find( { b:{ $gte:0 } }, { _id:0, b:1 } ).sort( { a:1 } )
-                           ._addSpecial( "$returnKey", true ) );
-
-// Ordered plan chosen, $returnKey specified, matching requires loading document.
-assertUnhintedExplain( { cursor:'BtreeCursor a_1_b_1', n:30, nscanned:30, nscannedObjects:30,
-                         scanAndOrder:false },
-                       t.find( { b:{ $gte:0 }, c:null }, { _id:0, b:1 } ).sort( { a:1 } )
-                           ._addSpecial( "$returnKey", true ) );
-
 // Unordered plan chosen, $returnKey specified.
-assertUnhintedExplain( { cursor:'BtreeCursor b_1', n:1, nscanned:1, nscannedObjects:0,
-                         nscannedObjectsAllPlans:1, scanAndOrder:true },
+assertUnhintedExplain( { cursor:'BtreeCursor b_1', n:1, nscanned:1, scanAndOrder:true },
                        t.find( { b:1 }, { _id:0, b:1 } ).sort( { a:1 } )
                            ._addSpecial( "$returnKey", true ) );
 
 // Unordered plan chosen, $returnKey specified, matching requires loading document.
 assertUnhintedExplain( { cursor:'BtreeCursor b_1', n:1, nscanned:1, nscannedObjects:1,
-                         nscannedObjectsAllPlans:2, scanAndOrder:true },
+                         scanAndOrder:true },
                        t.find( { b:1, c:null }, { _id:0, b:1 } ).sort( { a:1 } )
                            ._addSpecial( "$returnKey", true ) );
 
diff --git a/jstests/index_diag.js b/jstests/index_diag.js
index 22390c14901..21840682e7f 100644
--- a/jstests/index_diag.js
+++ b/jstests/index_diag.js
@@ -46,15 +46,5 @@ assert.eq( r( all ) ,
           ._addSpecial( "$returnKey" , true ).toArray()
           )
 
-// Descriptive test that a query involving multiple plans may return keys from multiple indexes.
-t.dropIndex( { _id : 1 , x : 1 } );
-ret =
-t.find( { _id : { $gt : -10 }, x : { $gt : -10 } } )._addSpecial( "$returnKey" , true ).toArray()
-keys = {};
-for( i in ret ) {
-    Object.extend( keys, ret[ i ] );
-}
-assert.eq( [ "_id" , "x" ], Object.keySet( keys ).sort() );
-
 assert.eq( [ {} , {} , {} ],
           t.find().hint( { $natural : 1 } )._addSpecial( "$returnKey" , true ).toArray() )
diff --git a/src/mongo/db/exec/and_common-inl.h b/src/mongo/db/exec/and_common-inl.h
index 932f74d8273..31f6fe33c87 100644
--- a/src/mongo/db/exec/and_common-inl.h
+++ b/src/mongo/db/exec/and_common-inl.h
@@ -52,12 +52,11 @@ namespace mongo {
             }
 
             // Merge computed data.
-            if (!dest->hasComputed(WSM_COMPUTED_TEXT_SCORE) && src->hasComputed(WSM_COMPUTED_TEXT_SCORE)) {
-                dest->addComputed(src->getComputed(WSM_COMPUTED_TEXT_SCORE)->clone());
-            }
-
-            if (!dest->hasComputed(WSM_COMPUTED_GEO_DISTANCE) && src->hasComputed(WSM_COMPUTED_GEO_DISTANCE)) {
-                dest->addComputed(src->getComputed(WSM_COMPUTED_GEO_DISTANCE)->clone());
+            typedef WorkingSetComputedDataType WSCD;
+            for (WSCD i = WSCD(0); i < WSM_COMPUTED_NUM_TYPES; i = WSCD(i + 1)) {
+                if (!dest->hasComputed(i) && src->hasComputed(i)) {
+                    dest->addComputed(src->getComputed(i)->clone());
+                }
             }
         }
     };
diff --git a/src/mongo/db/exec/index_scan.cpp b/src/mongo/db/exec/index_scan.cpp
index 333e1845503..55f756a5d62 100644
--- a/src/mongo/db/exec/index_scan.cpp
+++ b/src/mongo/db/exec/index_scan.cpp
@@ -29,6 +29,7 @@
 #include "mongo/db/exec/index_scan.h"
 
 #include "mongo/db/exec/filter.h"
+#include "mongo/db/exec/working_set_computed_data.h"
 #include "mongo/db/index/catalog_hack.h"
 #include "mongo/db/index/index_access_method.h"
 #include "mongo/db/index/index_cursor.h"
@@ -177,17 +178,23 @@ namespace mongo {
             }
         }
 
+        BSONObj ownedKeyObj = _indexCursor->getKey().getOwned();
+
         WorkingSetID id = _workingSet->allocate();
         WorkingSetMember* member = _workingSet->get(id);
         member->loc = loc;
-        member->keyData.push_back(IndexKeyDatum(_descriptor->keyPattern(),
-                                                _indexCursor->getKey().getOwned()));
+        member->keyData.push_back(IndexKeyDatum(_descriptor->keyPattern(), ownedKeyObj));
         member->state = WorkingSetMember::LOC_AND_IDX;
 
         if (Filter::passes(member, _filter)) {
             if (NULL != _filter) {
                 ++_specificStats.matchTested;
             }
+            if (_params.addKeyMetadata) {
+                BSONObjBuilder bob;
+                bob.appendKeys(_descriptor->keyPattern(), ownedKeyObj);
+                member->addComputed(new IndexKeyComputedData(bob.obj()));
+            }
             *out = id;
             ++_commonStats.advanced;
             return PlanStage::ADVANCED;
diff --git a/src/mongo/db/exec/index_scan.h b/src/mongo/db/exec/index_scan.h
index 8bbe7111f38..0848da1b46e 100644
--- a/src/mongo/db/exec/index_scan.h
+++ b/src/mongo/db/exec/index_scan.h
@@ -50,7 +50,8 @@ namespace mongo {
                             limit(0),
                             forceBtreeAccessMethod(false),
                             doNotDedup(false),
-                            maxScan(0) { }
+                            maxScan(0),
+                            addKeyMetadata(false) { }
 
         IndexDescriptor* descriptor;
 
@@ -68,6 +69,9 @@ namespace mongo {
 
         // How many keys will we look at?
         size_t maxScan;
+
+        // Do we want to add the key as metadata?
+        bool addKeyMetadata;
     };
 
     /**
diff --git a/src/mongo/db/exec/projection_exec.cpp b/src/mongo/db/exec/projection_exec.cpp
index 1ad47523970..269a1104fb0 100644
--- a/src/mongo/db/exec/projection_exec.cpp
+++ b/src/mongo/db/exec/projection_exec.cpp
@@ -44,7 +44,9 @@ namespace mongo {
           _arrayOpType(ARRAY_OP_NORMAL),
           _hasNonSimple(false),
           _hasDottedField(false),
-          _queryExpression(NULL) { }
+          _queryExpression(NULL),
+          _hasReturnKey(false) { }
+
 
     ProjectionExec::ProjectionExec(const BSONObj& spec, const MatchExpression* queryExpression)
         : _include(true),
@@ -56,7 +58,8 @@ namespace mongo {
           _arrayOpType(ARRAY_OP_NORMAL),
           _hasNonSimple(false),
           _hasDottedField(false),
-          _queryExpression(queryExpression) {
+          _queryExpression(queryExpression),
+          _hasReturnKey(false) {
 
         // Are we including or excluding fields?
         // -1 when we haven't initialized it.
@@ -124,6 +127,11 @@ namespace mongo {
                     else if (mongoutils::str::equals(e2.valuestr(), "diskloc")) {
                         _meta[e.fieldName()] = META_DISKLOC;
                     }
+                    else if (mongoutils::str::equals(e2.valuestr(), "indexKey")) {
+                        _hasReturnKey = true;
+                        // The index key clobbers everything so just stop parsing here.
+                        return;
+                    }
                     else {
                         // This shouldn't happen, should be caught by parsing.
                         verify(0);
@@ -218,8 +226,23 @@ namespace mongo {
     //
 
     Status ProjectionExec::transform(WorkingSetMember* member) const {
-        BSONObjBuilder bob;
+        if (_hasReturnKey) {
+            BSONObj keyObj;
+
+            if (member->hasComputed(WSM_INDEX_KEY)) {
+                const IndexKeyComputedData* key
+                    = static_cast<const IndexKeyComputedData*>(member->getComputed(WSM_INDEX_KEY));
+                keyObj = key->getKey();
+            }
 
+            member->state = WorkingSetMember::OWNED_OBJ;
+            member->obj = keyObj;
+            member->keyData.clear();
+            member->loc = DiskLoc();
+            return Status::OK();
+        }
+
+        BSONObjBuilder bob;
         if (!requiresDocument()) {
             // Go field by field.
             if (_includeID) {
diff --git a/src/mongo/db/exec/projection_exec.h b/src/mongo/db/exec/projection_exec.h
index 33c20582f96..8b6efb26003 100644
--- a/src/mongo/db/exec/projection_exec.h
+++ b/src/mongo/db/exec/projection_exec.h
@@ -182,6 +182,10 @@ namespace mongo {
 
         // Projections that aren't sourced from the document or index keys.
         MetaMap _meta;
+
+        // Do we have a returnKey projection?  If so we *only* output the index key metadata.  If
+        // it's not found we output nothing.
+        bool _hasReturnKey;
     };
 
 }  // namespace mongo
diff --git a/src/mongo/db/exec/working_set.h b/src/mongo/db/exec/working_set.h
index b3e17fea9e4..c32cf8be121 100644
--- a/src/mongo/db/exec/working_set.h
+++ b/src/mongo/db/exec/working_set.h
@@ -147,6 +147,7 @@ namespace mongo {
     enum WorkingSetComputedDataType {
         WSM_COMPUTED_TEXT_SCORE = 0,
         WSM_COMPUTED_GEO_DISTANCE = 1,
+        WSM_INDEX_KEY = 2,
         WSM_COMPUTED_NUM_TYPES,
     };
 
diff --git a/src/mongo/db/exec/working_set_common.cpp b/src/mongo/db/exec/working_set_common.cpp
index c50ab745ab8..1406f809c58 100644
--- a/src/mongo/db/exec/working_set_common.cpp
+++ b/src/mongo/db/exec/working_set_common.cpp
@@ -55,12 +55,12 @@ namespace mongo {
         dest->keyData = src.keyData;
         dest->state = src.state;
 
-        if (src.hasComputed(WSM_COMPUTED_TEXT_SCORE)) {
-            dest->addComputed(src.getComputed(WSM_COMPUTED_TEXT_SCORE)->clone());
-        }
-
-        if (src.hasComputed(WSM_COMPUTED_GEO_DISTANCE)) {
-            dest->addComputed(src.getComputed(WSM_COMPUTED_GEO_DISTANCE)->clone());
+        // Merge computed data.
+        typedef WorkingSetComputedDataType WSCD;
+        for (WSCD i = WSCD(0); i < WSM_COMPUTED_NUM_TYPES; i = WSCD(i + 1)) {
+            if (src.hasComputed(i)) {
+                dest->addComputed(src.getComputed(i)->clone());
+            }
         }
     }
 
diff --git a/src/mongo/db/exec/working_set_computed_data.h b/src/mongo/db/exec/working_set_computed_data.h
index 790115e4f12..887167167f5 100644
--- a/src/mongo/db/exec/working_set_computed_data.h
+++ b/src/mongo/db/exec/working_set_computed_data.h
@@ -64,4 +64,20 @@ namespace mongo {
         double _score;
     };
 
+    class IndexKeyComputedData : public WorkingSetComputedData {
+    public:
+        IndexKeyComputedData(BSONObj key)
+            : WorkingSetComputedData(WSM_INDEX_KEY),
+              _key(key) { }
+
+        BSONObj getKey() const { return _key; }
+
+        virtual IndexKeyComputedData* clone() const {
+            return new IndexKeyComputedData(_key);
+        }
+
+    private:
+        BSONObj _key;
+    };
+
 }  // namespace mongo
diff --git a/src/mongo/db/query/lite_parsed_query.cpp b/src/mongo/db/query/lite_parsed_query.cpp
index 4f1c979c3b5..639c2f20612 100644
--- a/src/mongo/db/query/lite_parsed_query.cpp
+++ b/src/mongo/db/query/lite_parsed_query.cpp
@@ -309,7 +309,15 @@ namespace mongo {
                 }
                 else if (str::equals("returnKey", name)) {
                     // Won't throw.
-                    _returnKey = e.trueValue();
+                    if (e.trueValue()) {
+                        _returnKey = true;
+                        BSONObjBuilder projBob;
+                        projBob.appendElements(_proj);
+                        // XXX: what's the syntax here?
+                        BSONObj indexKey = BSON("$$" << BSON("$meta" << "indexKey"));
+                        projBob.append(indexKey.firstElement());
+                        _proj = projBob.obj();
+                    }
                 }
                 else if (str::equals("maxScan", name)) {
                     // Won't throw.
diff --git a/src/mongo/db/query/new_find.cpp b/src/mongo/db/query/new_find.cpp
index 5fccdf4de77..3077b3f594c 100644
--- a/src/mongo/db/query/new_find.cpp
+++ b/src/mongo/db/query/new_find.cpp
@@ -100,6 +100,7 @@ namespace mongo {
     void enableNewQueryFramework() { newQueryFrameworkEnabled = true; }
 
     // Do we use the old or the new?  I call this the spigot.
+    // XXX: remove this
     bool canUseNewSystem(const QueryMessage& qm, CanonicalQuery** cqOut) {
         // This is a read lock.  We require this because if we're parsing a $where, the
         // where-specific parsing code assumes we have a lock and creates execution machinery that
@@ -114,15 +115,6 @@ namespace mongo {
         }
         verify(cq);
         auto_ptr<CanonicalQuery> scopedCq(cq);
-
-        const LiteParsedQuery& pq = cq->getParsed();
-
-        // We fail to deal well with obscure arguments to .find().
-        if (pq.returnKey()) {
-            QLOG() << "rejecting wacky query args query\n";
-            return false;
-        }
-
         *cqOut = scopedCq.release();
         return true;
     }
diff --git a/src/mongo/db/query/parsed_projection.cpp b/src/mongo/db/query/parsed_projection.cpp
index ef2bf600e5a..a61b5821b2d 100644
--- a/src/mongo/db/query/parsed_projection.cpp
+++ b/src/mongo/db/query/parsed_projection.cpp
@@ -52,6 +52,8 @@ namespace mongo {
 
         bool includeID = true;
 
+        bool hasIndexKeyProjection = false;
+
         // Until we see a positional or elemMatch operator we're normal.
         ArrayOpType arrayOpType = ARRAY_OP_NORMAL;
 
@@ -136,10 +138,15 @@ namespace mongo {
                     }
 
                     if (!mongoutils::str::equals(e2.valuestr(), "text")
-                        && !mongoutils::str::equals(e2.valuestr(), "diskloc")) {
+                        && !mongoutils::str::equals(e2.valuestr(), "diskloc")
+                        && !mongoutils::str::equals(e2.valuestr(), "indexKey")) {
                         return Status(ErrorCodes::BadValue,
                                       "unsupported $meta operator: " + e2.str());
                     }
+
+                    if (mongoutils::str::equals(e2.valuestr(), "indexKey")) {
+                        hasIndexKeyProjection = true;
+                    }
                 }
                 else {
                     return Status(ErrorCodes::BadValue,
@@ -197,6 +204,13 @@ namespace mongo {
         verify(spec.isOwned());
         pp->_source = spec;
 
+        // returnKey clobbers everything.
+        if (hasIndexKeyProjection) {
+            pp->_requiresDocument = false;
+            *out = pp.release();
+            return Status::OK();
+        }
+
         // Dotted fields aren't covered, non-simple require match details, and as for include, "if
         // we default to including then we can't use an index because we don't know what we're
         // missing."
diff --git a/src/mongo/db/query/planner_access.cpp b/src/mongo/db/query/planner_access.cpp
index 6f28702931f..744f17be890 100644
--- a/src/mongo/db/query/planner_access.cpp
+++ b/src/mongo/db/query/planner_access.cpp
@@ -136,6 +136,7 @@ namespace mongo {
             isn->indexIsMultiKey = index.multikey;
             isn->bounds.fields.resize(index.keyPattern.nFields());
             isn->maxScan = query.getParsed().getMaxScan();
+            isn->addKeyMetadata = query.getParsed().returnKey();
 
             IndexBoundsBuilder::translate(expr, index.keyPattern.firstElement(),
                                           &isn->bounds.fields[0], tightnessOut);
@@ -841,6 +842,7 @@ namespace mongo {
         isn->indexIsMultiKey = index.multikey;
         isn->bounds.fields.resize(index.keyPattern.nFields());
         isn->maxScan = query.getParsed().getMaxScan();
+        isn->addKeyMetadata = query.getParsed().returnKey();
 
         // TODO: can we use simple bounds with this compound idx?
         BSONObjIterator it(isn->indexKeyPattern);
@@ -924,6 +926,7 @@ namespace mongo {
         isn->indexIsMultiKey = index.multikey;
         isn->direction = 1;
         isn->maxScan = query.getParsed().getMaxScan();
+        isn->addKeyMetadata = query.getParsed().returnKey();
         isn->bounds.isSimpleRange = true;
         isn->bounds.startKey = startKey;
         isn->bounds.endKey = endKey;
diff --git a/src/mongo/db/query/query_solution.cpp b/src/mongo/db/query/query_solution.cpp
index af3f08ca4ef..3c2b28a0864 100644
--- a/src/mongo/db/query/query_solution.cpp
+++ b/src/mongo/db/query/query_solution.cpp
@@ -307,7 +307,7 @@ namespace mongo {
     //
 
     IndexScanNode::IndexScanNode()
-        : indexIsMultiKey(false), limit(0), direction(1), maxScan(0) { }
+        : indexIsMultiKey(false), limit(0), direction(1), maxScan(0), addKeyMetadata(false) { }
 
     void IndexScanNode::appendToString(stringstream* ss, int indent) const {
         addIndent(ss, indent);
diff --git a/src/mongo/db/query/query_solution.h b/src/mongo/db/query/query_solution.h
index a3fdceabbd8..8eee9ad74ab 100644
--- a/src/mongo/db/query/query_solution.h
+++ b/src/mongo/db/query/query_solution.h
@@ -358,6 +358,9 @@ namespace mongo {
         // maxScan option to .find() limits how many docs we look at.
         int maxScan;
 
+        // If there's a 'returnKey' projection we add key metadata.
+        bool addKeyMetadata;
+
         // BIG NOTE:
         // If you use simple bounds, we'll use whatever index access method the keypattern implies.
         // If you use the complex bounds, we force Btree access.
diff --git a/src/mongo/db/query/stage_builder.cpp b/src/mongo/db/query/stage_builder.cpp
index 68da614945b..63d29939626 100644
--- a/src/mongo/db/query/stage_builder.cpp
+++ b/src/mongo/db/query/stage_builder.cpp
@@ -89,6 +89,7 @@ namespace mongo {
             params.direction = ixn->direction;
             params.limit = ixn->limit;
             params.maxScan = ixn->maxScan;
+            params.addKeyMetadata = ixn->addKeyMetadata;
             return new IndexScan(params, ws, ixn->filter.get());
         }
         else if (STAGE_FETCH == root->getType()) {