feat: support getters on populate virtuals, including get option fo…

…r `Schema#virtual()` Fix #9343
Automattic · Nov 7, 2020 · ad1d18a · ad1d18a
1 parent df86f53
commit ad1d18a
Show file tree

Hide file tree

Showing 2 changed files with 94 additions and 28 deletions.
diff --git a/lib/schema.js b/lib/schema.js
@@ -1674,6 +1674,7 @@ Schema.prototype.indexes = function() {
  * @param {String|Function} [options.foreignField] Required for populate virtuals. See [populate virtual docs](populate.html#populate-virtuals) for more information.
  * @param {Boolean|Function} [options.justOne=false] Only works with populate virtuals. If [truthy](https://masteringjs.io/tutorials/fundamentals/truthy), will be a single doc or `null`. Otherwise, the populate virtual will be an array.
  * @param {Boolean} [options.count=false] Only works with populate virtuals. If [truthy](https://masteringjs.io/tutorials/fundamentals/truthy), this populate virtual will contain the number of documents rather than the documents themselves when you `populate()`.
+ * @param {Function|null} [options.get=null] Adds a [getter](/docs/tutorials/getters-setters.html) to this virtual to transform the populated doc.
  * @return {VirtualType}
  */
 
@@ -1716,38 +1717,46 @@ Schema.prototype.virtual = function(name, options) {
 
     const virtual = this.virtual(name);
     virtual.options = options;
-    return virtual.
-      get(function(_v) {
-        if (this.$$populatedVirtuals &&
-          this.$$populatedVirtuals.hasOwnProperty(name)) {
-          return this.$$populatedVirtuals[name];
-        }
-        if (_v == null) return undefined;
-        return _v;
-      }).
-      set(function(_v) {
-        if (!this.$$populatedVirtuals) {
-          this.$$populatedVirtuals = {};
-        }
+    process.nextTick(() => {
+      virtual.
+        get(function(_v) {
+          if (this.$$populatedVirtuals &&
+            this.$$populatedVirtuals.hasOwnProperty(name)) {
+            return this.$$populatedVirtuals[name];
+          }
+          if (_v == null) return undefined;
+          return _v;
+        }).
+        set(function(_v) {
+          if (!this.$$populatedVirtuals) {
+            this.$$populatedVirtuals = {};
+          }
 
-        if (options.justOne || options.count) {
-          this.$$populatedVirtuals[name] = Array.isArray(_v) ?
-            _v[0] :
-            _v;
+          if (options.justOne || options.count) {
+            this.$$populatedVirtuals[name] = Array.isArray(_v) ?
+              _v[0] :
+              _v;
 
-          if (typeof this.$$populatedVirtuals[name] !== 'object') {
-            this.$$populatedVirtuals[name] = options.count ? _v : null;
+            if (typeof this.$$populatedVirtuals[name] !== 'object') {
+              this.$$populatedVirtuals[name] = options.count ? _v : null;
+            }
+          } else {
+            this.$$populatedVirtuals[name] = Array.isArray(_v) ?
+              _v :
+              _v == null ? [] : [_v];
+
+            this.$$populatedVirtuals[name] = this.$$populatedVirtuals[name].filter(function(doc) {
+              return doc && typeof doc === 'object';
+            });
           }
-        } else {
-          this.$$populatedVirtuals[name] = Array.isArray(_v) ?
-            _v :
-            _v == null ? [] : [_v];
+        });
+    });
 
-          this.$$populatedVirtuals[name] = this.$$populatedVirtuals[name].filter(function(doc) {
-            return doc && typeof doc === 'object';
-          });
-        }
-      });
+    if (typeof options.get === 'function') {
+      virtual.get(options.get);
+    }
+
+    return virtual;
   }
 
   const virtuals = this.virtuals;

diff --git a/test/model.populate.test.js b/test/model.populate.test.js
@@ -4708,6 +4708,63 @@ describe('model: populate:', function() {
         });
       });
 
+      it('virtuals with getters (gh-9343)', function() {
+        const UserSchema = new Schema({
+          openId: String,
+          test: String
+        });
+        const CommentSchema = new Schema({
+          openId: String
+        });
+
+        CommentSchema.virtual('user', {
+          ref: 'User',
+          localField: 'openId',
+          foreignField: 'openId',
+          justOne: true
+        }).get(v => v.test);
+
+        const User = db.model('User', UserSchema);
+        const Comment = db.model('Comment', CommentSchema);
+
+        return co(function*() {
+          yield Comment.create({ openId: 'test' });
+          yield User.create({ openId: 'test', test: 'my string' });
+
+          const comment = yield Comment.findOne({ openId: 'test' }).populate('user');
+          assert.equal(comment.user, 'my string');
+        });
+      });
+
+      it('virtuals with `get` option (gh-9343)', function() {
+        const UserSchema = new Schema({
+          openId: String,
+          test: String
+        });
+        const CommentSchema = new Schema({
+          openId: String
+        });
+
+        CommentSchema.virtual('user', {
+          ref: 'User',
+          localField: 'openId',
+          foreignField: 'openId',
+          justOne: true,
+          get: v => v.test
+        });
+
+        const User = db.model('User', UserSchema);
+        const Comment = db.model('Comment', CommentSchema);
+
+        return co(function*() {
+          yield Comment.create({ openId: 'test' });
+          yield User.create({ openId: 'test', test: 'my string' });
+
+          const comment = yield Comment.findOne({ openId: 'test' }).populate('user');
+          assert.equal(comment.user, 'my string');
+        });
+      });
+
       it('hydrates properly (gh-4618)', function(done) {
         const ASchema = new Schema({
           name: { type: String }