From e5e2937f8f0cd981308f5361cfb8e1b143b75e33 Mon Sep 17 00:00:00 2001
From: Valeri Karpov <val@karpov.io>
Date: Thu, 25 Feb 2021 09:55:44 -0500
Subject: [PATCH] fix(populate): mostly working `transform` option for virtual
 populate re: #3775

---
 lib/helpers/populate/assignVals.js | 11 ++--
 test/model.populate.test.js        | 84 ++++++++++++++++++++++++++++++
 2 files changed, 91 insertions(+), 4 deletions(-)

diff --git a/lib/helpers/populate/assignVals.js b/lib/helpers/populate/assignVals.js
index 2be723792d3..d5e3ccef04b 100644
--- a/lib/helpers/populate/assignVals.js
+++ b/lib/helpers/populate/assignVals.js
@@ -25,6 +25,7 @@ module.exports = function assignVals(o) {
 
   // replace the original ids in our intermediate _ids structure
   // with the documents found by query
+  o.allIds = [].concat(o.allIds);
   assignRawDocsToIdStructure(o.rawIds, o.rawDocs, o.rawOrder, populateOptions);
 
   // now update the original documents being populated using the
@@ -44,6 +45,8 @@ module.exports = function assignVals(o) {
       return val.val;
     }
 
+    const _allIds = o.allIds[i];
+
     if (o.justOne === true && Array.isArray(val)) {
       // Might be an embedded discriminator (re: gh-9244) with multiple models, so make sure to pick the right
       // model before assigning.
@@ -63,11 +66,11 @@ module.exports = function assignVals(o) {
         val[i] = ret[i];
       }
 
-      return valueFilter(val[0], options, populateOptions, o.allIds[i]);
+      return valueFilter(val[0], options, populateOptions, _allIds);
     } else if (o.justOne === false && !Array.isArray(val)) {
-      return valueFilter([val], options, populateOptions, o.allIds[i]);
+      return valueFilter([val], options, populateOptions, _allIds);
     }
-    return valueFilter(val, options, populateOptions, o.allIds[i]);
+    return valueFilter(val, options, populateOptions, _allIds);
   }
 
   for (i = 0; i < docs.length; ++i) {
@@ -209,7 +212,7 @@ function valueFilter(val, assignmentOpts, populateOptions, allIds) {
       if (!isPopulatedObject(subdoc) && (!populateOptions.retainNullValues || subdoc != null) && !userSpecifiedTransform) {
         continue;
       } else if (userSpecifiedTransform) {
-        subdoc = transform(isPopulatedObject(subdoc) ? subdoc : null, allIds[i]);
+        subdoc = transform(isPopulatedObject(subdoc) ? subdoc : null, allIds);
       }
       maybeRemoveId(subdoc, assignmentOpts);
       ret.push(subdoc);
diff --git a/test/model.populate.test.js b/test/model.populate.test.js
index ec262c64672..da7cd59ea4a 100644
--- a/test/model.populate.test.js
+++ b/test/model.populate.test.js
@@ -10078,4 +10078,88 @@ describe('model: populate:', function() {
       assert.equal(called[0].id.toHexString(), newId.toHexString());
     });
   });
+
+  it('transform with virtual populate, justOne = true (gh-3375)', function() {
+    const parentSchema = new Schema({
+      name: String
+    });
+    parentSchema.virtual('child', {
+      ref: 'Child',
+      localField: '_id',
+      foreignField: 'parentId',
+      justOne: true
+    });
+    const Parent = db.model('Parent', parentSchema);
+
+    const Child = db.model('Child', Schema({ name: String, parentId: 'ObjectId' }));
+
+    return co(function*() {
+      let p = yield Parent.create({ name: 'Anakin' });
+      const child = yield Child.create({ name: 'Luke', parentId: p._id });
+
+      let called = [];
+
+      p = yield Parent.findById(p).populate({
+        path: 'child',
+        transform: function(doc, id) {
+          called.push({
+            doc: doc,
+            id: id
+          });
+
+          return id;
+        }
+      });
+
+      assert.equal(called.length, 1);
+      assert.strictEqual(called[0].doc.parentId.toHexString(), p._id.toHexString());
+      assert.equal(called[0].id.toHexString(), p._id.toHexString());
+    });
+  });
+
+  it('transform with virtual populate, justOne = false (gh-3375) XYZ', function() {
+    const parentSchema = new Schema({
+      name: String
+    });
+    parentSchema.virtual('children', {
+      ref: 'Child',
+      localField: '_id',
+      foreignField: 'parentId',
+      justOne: false
+    });
+    const Parent = db.model('Parent', parentSchema);
+
+    const Child = db.model('Child', Schema({ name: String, parentId: 'ObjectId' }));
+
+    return co(function*() {
+      let p = yield Parent.create({ name: 'Anakin' });
+      const children = yield Child.create([
+        { name: 'Luke', parentId: p._id },
+        { name: 'Leia', parentId: p._id }
+      ]);
+
+      let called = [];
+
+      p = yield Parent.findById(p).populate({
+        path: 'children',
+        transform: function(doc, id) {
+          called.push({
+            doc: doc,
+            id: id
+          });
+
+          return id;
+        }
+      });
+
+      assert.equal(called.length, 2);
+      assert.deepEqual(called.map(c => c.doc.name).sort(), ['Leia', 'Luke']);
+
+      assert.strictEqual(called[0].doc.parentId.toHexString(), p._id.toHexString());
+      assert.equal(called[0].id.toHexString(), p._id.toHexString());
+
+      assert.strictEqual(called[1].doc.parentId.toHexString(), p._id.toHexString());
+      assert.equal(called[1].id.toHexString(), p._id.toHexString());
+    });
+  });
 });