SERVER-5022 allow covered queries with shard filter

16 files changed, 313 insertions, 40 deletions

diff --git a/src/mongo/db/commands/mr.cpp b/src/mongo/db/commands/mr.cpp
index 950df2d6d02..49857488785 100644
--- a/src/mongo/db/commands/mr.cpp
+++ b/src/mongo/db/commands/mr.cpp
@@ -1369,7 +1369,7 @@ namespace mongo {
                             // because of a chunk migration
                             if ( collMetadata ) {
                                 KeyPattern kp( collMetadata->getKeyPattern() );
-                                if ( !collMetadata->keyBelongsToMe( kp.extractSingleKey( o ) ) ) {
+                                if (!collMetadata->keyBelongsToMe(kp.extractShardKeyFromDoc(o))) {
                                     continue;
                                 }
                             }
diff --git a/src/mongo/db/dbhelpers.cpp b/src/mongo/db/dbhelpers.cpp
index db184d1e7fd..2ebffc2c553 100644
--- a/src/mongo/db/dbhelpers.cpp
+++ b/src/mongo/db/dbhelpers.cpp
@@ -419,7 +419,7 @@ namespace mongo {
                     bool docIsOrphan;
                     if ( metadataNow ) {
                         KeyPattern kp( metadataNow->getKeyPattern() );
-                        BSONObj key = kp.extractSingleKey( obj );
+                        BSONObj key = kp.extractShardKeyFromDoc(obj);
                         docIsOrphan = !metadataNow->keyBelongsToMe( key )
                             && !metadataNow->keyIsPending( key );
                     }
diff --git a/src/mongo/db/exec/shard_filter.cpp b/src/mongo/db/exec/shard_filter.cpp
index c03fb6b3e95..8d889098b2d 100644
--- a/src/mongo/db/exec/shard_filter.cpp
+++ b/src/mongo/db/exec/shard_filter.cpp
@@ -29,6 +29,10 @@
 #include "mongo/db/exec/shard_filter.h"
 
 #include "mongo/db/keypattern.h"
+#include "mongo/db/matcher/expression.h"
+#include "mongo/db/exec/filter.h"
+#include "mongo/db/exec/working_set_common.h"
+#include "mongo/util/log.h"
 
 namespace mongo {
 
@@ -56,15 +60,44 @@ namespace mongo {
         StageState status = _child->work(out);
 
         if (PlanStage::ADVANCED == status) {
-            // If we're sharded make sure that we don't return any data that hasn't been migrated
-            // off of our shared yet.
+
+            // If we're sharded make sure that we don't return data that is not owned by us,
+            // including pending documents from in-progress migrations and orphaned documents from
+            // aborted migrations
             if (_metadata) {
-                KeyPattern kp(_metadata->getKeyPattern());
 
+                KeyPattern shardKeyPattern(_metadata->getKeyPattern());
                 WorkingSetMember* member = _ws->get(*out);
+                WorkingSetMatchableDocument matchable(member);
+                BSONObj shardKey = shardKeyPattern.extractShardKeyFromMatchable(matchable);
+
+                if (shardKey.isEmpty()) {
+
+                    // We can't find a shard key for this document - this should never happen with
+                    // a non-fetched result unless our query planning is screwed up
+                    if (!member->hasObj()) {
+
+                        Status status(ErrorCodes::InternalError,
+                                     "shard key not found after a covered stage, "
+                                     "query planning has failed");
+
+                        // Fail loudly and cleanly in production, fatally in debug
+                        error() << status.toString();
+                        dassert(false);
+
+                        _ws->free(*out);
+                        *out = WorkingSetCommon::allocateStatusMember(_ws, status);
+                        return PlanStage::FAILURE;
+                    }
+
+                    // Skip this document with a warning - no shard key should not be possible
+                    // unless manually inserting data into a shard
+                    warning() << "no shard key found in document " << member->obj.toString() << " "
+                              << "for shard key pattern " << _metadata->getKeyPattern() << ", "
+                              << "document may have been inserted manually into shard";
+                }
 
-                // This performs excessive BSONObj creation but that's OK for now.
-                if (!_metadata->keyBelongsToMe(kp.extractSingleKey(member->obj))) {
+                if (!_metadata->keyBelongsToMe(shardKey)) {
                     _ws->free(*out);
                     ++_specificStats.chunkSkips;
                     return PlanStage::NEED_TIME;
diff --git a/src/mongo/db/keypattern.cpp b/src/mongo/db/keypattern.cpp
index 222253e4c4d..18a0d270a7f 100644
--- a/src/mongo/db/keypattern.cpp
+++ b/src/mongo/db/keypattern.cpp
@@ -62,20 +62,62 @@ namespace mongo {
                && i.next().eoo();
     }
 
-    BSONObj KeyPattern::extractSingleKey(const BSONObj& doc ) const {
+    BSONObj KeyPattern::extractShardKeyFromQuery(const BSONObj& query) const {
+        if (_pattern.isEmpty())
+            return BSONObj();
+
+        if (mongoutils::str::equals(_pattern.firstElement().valuestrsafe(), "hashed")) {
+            BSONElement fieldVal = query.getFieldDotted(_pattern.firstElementFieldName());
+            return BSON(_pattern.firstElementFieldName() <<
+                        BSONElementHasher::hash64(fieldVal , BSONElementHasher::DEFAULT_HASH_SEED));
+        }
+
+        return query.extractFields(_pattern);
+    }
+
+    BSONObj KeyPattern::extractShardKeyFromDoc(const BSONObj& doc) const {
+        BSONMatchableDocument matchable(doc);
+        return extractShardKeyFromMatchable(matchable);
+    }
+
+    BSONObj KeyPattern::extractShardKeyFromMatchable(const MatchableDocument& matchable) const {
+
         if ( _pattern.isEmpty() )
             return BSONObj();
 
-        if ( mongoutils::str::equals( _pattern.firstElement().valuestrsafe() , "hashed" ) ){
-            BSONElement fieldVal = doc.getFieldDotted( _pattern.firstElementFieldName() );
-            return BSON( _pattern.firstElementFieldName() <<
-                         BSONElementHasher::hash64( fieldVal ,
-                                                    BSONElementHasher::DEFAULT_HASH_SEED ) );
+        BSONObjBuilder keyBuilder;
+
+        BSONObjIterator patternIt(_pattern);
+        while (patternIt.more()) {
+
+            BSONElement patternEl = patternIt.next();
+            ElementPath path;
+            path.init(patternEl.fieldName());
+
+            MatchableDocument::IteratorHolder matchIt(&matchable, &path);
+            if (!matchIt->more())
+                return BSONObj();
+            BSONElement matchEl = matchIt->next().element();
+            // We sometimes get eoo(), apparently
+            if (matchEl.eoo() || matchIt->more())
+                return BSONObj();
+
+            if (mongoutils::str::equals(patternEl.valuestrsafe(), "hashed")) {
+                keyBuilder.append(patternEl.fieldName(),
+                                  BSONElementHasher::hash64(matchEl,
+                                                            BSONElementHasher::DEFAULT_HASH_SEED));
+            }
+            else {
+                // NOTE: The matched element may *not* have the same field name as the path -
+                // index keys don't contain field names, for example
+                keyBuilder.appendAs(matchEl, patternEl.fieldName());
+            }
         }
 
-        return doc.extractFields( _pattern );
+        return keyBuilder.obj();
     }
 
+
     bool KeyPattern::isSpecial() const {
         BSONForEach(e, _pattern) {
             int fieldVal = e.numberInt();
diff --git a/src/mongo/db/keypattern.h b/src/mongo/db/keypattern.h
index 40454e3b672..c4c8e849fd1 100644
--- a/src/mongo/db/keypattern.h
+++ b/src/mongo/db/keypattern.h
@@ -35,6 +35,7 @@
 #include "mongo/platform/unordered_set.h"
 #include "mongo/util/mongoutils/str.h"
 #include "mongo/db/query/index_bounds.h"
+#include "mongo/db/matcher/matchable.h"
 
 namespace mongo {
 
@@ -134,7 +135,8 @@ namespace mongo {
 
         std::string toString() const{ return toBSON().toString(); }
 
-        /* Given a document, extracts the index key corresponding to this KeyPattern
+        /**
+         * Given a document, extracts the shard key corresponding to the key pattern.
          * Warning: assumes that there is a *single* key to be extracted!
          *
          * Examples:
@@ -143,11 +145,33 @@ namespace mongo {
          *   { c : 4 , a : 2 } -->  returns { a : 2 }
          *   { b : 2 }  (bad input, don't call with this)
          *   { a : [1,2] }  (bad input, don't call with this)
-         *
          *  If 'this' KeyPattern is { a  : "hashed" }
          *   { a: 1 } --> returns { a : NumberLong("5902408780260971510")  }
+         *  If 'this' KeyPattern is { 'a.b' : 1 }
+         *   { a : { b : "hi" } } --> returns { a : "hi" }
+         */
+        BSONObj extractShardKeyFromDoc(const BSONObj& doc) const;
+
+        /**
+         * Given a MatchableDocument, extracts the shard key corresponding to the key pattern.
+         * See above.
+         */
+        BSONObj extractShardKeyFromMatchable(const MatchableDocument& matchable) const;
+
+        /**
+         * Given a query expression, extracts the shard key corresponding to the key pattern.
+         *
+         * NOTE: This generally is similar to the above, however "a.b" fields in the query (which
+         * are invalid document fields) may match "a.b" fields in the shard key pattern.
+         *
+         * Examples:
+         *  If the key pattern is { a : 1 }
+         *   { a : "hi", b : 4 } --> returns { a : "hi" }
+         *  If the key pattern is { 'a.b' : 1 }
+         *   { a : { b : "hi" } } --> returns { 'a.b' : "hi" }
+         *   { 'a.b' : "hi" } --> returns { 'a.b' : "hi" }
          */
-        BSONObj extractSingleKey( const BSONObj& doc ) const;
+        BSONObj extractShardKeyFromQuery(const BSONObj& query) const;
 
         /**
          * Return an ordered list of bounds generated using this KeyPattern and the
diff --git a/src/mongo/db/query/planner_analysis.cpp b/src/mongo/db/query/planner_analysis.cpp
index 68767dd1cb9..233e3b3e4b8 100644
--- a/src/mongo/db/query/planner_analysis.cpp
+++ b/src/mongo/db/query/planner_analysis.cpp
@@ -501,12 +501,29 @@ namespace mongo {
         // If we're answering a query on a sharded system, we need to drop documents that aren't
         // logically part of our shard.
         if (params.options & QueryPlannerParams::INCLUDE_SHARD_FILTER) {
-            // TODO: We could use params.shardKey to do fetch analysis instead of always fetching.
+
             if (!solnRoot->fetched()) {
-                FetchNode* fetch = new FetchNode();
-                fetch->children.push_back(solnRoot);
-                solnRoot = fetch;
+
+                // See if we need to fetch information for our shard key.
+                // NOTE: Solution nodes only list ordinary, non-transformed index keys for now
+
+                bool fetch = false;
+                BSONObjIterator it(params.shardKey);
+                while (it.more()) {
+                    BSONElement nextEl = it.next();
+                    if (!solnRoot->hasField(nextEl.fieldName())) {
+                        fetch = true;
+                        break;
+                    }
+                }
+
+                if (fetch) {
+                    FetchNode* fetch = new FetchNode();
+                    fetch->children.push_back(solnRoot);
+                    solnRoot = fetch;
+                }
             }
+
             ShardingFilterNode* sfn = new ShardingFilterNode();
             sfn->children.push_back(solnRoot);
             solnRoot = sfn;
diff --git a/src/mongo/db/query/query_planner_test.cpp b/src/mongo/db/query/query_planner_test.cpp
index 20d2a01c2ce..4e4c92f930a 100644
--- a/src/mongo/db/query/query_planner_test.cpp
+++ b/src/mongo/db/query/query_planner_test.cpp
@@ -4588,6 +4588,132 @@ namespace {
     }
 
     //
+    // Test shard filter query planning
+    //
+
+    TEST_F(QueryPlannerTest, ShardFilterCollScan) {
+        params.options = QueryPlannerParams::INCLUDE_SHARD_FILTER;
+        params.shardKey = BSON("a" << 1);
+        addIndex(BSON("a" << 1));
+
+        runQuery(fromjson("{b: 1}"));
+
+        assertNumSolutions(1U);
+        assertSolutionExists("{sharding_filter: {node: "
+                               "{cscan: {dir: 1}}}}}}}");
+    }
+
+    TEST_F(QueryPlannerTest, ShardFilterBasicIndex) {
+        params.options = QueryPlannerParams::INCLUDE_SHARD_FILTER;
+        params.shardKey = BSON("a" << 1);
+        addIndex(BSON("a" << 1));
+        addIndex(BSON("b" << 1));
+
+        runQuery(fromjson("{b: 1}"));
+
+        assertNumSolutions(1U);
+        assertSolutionExists("{sharding_filter: {node: "
+                               "{fetch: {node: "
+                                 "{ixscan: {pattern: {b: 1}}}}}}}");
+    }
+
+    TEST_F(QueryPlannerTest, ShardFilterBasicCovered) {
+        params.options = QueryPlannerParams::INCLUDE_SHARD_FILTER;
+        params.shardKey = BSON("a" << 1);
+        addIndex(BSON("a" << 1));
+
+        runQuery(fromjson("{a: 1}"));
+
+        assertNumSolutions(1U);
+        assertSolutionExists("{fetch: {node: "
+                               "{sharding_filter: {node: "
+                                 "{ixscan: {pattern: {a: 1}}}}}}}");
+    }
+
+    TEST_F(QueryPlannerTest, ShardFilterBasicProjCovered) {
+        params.options = QueryPlannerParams::INCLUDE_SHARD_FILTER;
+        params.shardKey = BSON("a" << 1);
+        addIndex(BSON("a" << 1));
+
+        runQuerySortProj(fromjson("{a: 1}"), BSONObj(), fromjson("{_id : 0, a : 1}"));
+
+        assertNumSolutions(1U);
+        assertSolutionExists("{proj: {spec: {_id: 0, a: 1}, type: 'coveredIndex', node: "
+                               "{sharding_filter: {node: "
+                                 "{ixscan: {pattern: {a: 1}}}}}}}");
+    }
+
+    TEST_F(QueryPlannerTest, ShardFilterCompoundProjCovered) {
+        params.options = QueryPlannerParams::INCLUDE_SHARD_FILTER;
+        params.shardKey = BSON("a" << 1 << "b" << 1);
+        addIndex(BSON("a" << 1 << "b" << 1));
+
+        runQuerySortProj(fromjson("{a: 1}"), BSONObj(), fromjson("{_id: 0, a: 1, b: 1}"));
+
+        assertNumSolutions(1U);
+        assertSolutionExists("{proj: {spec: {_id: 0, a: 1, b: 1 }, type: 'coveredIndex', node: "
+                               "{sharding_filter: {node: "
+                                 "{ixscan: {pattern: {a: 1, b: 1}}}}}}}");
+    }
+
+    TEST_F(QueryPlannerTest, ShardFilterNestedProjNotCovered) {
+        // Nested projections can't be covered currently, though the shard key filter shouldn't need
+        // to fetch.
+        params.options = QueryPlannerParams::INCLUDE_SHARD_FILTER;
+        params.shardKey = BSON("a" << 1 << "b.c" << 1);
+        addIndex(BSON("a" << 1 << "b.c" << 1));
+
+        runQuerySortProj(fromjson("{a: 1}"), BSONObj(), fromjson("{_id: 0, a: 1, 'b.c': 1}"));
+
+        assertNumSolutions(1U);
+        assertSolutionExists("{proj: {spec: {_id: 0, a: 1, 'b.c': 1 }, type: 'default', node: "
+                               "{fetch: {node: "
+                                 "{sharding_filter: {node: "
+                                   "{ixscan: {pattern: {a: 1, 'b.c': 1}}}}}}}}}");
+    }
+
+    TEST_F(QueryPlannerTest, ShardFilterHashProjNotCovered) {
+        params.options = QueryPlannerParams::INCLUDE_SHARD_FILTER;
+        params.shardKey = BSON("a" << "hashed");
+        addIndex(BSON("a" << "hashed"));
+
+        runQuerySortProj(fromjson("{a: 1}"), BSONObj(), fromjson("{_id : 0, a : 1}"));
+
+        assertNumSolutions(1U);
+        assertSolutionExists("{proj: {spec: {_id: 0,a: 1}, type: 'simple', node: "
+                               "{sharding_filter : {node: "
+                                 "{fetch: {node: "
+                                   "{ixscan: {pattern: {a: 'hashed'}}}}}}}}}");
+    }
+
+    TEST_F(QueryPlannerTest, ShardFilterKeyPrefixIndexCovered) {
+        params.options = QueryPlannerParams::INCLUDE_SHARD_FILTER;
+        params.shardKey = BSON("a" << 1);
+        addIndex(BSON("a" << 1 << "b" << 1 << "_id" << 1));
+
+        runQuerySortProj(fromjson("{a: 1}"), BSONObj(), fromjson("{a : 1}"));
+
+        assertNumSolutions(1U);
+        assertSolutionExists("{proj: {spec: {a: 1}, type: 'coveredIndex', node: "
+                               "{sharding_filter : {node: "
+                                 "{ixscan: {pattern: {a: 1, b: 1, _id: 1}}}}}}}");
+    }
+
+    TEST_F(QueryPlannerTest, ShardFilterNoIndexNotCovered) {
+        params.options = QueryPlannerParams::INCLUDE_SHARD_FILTER;
+        params.shardKey = BSON("a" << "hashed");
+        addIndex(BSON("b" << 1));
+
+        runQuerySortProj(fromjson("{b: 1}"), BSONObj(), fromjson("{_id : 0, a : 1}"));
+
+        assertNumSolutions(1U);
+        assertSolutionExists("{proj: {spec: {_id: 0,a: 1}, type: 'simple', node: "
+                               "{sharding_filter : {node: "
+                                 "{fetch: {node: "
+                                   "{ixscan: {pattern: {b: 1}}}}}}}}}");
+    }
+
+    //
     // Test bad input to query planner helpers.
     //
 
diff --git a/src/mongo/db/query/query_planner_test_lib.cpp b/src/mongo/db/query/query_planner_test_lib.cpp
index 134392e34d9..8b150aaf19c 100644
--- a/src/mongo/db/query/query_planner_test_lib.cpp
+++ b/src/mongo/db/query/query_planner_test_lib.cpp
@@ -393,6 +393,17 @@ namespace mongo {
             if (el.eoo() || !el.isABSONObj()) { return false; }
             BSONObj projObj = el.Obj();
 
+            BSONElement projType = projObj["type"];
+            if (!projType.eoo()) {
+                string projTypeStr = projType.str();
+                if (!((pn->projType == ProjectionNode::DEFAULT && projTypeStr == "default") ||
+                      (pn->projType == ProjectionNode::SIMPLE_DOC && projTypeStr == "simple") ||
+                      (pn->projType == ProjectionNode::COVERED_ONE_INDEX &&
+                           projTypeStr == "coveredIndex"))) {
+                    return false;
+                }
+            }
+
             BSONElement spec = projObj["spec"];
             if (spec.eoo() || !spec.isABSONObj()) { return false; }
             BSONElement child = projObj["node"];
@@ -467,6 +478,18 @@ namespace mongo {
 
             return solutionMatches(child.Obj(), kn->children[0]);
         }
+        else if (STAGE_SHARDING_FILTER == trueSoln->getType()) {
+            const ShardingFilterNode* fn = static_cast<const ShardingFilterNode*>(trueSoln);
+
+            BSONElement el = testSoln["sharding_filter"];
+            if (el.eoo() || !el.isABSONObj()) { return false; }
+            BSONObj keepObj = el.Obj();
+
+            BSONElement child = keepObj["node"];
+            if (child.eoo() || !child.isABSONObj()) { return false; }
+
+            return solutionMatches(child.Obj(), fn->children[0]);
+        }
 
         return false;
     }
diff --git a/src/mongo/db/query/query_solution.cpp b/src/mongo/db/query/query_solution.cpp
index deee6cca95a..3e042cb2b7c 100644
--- a/src/mongo/db/query/query_solution.cpp
+++ b/src/mongo/db/query/query_solution.cpp
@@ -26,6 +26,7 @@
  *    it in the license file.
  */
 
+#include "mongo/db/index_names.h"
 #include "mongo/db/query/query_solution.h"
 #include "mongo/db/query/lite_parsed_query.h"
 #include "mongo/db/matcher/expression_geo.h"
@@ -412,6 +413,10 @@ namespace mongo {
         // in the key was extracted from an array in the original document.
         if (indexIsMultiKey) { return false; }
 
+        // Custom index access methods may return non-exact key data - this function is currently
+        // used for covering exact key data only.
+        if (IndexNames::BTREE != IndexNames::findPluginName(indexKeyPattern)) { return false; }
+
         BSONObjIterator it(indexKeyPattern);
         while (it.more()) {
             if (field == it.next().fieldName()) {
diff --git a/src/mongo/s/chunk.cpp b/src/mongo/s/chunk.cpp
index 2c8960a495a..46afd515eea 100644
--- a/src/mongo/s/chunk.cpp
+++ b/src/mongo/s/chunk.cpp
@@ -240,7 +240,7 @@ namespace mongo {
         conn.done();
         if ( end.isEmpty() )
             return BSONObj();
-        return _manager->getShardKey().extractKey( end );
+        return _manager->getShardKey().extractKeyFromQueryOrDoc( end );
     }
 
     void Chunk::pickMedianKey( BSONObj& medianKey ) const {
@@ -1228,7 +1228,7 @@ namespace mongo {
     }
 
     ChunkPtr ChunkManager::findChunkForDoc( const BSONObj& doc ) const {
-        BSONObj key = _key.extractKey( doc );
+        BSONObj key = _key.extractKeyFromQueryOrDoc( doc );
         return findIntersectingChunk( key );
     }
 
diff --git a/src/mongo/s/d_migrate.cpp b/src/mongo/s/d_migrate.cpp
index 489c651e00f..d3782dd0a53 100644
--- a/src/mongo/s/d_migrate.cpp
+++ b/src/mongo/s/d_migrate.cpp
@@ -216,7 +216,7 @@ namespace mongo {
                     const BSONObj& max ,
                     const BSONObj& shardKeyPattern ) {
         ShardKeyPattern shardKey( shardKeyPattern );
-        BSONObj k = shardKey.extractKey( obj );
+        BSONObj k = shardKey.extractKeyFromQueryOrDoc( obj );
         return k.woCompare( min ) >= 0 && k.woCompare( max ) < 0;
     }
 
diff --git a/src/mongo/s/d_split.cpp b/src/mongo/s/d_split.cpp
index c6d8a6272ae..c5f6fe331a3 100644
--- a/src/mongo/s/d_split.cpp
+++ b/src/mongo/s/d_split.cpp
@@ -696,7 +696,7 @@ namespace mongo {
                     return false;
                 }
 
-                if (!isShardDocSizeValid(collMetadata->getKeyPattern(), endKey, &errmsg)) {
+                if (!isShardKeySizeValid(endKey, &errmsg)) {
                     warning() << errmsg << endl;
                     return false;
                 }
diff --git a/src/mongo/s/shard_key_pattern.cpp b/src/mongo/s/shard_key_pattern.cpp
index e0289dc2fa2..d1dbf9cd1ab 100644
--- a/src/mongo/s/shard_key_pattern.cpp
+++ b/src/mongo/s/shard_key_pattern.cpp
@@ -44,18 +44,15 @@ namespace mongo {
         return shardKeyPattern.isFieldNamePrefixOf( uIndexKeyPattern );
     }
 
-    bool isShardDocSizeValid(const BSONObj& shardKey, const BSONObj& doc, string* errMsg) {
+    bool isShardKeySizeValid(const BSONObj& shardKey, string* errMsg) {
         string dummy;
         if (errMsg == NULL) {
             errMsg = &dummy;
         }
 
-        const KeyPattern shardKeyPattern(shardKey);
-        BSONObj key = shardKeyPattern.extractSingleKey(doc);
-
-        if (key.objsize() > kMaxShardKeySize) {
-            *errMsg = str::stream() << "shard key is larger than " << kMaxShardKeySize
-                                    << " bytes: " << key.objsize();
+        if (shardKey.objsize() > kMaxShardKeySize) {
+            *errMsg = str::stream() << "shard key is larger than " << kMaxShardKeySize << " bytes: "
+                                    << shardKey.objsize();
             return false;
         }
 
diff --git a/src/mongo/s/shard_key_pattern.h b/src/mongo/s/shard_key_pattern.h
index e7c98ffb027..3f2c034a5e9 100644
--- a/src/mongo/s/shard_key_pattern.h
+++ b/src/mongo/s/shard_key_pattern.h
@@ -43,5 +43,5 @@ namespace mongo {
 
     bool isUniqueIndexCompatible(const BSONObj& shardKeyPattern, const BSONObj& uIndexKeyPattern);
 
-    bool isShardDocSizeValid(const BSONObj& shardKey, const BSONObj& doc, std::string* errMsg);
+    bool isShardKeySizeValid(const BSONObj& shardKey, std::string* errMsg);
 }
diff --git a/src/mongo/s/shardkey.cpp b/src/mongo/s/shardkey.cpp
index 4ad9054ccda..30a92a66d75 100644
--- a/src/mongo/s/shardkey.cpp
+++ b/src/mongo/s/shardkey.cpp
@@ -217,8 +217,8 @@ namespace mongo {
             ShardKeyPattern k( fromjson("{a:1,'sub.b':-1,'sub.c':1}") );
 
             BSONObj x = fromjson("{a:1,'sub.b':2,'sub.c':3}");
-            verify( k.extractKey( fromjson("{a:1,sub:{b:2,c:3}}") ).binaryEqual(x) );
-            verify( k.extractKey( fromjson("{sub:{b:2,c:3},a:1}") ).binaryEqual(x) );
+            verify( k.extractKeyFromQueryOrDoc( fromjson("{a:1,sub:{b:2,c:3}}") ).binaryEqual(x) );
+            verify( k.extractKeyFromQueryOrDoc( fromjson("{sub:{b:2,c:3},a:1}") ).binaryEqual(x) );
         }
 
         void isSpecialTest() {
@@ -298,7 +298,7 @@ namespace mongo {
             BSONObj k1 = BSON( "key" << 5 );
 
             verify( min < max );
-            verify( min < k.extractKey( k1 ) );
+            verify( min < k.extractKeyFromQueryOrDoc( k1 ) );
             verify( max > min );
 
             hasshardkeytest();
@@ -308,7 +308,7 @@ namespace mongo {
             BSONObj a = k1;
             BSONObj b = BSON( "key" << 999 );
 
-            verify( k.extractKey( a ) <  k.extractKey( b ) );
+            verify( k.extractKeyFromQueryOrDoc( a ) <  k.extractKeyFromQueryOrDoc( b ) );
 
             isSpecialTest();
 
diff --git a/src/mongo/s/shardkey.h b/src/mongo/s/shardkey.h
index c6abc9d2c58..bfbaf532fa8 100644
--- a/src/mongo/s/shardkey.h
+++ b/src/mongo/s/shardkey.h
@@ -99,7 +99,12 @@ namespace mongo {
 
         std::string toString() const;
 
-        BSONObj extractKey(const BSONObj& from) const;
+        /**
+         * DEPRECATED function to return a shard key from either a document or a query expression.
+         * Always prefer the more specific keypattern.h extractKeyFromXXX functions instead.
+         * TODO: Eliminate completely.
+         */
+        BSONObj extractKeyFromQueryOrDoc(const BSONObj& from) const;
 
         bool partOfShardKey(const StringData& key ) const {
             return pattern.hasField(key);
@@ -155,8 +160,9 @@ namespace mongo {
         std::set<std::string> patternfields;
     };
 
-    inline BSONObj ShardKeyPattern::extractKey(const BSONObj& from) const {
-        BSONObj k = pattern.extractSingleKey( from );
+    // See note above - do not use in new code
+    inline BSONObj ShardKeyPattern::extractKeyFromQueryOrDoc(const BSONObj& from) const {
+        BSONObj k = pattern.extractShardKeyFromQuery( from );
         uassert(13334, "Shard Key must be less than 512 bytes", k.objsize() < kMaxShardKeySize);
         return k;
     }