SERVER-14302 SERVER-14304 Fix IDHackRunner's use of sharding filter

3 files changed, 80 insertions, 39 deletions

diff --git a/jstests/sharding/idhack_sharded.js b/jstests/sharding/idhack_sharded.js
new file mode 100644
index 00000000000..8c45f5d0f00
--- /dev/null
+++ b/jstests/sharding/idhack_sharded.js
@@ -0,0 +1,34 @@
+// Test that idhack queries with projections obey the sharding filter.  SERVER-14032, SERVER-14034.
+
+var st = new ShardingTest({shards: 2});
+var coll = st.s0.getCollection("test.foo");
+
+//
+// Pre-split collection: shard 0 takes {x: {$lt: 0}}, shard 1 takes {x: {$gte: 0}}.
+//
+assert.commandWorked(coll.getDB().adminCommand({enableSharding: coll.getDB().getName()}));
+coll.getDB().adminCommand({movePrimary: coll.getDB().getName(), to: "shard0000"});
+assert.commandWorked(coll.getDB().adminCommand({shardCollection: coll.getFullName(), key: {x: 1}}));
+assert.commandWorked(coll.getDB().adminCommand({split: coll.getFullName(), middle: {x: 0}}));
+assert.commandWorked(coll.getDB().adminCommand({moveChunk: coll.getFullName(), find: {x: 0},
+                                                to: "shard0001"}));
+
+//
+// Test that idhack queries with projections that remove the shard key return correct results.
+// SERVER-14032.
+//
+assert.writeOK(coll.insert({_id: 1, x: 1, y: 1}));
+assert.eq(1, coll.find().itcount());
+assert.eq(1, coll.find({_id: 1}, {x: 0}).itcount());
+assert.eq(1, coll.find({_id: 1}, {y: 1}).itcount());
+coll.remove({});
+
+//
+// Test that idhack queries with covered projections do not return orphan documents.  SERVER-14034.
+//
+assert.writeOK(st.shard0.getCollection(coll.getFullName()).insert({_id: 1, x: 1}));
+assert.writeOK(st.shard1.getCollection(coll.getFullName()).insert({_id: 1, x: 1}));
+assert.eq(2, coll.count());
+assert.eq(1, coll.find().itcount());
+assert.eq(1, coll.find({_id: 1}, {_id: 1}).itcount());
+coll.remove({});
diff --git a/src/mongo/db/query/idhack_runner.cpp b/src/mongo/db/query/idhack_runner.cpp
index ca9124afaae..668707593a1 100644
--- a/src/mongo/db/query/idhack_runner.cpp
+++ b/src/mongo/db/query/idhack_runner.cpp
@@ -95,35 +95,49 @@ namespace mongo {
         if (NULL == objOut) {
             // No object requested - nothing to do.
         }
-        else if (hasCoveredProjection()) {
-            // Covered query on _id field only.
-            // Set object to search key.
-            // Search key is retrieved from the canonical query at
-            // construction and always contains the _id field name.
-            // It is possible to construct the ID hack runner with just the collection
-            // and the key object (which could be {"": my_obj_id}) but _query would be null
-            // in that case and the query would never be seen as covered.
-            *objOut = _key.getOwned();
-        }
         else {
-            _nscannedObjects++;
+            // If we're sharded, get the config metadata for this collection.  This will be used
+            // later to see if we own the document to be returned.
+            //
+            // Query execution machinery should generally delegate to ShardFilterStage in order to
+            // accomplish this task.  It is only safe to rely on the state of the config metadata
+            // here because it is not possible for the config metadata to change during the lifetime
+            // of the IDHackRunner (since the IDHackRunner returns only a single document, the
+            // config metadata must be the same as it was when the query started).
+            CollectionMetadataPtr collectionMetadata;
+            if (shardingState.needCollectionMetadata(_collection->ns().ns())) {
+                collectionMetadata = shardingState.getCollectionMetadata(_collection->ns().ns());
+            }
 
-            if (!applyProjection(loc, objOut)) {
-                // No projection. Just return the object inside the diskloc.
-                *objOut = _collection->docFor(loc);
+            // If we're not sharded, consider a covered projection (we can't if we're sharded, since
+            // we require a fetch in order to apply the sharding filter).
+            if (!collectionMetadata && hasCoveredProjection()) {
+                // Covered query on _id field only.  Set object to search key.  Search key is
+                // retrieved from the canonical query at construction and always contains the _id
+                // field name.  It is possible to construct the ID hack runner with just the
+                // collection and the key object (which could be {"": my_obj_id}) but _query would
+                // be null in that case and the query would never be seen as covered.
+                *objOut = _key.getOwned();
             }
+            // Otherwise, fetch the document.
+            else {
+                *objOut = _collection->docFor(loc);
+                _nscannedObjects++;
 
-            // If we're sharded make sure the key belongs to us.  We need the object to do this.
-            if (shardingState.needCollectionMetadata(_collection->ns().ns())) {
-                CollectionMetadataPtr m = shardingState.getCollectionMetadata(_collection->ns().ns());
-                if (m) {
-                    KeyPattern kp(m->getKeyPattern());
-                    if (!m->keyBelongsToMe( kp.extractSingleKey(*objOut))) {
+                // If we're sharded, make sure the key belongs to us.
+                if (collectionMetadata) {
+                    KeyPattern kp(collectionMetadata->getKeyPattern());
+                    if (!collectionMetadata->keyBelongsToMe(kp.extractSingleKey(*objOut))) {
                         // We have something with a matching _id but it doesn't belong to me.
                         _done = true;
                         return Runner::RUNNER_EOF;
                     }
                 }
+
+                // Apply the projection if one was requested.
+                if (_query && _query->getProj()) {
+                    *objOut = applyProjection(*objOut);
+                }
             }
         }
 
@@ -136,13 +150,8 @@ namespace mongo {
         return Runner::RUNNER_ADVANCED;
     }
 
-     bool IDHackRunner::applyProjection(const DiskLoc& loc, BSONObj* objOut) const {
-        if (NULL == _query.get() || NULL == _query->getProj()) {
-            // This idhack query does not have a projection.
-            return false;
-        }
-
-        const BSONObj& docAtLoc = _collection->docFor(loc);
+    BSONObj IDHackRunner::applyProjection(const BSONObj& docObj) const {
+        invariant(_query && _query->getProj());
 
         // We have a non-covered projection (covered projections should be handled earlier,
         // in getNext(..). For simple inclusion projections we use a fast path similar to that
@@ -155,14 +164,16 @@ namespace mongo {
             BSONObjBuilder bob;
             const BSONObj& queryObj = _query->getParsed().getFilter();
             bob.append(queryObj["_id"]);
-            *objOut = bob.obj();
+            return bob.obj();
         }
         else if (_query->getProj()->requiresDocument() || _query->getProj()->wantIndexKey()) {
             // Not a simple projection, so fallback on the regular projection path.
+            BSONObj projectedObj;
             ProjectionExec projExec(projObj,
                                     _query->root(),
                                     WhereCallbackReal(_collection->ns().db()));
-            projExec.transform(docAtLoc, objOut);
+            projExec.transform(docObj, &projectedObj);
+            return projectedObj;
         }
         else {
             // This is a simple inclusion projection. Start by getting the set
@@ -172,13 +183,11 @@ namespace mongo {
 
             // Apply the simple inclusion projection.
             BSONObjBuilder bob;
-            ProjectionStage::transformSimpleInclusion(docAtLoc, includedFields, bob);
+            ProjectionStage::transformSimpleInclusion(docObj, includedFields, bob);
 
-            *objOut = bob.obj();
+            return bob.obj();
         }
-
-        return true;
-     }
+    }
 
     bool IDHackRunner::isEOF() {
         return _killed || _done;
diff --git a/src/mongo/db/query/idhack_runner.h b/src/mongo/db/query/idhack_runner.h
index dd2027e19bb..89a6df33f55 100644
--- a/src/mongo/db/query/idhack_runner.h
+++ b/src/mongo/db/query/idhack_runner.h
@@ -88,12 +88,10 @@ namespace mongo {
         bool hasCoveredProjection() const;
 
         /**
-         * If '_query' has a projection, then apply it, returning the result in 'objOut'.
-         * The diskloc 'loc' contains the BSONObj to transform.
-         *
-         * Otherwise do nothing and return false.
+         * Apply the projection from '_query' to the given object and return the result.
+         * '_query->getProj()' must be non-NULL.
          */
-         bool applyProjection(const DiskLoc& loc, BSONObj* objOut) const;
+        BSONObj applyProjection(const BSONObj& docObj) const;
 
         // Not owned here.
         const Collection* _collection;