SERVER-7491: Correctly generate projections for subfields of _id

This requires special casing due to SERVER-7502

2 files changed, 68 insertions, 2 deletions

diff --git a/src/mongo/db/pipeline/document_source.cpp b/src/mongo/db/pipeline/document_source.cpp
index e1f033192a9..a459ec0ad10 100755
--- a/src/mongo/db/pipeline/document_source.cpp
+++ b/src/mongo/db/pipeline/document_source.cpp
@@ -87,11 +87,16 @@ namespace mongo {
 
     BSONObj DocumentSource::depsToProjection(const set<string>& deps) {
         BSONObjBuilder bb;
-        if (deps.count("_id") == 0)
-            bb.append("_id", 0);
+
+        bool needId = false;
 
         string last;
         for (set<string>::const_iterator it(deps.begin()), end(deps.end()); it!=end; ++it) {
+            if (str::startsWith(*it, "_id") && (it->size() == 3 || (*it)[3] == '.')) {
+                // _id and subfields are handled specially due in part to SERVER-7502
+                needId = true;
+                continue;
+            }
             if (!last.empty() && str::startsWith(*it, last)) {
                 // we are including a parent of *it so we don't need to
                 // include this field explicitly. In fact, due to
@@ -102,6 +107,12 @@ namespace mongo {
             last = *it + '.';
             bb.append(*it, 1);
         }
+
+        if (needId) // we are explicit either way
+            bb.append("_id", 1);
+        else
+            bb.append("_id", 0);
+
         return bb.obj();
     }
 }
diff --git a/src/mongo/dbtests/documentsourcetests.cpp b/src/mongo/dbtests/documentsourcetests.cpp
index bd031a53dc0..63eb71be846 100644
--- a/src/mongo/dbtests/documentsourcetests.cpp
+++ b/src/mongo/dbtests/documentsourcetests.cpp
@@ -44,6 +44,59 @@ namespace DocumentSourceTests {
         }
     };
 
+    namespace DocumentSourceClass {
+        using mongo::DocumentSource;
+
+        template<size_t ArrayLen>
+        set<string> arrayToSet(const char* (&array) [ArrayLen]) {
+            set<string> out;
+            for (size_t i = 0; i < ArrayLen; i++)
+                out.insert(array[i]);
+            return out;
+        }
+
+        class Deps {
+        public:
+            void run() {
+                {
+                    const char* array[] = {"a", "b"}; // basic
+                    BSONObj proj = DocumentSource::depsToProjection(arrayToSet(array));
+                    ASSERT_EQUALS(proj, BSON("a" << 1 << "b" << 1 << "_id" << 0));
+                }
+                {
+                    const char* array[] = {"a", "ab"}; // prefixed but not subfield
+                    BSONObj proj = DocumentSource::depsToProjection(arrayToSet(array));
+                    ASSERT_EQUALS(proj, BSON("a" << 1 << "ab" << 1 << "_id" << 0));
+                }
+                {
+                    const char* array[] = {"a", "b", "a.b"}; // a.b included by a
+                    BSONObj proj = DocumentSource::depsToProjection(arrayToSet(array));
+                    ASSERT_EQUALS(proj, BSON("a" << 1 << "b" << 1 << "_id" << 0));
+                }
+                {
+                    const char* array[] = {"a", "_id"}; // _id now included
+                    BSONObj proj = DocumentSource::depsToProjection(arrayToSet(array));
+                    ASSERT_EQUALS(proj, BSON("a" << 1 << "_id" << 1));
+                }
+                {
+                    const char* array[] = {"a", "_id.a"}; // still include whole _id (SERVER-7502)
+                    BSONObj proj = DocumentSource::depsToProjection(arrayToSet(array));
+                    ASSERT_EQUALS(proj, BSON("a" << 1 << "_id" << 1));
+                }
+                {
+                    const char* array[] = {"a", "_id", "_id.a"}; // handle both _id and subfield
+                    BSONObj proj = DocumentSource::depsToProjection(arrayToSet(array));
+                    ASSERT_EQUALS(proj, BSON("a" << 1 << "_id" << 1));
+                }
+                {
+                    const char* array[] = {"a", "_id", "_id_a"}; // _id prefixed but non-subfield
+                    BSONObj proj = DocumentSource::depsToProjection(arrayToSet(array));
+                    ASSERT_EQUALS(proj, BSON("_id_a" << 1 << "a" << 1 << "_id" << 1));
+                }
+            }
+        };
+    }
+
     namespace DocumentSourceCursor {
 
         using mongo::DocumentSourceCursor;
@@ -1690,6 +1743,8 @@ namespace DocumentSourceTests {
         All() : Suite( "documentsource" ) {
         }
         void setupTests() {
+            add<DocumentSourceClass::Deps>();
+
             add<DocumentSourceCursor::Create>();
             add<DocumentSourceCursor::Iterate>();
             add<DocumentSourceCursor::Dispose>();