SERVER-7502 test that partial projection of _id works correctly

4 files changed, 48 insertions, 22 deletions

diff --git a/jstests/aggregation/sources/project/remove_redundant_projects.js b/jstests/aggregation/sources/project/remove_redundant_projects.js
index e4584df4f65..ea620be2c2a 100644
--- a/jstests/aggregation/sources/project/remove_redundant_projects.js
+++ b/jstests/aggregation/sources/project/remove_redundant_projects.js
@@ -140,13 +140,15 @@ assertResultsMatch({
     removedProjectStage: {_id: 0, a: 1}
 });
 
-// Test that projections on _id with nested fields are not removed from pipeline. Due to
-// SERVER-7502, the dependency analysis does not generate a covered projection for nested
-// fields in _id and thus we cannot remove the stage.
+// Test that projections on _id with nested fields are removed from pipeline.
 indexSpec = {
     '_id.a': 1,
     a: 1
 };
-assertResultsMatch(
-    {pipeline: [{$project: {'_id.a': 1}}], expectProjectToCoalesce: false, index: indexSpec});
+assertResultsMatch({
+    pipeline: [{$match: {"_id.a": 1}}, {$project: {'_id.a': 1}}],
+    expectProjectToCoalesce: true,
+    index: indexSpec,
+    pipelineOptimizedAway: true,
+});
 }());
diff --git a/jstests/core/id_partial_projection.js b/jstests/core/id_partial_projection.js
new file mode 100644
index 00000000000..7affbcfb875
--- /dev/null
+++ b/jstests/core/id_partial_projection.js
@@ -0,0 +1,21 @@
+/**
+ * Tests partial inclusion/exclusion of _id.
+ * See SERVER-7502 for details.
+ */
+(function() {
+"use strict";
+
+const coll = db.id_partial_projection;
+coll.drop();
+
+assert.commandWorked(coll.insert({_id: {a: 1, b: 1}, otherField: 1}));
+assert.commandWorked(coll.insert({_id: 3, otherField: 2}));
+
+assert.eq(coll.find({}, {"_id": 1}).toArray(), [{_id: {a: 1, b: 1}}, {_id: 3}]);
+assert.eq(coll.find({}, {"_id.a": 1}).toArray(), [{_id: {a: 1}}, {}]);
+assert.eq(coll.find({}, {"_id.b": 1}).toArray(), [{_id: {b: 1}}, {}]);
+
+assert.eq(coll.find({}, {"_id.a": 0}).toArray(),
+          [{_id: {b: 1}, otherField: 1}, {_id: 3, otherField: 2}]);
+assert.eq(coll.find({}, {_id: 0}).toArray(), [{otherField: 1}, {otherField: 2}]);
+})();
diff --git a/src/mongo/db/pipeline/dependencies.cpp b/src/mongo/db/pipeline/dependencies.cpp
index 7bc9532ed6f..cfe063b976e 100644
--- a/src/mongo/db/pipeline/dependencies.cpp
+++ b/src/mongo/db/pipeline/dependencies.cpp
@@ -51,21 +51,17 @@ BSONObj DepsTracker::toProjectionWithoutMetadata() const {
         return bb.obj();
     }
 
-    bool needId = false;
+    bool idSpecified = false;
     std::string last;
     for (const auto& field : fields) {
         if (str::startsWith(field, "_id") && (field.size() == 3 || field[3] == '.')) {
-            // _id and subfields are handled specially due in part to SERVER-7502
-            needId = true;
-            continue;
+            idSpecified = true;
         }
 
         if (!last.empty() && str::startsWith(field, last)) {
             // we are including a parent of *it so we don't need to include this field
-            // explicitly. In fact, due to SERVER-6527 if we included this field, the parent
-            // wouldn't be fully included.  This logic relies on on set iterators going in
-            // lexicographic order so that a string is always directly before of all fields it
-            // prefixes.
+            // explicitly. This logic relies on on set iterators going in lexicographic order so
+            // that a string is always directly before of all fields it prefixes.
             continue;
         }
 
@@ -79,10 +75,9 @@ BSONObj DepsTracker::toProjectionWithoutMetadata() const {
         bb.append(field, 1);
     }
 
-    if (needId)  // we are explicit either way
-        bb.append("_id", 1);
-    else
+    if (!idSpecified) {
         bb.append("_id", 0);
+    }
 
     return bb.obj();
 }
