SERVER-19355 allow simultaneous returnKey and sortKey meta-projections

This is necessary in order to support queries with returnKey and sort routed through mongos.

3 files changed, 48 insertions, 17 deletions

diff --git a/src/mongo/db/exec/projection_exec.cpp b/src/mongo/db/exec/projection_exec.cpp
index fc20f4e53bf..1d996d1c6eb 100644
--- a/src/mongo/db/exec/projection_exec.cpp
+++ b/src/mongo/db/exec/projection_exec.cpp
@@ -39,6 +39,29 @@ namespace mongo {
 using std::max;
 using std::string;
 
+namespace {
+
+/**
+ * Adds sort key metadata inside 'member' to 'builder' with field name 'fieldName'.
+ *
+ * Returns a non-OK status if sort key metadata is missing from 'member'.
+ */
+Status addSortKeyMetaProj(StringData fieldName,
+                          const WorkingSetMember& member,
+                          BSONObjBuilder* builder) {
+    if (!member.hasComputed(WSM_SORT_KEY)) {
+        return Status(ErrorCodes::InternalError,
+                      "sortKey meta-projection requested but no data available");
+    }
+
+    const SortKeyComputedData* sortKeyData =
+        static_cast<const SortKeyComputedData*>(member.getComputed(WSM_SORT_KEY));
+    builder->append(fieldName, sortKeyData->getSortKey());
+    return Status::OK();
+}
+
+}  // namespace
+
 ProjectionExec::ProjectionExec()
     : _include(true),
       _special(false),
@@ -51,7 +74,6 @@ ProjectionExec::ProjectionExec()
       _queryExpression(NULL),
       _hasReturnKey(false) {}
 
-
 ProjectionExec::ProjectionExec(const BSONObj& spec,
                                const MatchExpression* queryExpression,
                                const MatchExpressionParser::WhereCallback& whereCallback)
@@ -126,7 +148,8 @@ ProjectionExec::ProjectionExec(const BSONObj& spec,
                 if (e2.valuestr() == LiteParsedQuery::metaTextScore) {
                     _meta[e.fieldName()] = META_TEXT_SCORE;
                 } else if (e2.valuestr() == LiteParsedQuery::metaSortKey) {
-                    _meta[e.fieldName()] = META_SORT_KEY;
+                    _sortKeyMetaFields.push_back(e.fieldName());
+                    _meta[_sortKeyMetaFields.back()] = META_SORT_KEY;
                 } else if (e2.valuestr() == LiteParsedQuery::metaRecordId) {
                     _meta[e.fieldName()] = META_RECORDID;
                 } else if (e2.valuestr() == LiteParsedQuery::metaGeoNearPoint) {
@@ -135,8 +158,6 @@ ProjectionExec::ProjectionExec(const BSONObj& spec,
                     _meta[e.fieldName()] = META_GEONEAR_DIST;
                 } else if (e2.valuestr() == LiteParsedQuery::metaIndexKey) {
                     _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);
@@ -226,15 +247,24 @@ void ProjectionExec::add(const string& field, int skip, int limit) {
 
 Status ProjectionExec::transform(WorkingSetMember* member) const {
     if (_hasReturnKey) {
-        BSONObj keyObj;
+        BSONObjBuilder builder;
 
         if (member->hasComputed(WSM_INDEX_KEY)) {
             const IndexKeyComputedData* key =
                 static_cast<const IndexKeyComputedData*>(member->getComputed(WSM_INDEX_KEY));
-            keyObj = key->getKey();
+            builder.appendElements(key->getKey());
         }
 
-        member->obj = Snapshotted<BSONObj>(SnapshotId(), keyObj.getOwned());
+        // Must be possible to do both returnKey meta-projection and sortKey meta-projection so that
+        // mongos can support returnKey.
+        for (auto fieldName : _sortKeyMetaFields) {
+            auto sortKeyMetaStatus = addSortKeyMetaProj(fieldName, *member, &builder);
+            if (!sortKeyMetaStatus.isOK()) {
+                return sortKeyMetaStatus;
+            }
+        }
+
+        member->obj = Snapshotted<BSONObj>(SnapshotId(), builder.obj());
         member->keyData.clear();
         member->loc = RecordId();
         member->transitionToOwnedObj();
@@ -316,13 +346,10 @@ Status ProjectionExec::transform(WorkingSetMember* member) const {
                 bob.append(it->first, 0.0);
             }
         } else if (META_SORT_KEY == it->second) {
-            if (!member->hasComputed(WSM_SORT_KEY)) {
-                return Status(ErrorCodes::InternalError,
-                              "sortKey meta-projection requested but no data available");
+            auto sortKeyMetaStatus = addSortKeyMetaProj(it->first, *member, &bob);
+            if (!sortKeyMetaStatus.isOK()) {
+                return sortKeyMetaStatus;
             }
-            const SortKeyComputedData* sortKeyData =
-                static_cast<const SortKeyComputedData*>(member->getComputed(WSM_SORT_KEY));
-            bob.append(it->first, sortKeyData->getSortKey());
         } else if (META_RECORDID == it->second) {
             bob.append(it->first, static_cast<long long>(member->loc.repr()));
         }
diff --git a/src/mongo/db/exec/projection_exec.h b/src/mongo/db/exec/projection_exec.h
index 8079ee8afa8..906d2fe12d8 100644
--- a/src/mongo/db/exec/projection_exec.h
+++ b/src/mongo/db/exec/projection_exec.h
@@ -190,9 +190,13 @@ private:
     // 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.
+    // Do we have a returnKey projection?  If so we *only* output the index key metadata, and
+    // possibly the sort key for mongos to use.  If it's not found we output nothing.
     bool _hasReturnKey;
+
+    // The field names associated with any sortKey meta-projection(s). Empty if there is no sortKey
+    // meta-projection.
+    std::vector<StringData> _sortKeyMetaFields;
 };
 
 }  // namespace mongo
diff --git a/src/mongo/db/query/parsed_projection.cpp b/src/mongo/db/query/parsed_projection.cpp
index 32979e9693e..4f6a48c2f9e 100644
--- a/src/mongo/db/query/parsed_projection.cpp
+++ b/src/mongo/db/query/parsed_projection.cpp
@@ -273,8 +273,8 @@ Status ParsedProjection::make(const BSONObj& spec,
         }
     }
 
-    // returnKey clobbers everything.
-    if (hasIndexKeyProjection) {
+    // returnKey clobbers everything except for sortKey meta-projection.
+    if (hasIndexKeyProjection && !wantSortKey) {
         pp->_requiresDocument = false;
     }