fix(map): support populate with maps

Re: #681

lib/document.js

@@ -2698,7 +2698,6 @@ Document.prototype.execPopulate = function() {
 
 Document.prototype.populated = function(path, val, options) {
   // val and options are internal
-
   if (val === null || val === void 0) {
     if (!this.$__.populated) {
       return undefined;

lib/model.js

@@ -3042,6 +3042,7 @@ function populate(model, docs, options, callback) {
   }
 
   modelsMap = getModelsMapForPopulate(model, docs, options);
+
   if (modelsMap instanceof Error) {
     return utils.immediate(function() {
       callback(modelsMap);
@@ -3273,6 +3274,19 @@ function assignVals(o) {
       rawIds[i] = rawIds[i][0];
     }
 
+    const existingVal = utils.getValue(o.path, docs[i]);
+    // If we're populating a map, the existing value will be an object, so
+    // we need to transform again
+    if (!o.isVirtual &&
+        existingVal != null &&
+        existingVal.constructor.name === 'Object') {
+      const _keys = Object.keys(existingVal);
+      rawIds[i] = rawIds[i].reduce((cur, v, i) => {
+        cur[_keys[i]] = v;
+        return cur;
+      }, {});
+    }
+
     if (o.isVirtual && docs[i].constructor.name === 'model') {
       // If virtual populate and doc is already init-ed, need to walk through
       // the actual doc to set rather than setting `_doc` directly
@@ -3289,6 +3303,7 @@ function assignVals(o) {
       if (docs[i].$__) {
         docs[i].populated(o.path, o.allIds[i], o.allOptions);
       }
+
       utils.setValue(o.path, rawIds[i], docs[i], setValue, false);
     }
   }
@@ -3472,6 +3487,7 @@ function getModelsMapForPopulate(model, docs, options) {
     if (typeof foreignField === 'function') {
       foreignField = foreignField.call(doc);
     }
+
     const ret = convertTo_id(utils.getValue(localField, doc));
     const id = String(utils.getValue(foreignField, doc));
     options._docs[id] = Array.isArray(ret) ? ret.slice() : ret;
@@ -3534,6 +3550,9 @@ function getModelsMapForPopulate(model, docs, options) {
     if (schema && schema.caster) {
       schema = schema.caster;
     }
+    if (schema && schema.$isSchemaMap) {
+      schema = schema.$__schemaType;
+    }
 
     if (!schema && model.discriminators) {
       discriminatorKey = model.schema.discriminatorMapping.key;
@@ -3622,7 +3641,7 @@ function convertTo_id(val) {
   if (val instanceof Model) return val._id;
 
   if (Array.isArray(val)) {
-    for (var i = 0; i < val.length; ++i) {
+    for (let i = 0; i < val.length; ++i) {
       if (val[i] instanceof Model) {
         val[i] = val[i]._id;
       }
@@ -3634,6 +3653,15 @@ function convertTo_id(val) {
     return [].concat(val);
   }
 
+  // If `populate('map')`, `val` will be an object
+  if (val != null && val.constructor.name === 'Object') {
+    const ret = [];
+    for (const key of Object.keys(val)) {
+      ret.push(val[key]);
+    }
+    return ret;
+  }
+
   return val;
 }
 

lib/types/map.js

@@ -38,12 +38,24 @@ class MongooseMap extends Map {
       return;
     }
 
-    try {
-      value = this.$__schemaType.
-        applySetters(value, this.$__parent, false, this.get(key));
-    } catch (error) {
-      this.$__parent.invalidate(this.$__path + '.' + key, error);
-      return;
+    const fullPath = this.$__path + '.' + key;
+    const populated = this.$__parent != null && this.$__parent.$__ ?
+      this.$__parent.populated(fullPath) || this.$__parent.populated(this.$__path) :
+      null;
+
+    if (populated != null) {
+      if (value.$__ == null) {
+        value = new populated.options.model(value);
+      }
+      value.$__.wasPopulated = true;
+    } else {
+      try {
+        value = this.$__schemaType.
+          applySetters(value, this.$__parent, false, this.get(key));
+      } catch (error) {
+        this.$__parent.invalidate(fullPath, error);
+        return;
+      }
     }
 
     super.set(key, value);

test/types.map.test.js

@@ -218,6 +218,70 @@ describe('Map', function() {
     });
   });
 
+  it('populate', function() {
+    const UserSchema = new mongoose.Schema({
+      keys: {
+        type: Map,
+        of: {
+          type: mongoose.Schema.Types.ObjectId,
+          ref: 'MapPopulateTest'
+        }
+      }
+    });
+
+    const KeySchema = new mongoose.Schema({ key: String });
+
+    const User = db.model('MapPopulateTest_0', UserSchema);
+    const Key = db.model('MapPopulateTest', KeySchema);
+
+    return co(function*() {
+      const key = yield Key.create({ key: 'abc123' });
+      const key2 = yield Key.create({ key: 'key' });
+
+      const doc = yield User.create({ keys: { github: key._id } });
+
+      const populated = yield User.findById(doc).populate('keys.github');
+
+      assert.equal(populated.keys.get('github').key, 'abc123');
+
+      populated.keys.set('twitter', key2._id);
+
+      yield populated.save();
+
+      const rawDoc = yield User.collection.findOne({ _id: doc._id });
+      assert.deepEqual(rawDoc.keys, { github: key._id, twitter: key2._id });
+    });
+  });
+
+  it('populate with wildcard', function() {
+    const UserSchema = new mongoose.Schema({
+      apiKeys: {
+        type: Map,
+        of: {
+          type: mongoose.Schema.Types.ObjectId,
+          ref: 'MapPopulateWildcardTest'
+        }
+      }
+    });
+
+    const KeySchema = new mongoose.Schema({ key: String });
+
+    const User = db.model('MapPopulateWildcardTest_0', UserSchema);
+    const Key = db.model('MapPopulateWildcardTest', KeySchema);
+
+    return co(function*() {
+      const key = yield Key.create({ key: 'abc123' });
+      const key2 = yield Key.create({ key: 'key' });
+
+      const doc = yield User.create({ apiKeys: { github: key._id, twitter: key2._id } });
+
+      const populated = yield User.findById(doc).populate('apiKeys');
+
+      assert.equal(populated.apiKeys.get('github').key, 'abc123');
+      assert.equal(populated.apiKeys.get('twitter').key, 'key');
+    });
+  });
+
   it('discriminators', function() {
     const TestSchema = new mongoose.Schema({
       n: Number