diff --git a/src/mongo/db/pipeline/dependencies_test.cpp b/src/mongo/db/pipeline/dependencies_test.cpp
index f5b9c5a37a5..b6c688da8b0 100644
--- a/src/mongo/db/pipeline/dependencies_test.cpp
+++ b/src/mongo/db/pipeline/dependencies_test.cpp
@@ -87,21 +87,21 @@ TEST(DependenciesToProjectionTest, ShouldIncludeIdIfNeeded) {
     const char* array[] = {"a", "_id"};
     DepsTracker deps;
     deps.fields = arrayToSet(array);
-    ASSERT_BSONOBJ_EQ(deps.toProjectionWithoutMetadata(), BSON("a" << 1 << "_id" << 1));
+    ASSERT_BSONOBJ_EQ(deps.toProjectionWithoutMetadata(), BSON("_id" << 1 << "a" << 1));
 }
 
-TEST(DependenciesToProjectionTest, ShouldIncludeEntireIdEvenIfOnlyASubFieldIsNeeded) {
-    const char* array[] = {"a", "_id.a"};  // still include whole _id (SERVER-7502)
+TEST(DependenciesToProjectionTest, ShouldNotIncludeEntireIdIfOnlyASubFieldIsNeeded) {
+    const char* array[] = {"a", "_id.a"};
     DepsTracker deps;
     deps.fields = arrayToSet(array);
-    ASSERT_BSONOBJ_EQ(deps.toProjectionWithoutMetadata(), BSON("a" << 1 << "_id" << 1));
+    ASSERT_BSONOBJ_EQ(deps.toProjectionWithoutMetadata(), BSON("_id.a" << 1 << "a" << 1));
 }
 
 TEST(DependenciesToProjectionTest, ShouldNotIncludeSubFieldOfIdIfIdIncluded) {
     const char* array[] = {"a", "_id", "_id.a"};  // handle both _id and subfield
     DepsTracker deps;
     deps.fields = arrayToSet(array);
-    ASSERT_BSONOBJ_EQ(deps.toProjectionWithoutMetadata(), BSON("a" << 1 << "_id" << 1));
+    ASSERT_BSONOBJ_EQ(deps.toProjectionWithoutMetadata(), BSON("_id" << 1 << "a" << 1));
 }
 
 TEST(DependenciesToProjectionTest, ShouldIncludeFieldPrefixedById) {
@@ -109,7 +109,15 @@ TEST(DependenciesToProjectionTest, ShouldIncludeFieldPrefixedById) {
     DepsTracker deps;
     deps.fields = arrayToSet(array);
     ASSERT_BSONOBJ_EQ(deps.toProjectionWithoutMetadata(),
-                      BSON("_id_a" << 1 << "a" << 1 << "_id" << 1));
+                      BSON("_id" << 1 << "_id_a" << 1 << "a" << 1));
+}
+
+TEST(DependenciesToProjectionTest, ShouldIncludeFieldPrefixedByIdWhenIdSubfieldIsIncluded) {
+    const char* array[] = {"a", "_id.a", "_id_a"};  // _id prefixed but non-subfield
+    DepsTracker deps;
+    deps.fields = arrayToSet(array);
+    ASSERT_BSONOBJ_EQ(deps.toProjectionWithoutMetadata(),
+                      BSON("_id.a" << 1 << "_id_a" << 1 << "a" << 1));
 }
 
 TEST(DependenciesToProjectionTest, ShouldOutputEmptyObjectIfEntireDocumentNeeded) {