Merge pull request #2773 from chetverikov/gh2719

Fix #2719
Automattic · Mar 19, 2015 · ad4462c · ad4462c
2 parents 5cf9fd4 + 80f8332
commit ad4462c
Show file tree

Hide file tree

Showing 2 changed files with 113 additions and 51 deletions.
diff --git a/lib/model.js b/lib/model.js
@@ -2182,12 +2182,7 @@ var excludeIdReg = /\s?-_id\s?/,
   excludeIdRegGlobal = /\s?-_id\s?/g;
 
 function populate(model, docs, options, cb) {
-  var schema = model._getSchema(options.path),
-    modelsMap, rawIds;
-
-  if (schema && schema.caster) {
-    schema = schema.caster;
-  }
+  var modelsMap, rawIds;
 
   // normalize single / multiple docs passed
   if (!Array.isArray(docs)) {
@@ -2198,7 +2193,7 @@ function populate(model, docs, options, cb) {
     return cb();
   }
 
-  modelsMap = getModelsMapForPopulate(model, schema, docs, options);
+  modelsMap = getModelsMapForPopulate(model, docs, options);
   rawIds = getIdsForAndAddIdsInMapPopulate(modelsMap);
 
   var i, len = modelsMap.length,
@@ -2294,65 +2289,91 @@ function populate(model, docs, options, cb) {
   }
 }
 
-function getModelsMapForPopulate(model, schema, docs, options) {
+function getModelsMapForPopulate(model, docs, options) {
   var i, doc, len = docs.length,
-    available = {}, map = [],
-    refPath = schema && schema.options && schema.options.refPath,
-    Model, currentOptions, modelNames, modelName;
+    available = {},
+    map = [],
+    modelNameFromQuery = options.model && options.model.modelName || options.model,
+    schema, refPath, Model, currentOptions, modelNames, modelName, discriminatorKey, modelForFindSchema;
 
-  if (refPath) {
-    for (i = 0; i < len; i++) {
-      doc = docs[i];
+  schema = model._getSchema(options.path);
+
+  if(schema && schema.caster){
+    schema = schema.caster;
+  }
+
+  if (!schema && model.discriminators){
+    discriminatorKey = model.schema.discriminatorMapping.key
+  }
+
+  refPath = schema && schema.options && schema.options.refPath;
+
+  for (i = 0; i < len; i++) {
+    doc = docs[i];
+
+    if(refPath){
       modelNames = utils.getValue(refPath, doc);
+    }else{
+      if(!modelNameFromQuery){
+        var schemaForCurrentDoc;
 
-      if (!modelNames)
-        continue;
+        if(!schema && discriminatorKey){
+          modelForFindSchema = utils.getValue(discriminatorKey, doc);
 
-      if (!Array.isArray(modelNames)) {
-        modelNames = [modelNames];
-      }
+          if(modelForFindSchema){
+            schemaForCurrentDoc = model.db.model(modelForFindSchema)._getSchema(options.path);
 
-      var k = modelNames.length;
-      while (k--) {
-        modelName = modelNames[k];
-        if (!available[modelName]) {
-          Model = model.db.model(modelName);
-          currentOptions = {
-            model: Model
-          };
-
-          utils.merge(currentOptions, options);
-
-          available[modelName] = {
-            Model: Model,
-            options: currentOptions,
-            docs: [doc],
-            ids: []
-          };
-          map.push(available[modelName]);
+            if(schema && schema.caster){
+              schema = schema.caster;
+            }
+          }
         } else {
-          available[modelName].docs.push(doc);
+          schemaForCurrentDoc = schema;
         }
 
+        modelNames = [
+          schemaForCurrentDoc && schemaForCurrentDoc.options && schemaForCurrentDoc.options.ref            // declared in schema
+          || model.modelName                                           // an ad-hoc structure
+        ]
+      }else{
+        modelNames = [modelNameFromQuery];  // query options
       }
     }
 
-  } else {
+    if (!modelNames)
+      continue;
 
-    // model name for the populate query
-    modelName = options.model && options.model.modelName || options.model // query options
-    || schema && schema.options && schema.options.ref // declared in schema
-    || model.modelName; // an ad-hoc structure
+    if (!Array.isArray(modelNames)) {
+      modelNames = [modelNames];
+    }
 
-    Model = model.db.model(modelName);
-    options.model = Model;
+    var k = modelNames.length;
+    while (k--) {
+      modelName = modelNames[k];
+      if (!available[modelName]) {
+        Model = model.db.model(modelName);
+        currentOptions = {
+          model: Model
+        };
+
+        if(schema && !discriminatorKey){
+          options.model = Model;
+        }
 
-    map = [{
-      Model: Model,
-      options: options,
-      docs: docs,
-      ids: []
-    }];
+        utils.merge(currentOptions, options);
+
+        available[modelName] = {
+          Model: Model,
+          options: currentOptions,
+          docs: [doc],
+          ids: []
+        };
+        map.push(available[modelName]);
+      } else {
+        available[modelName].docs.push(doc);
+      }
+
+    }
   }
 
   return map;

diff --git a/test/model.discriminator.querying.test.js b/test/model.discriminator.querying.test.js
@@ -705,6 +705,47 @@ describe('model', function() {
           });
         });
       });
+
+      it('reference in child schemas (gh-2719)', function(done){
+        var vehicleSchema = new Schema({});
+        var carSchema = new Schema({
+          speed: Number,
+          garage: {type: Schema.Types.ObjectId, ref: 'gh2719PopulationGarage'}
+        });
+        var busSchema = new Schema({
+          speed: Number,
+          garage: {type: Schema.Types.ObjectId, ref: 'gh2719PopulationGarage'}
+        });
+
+        var garageSchema = new Schema({
+          name: String,
+          num_of_places: Number
+        });
+
+        var Vehicle = db.model('gh2719PopulationVehicle', vehicleSchema)
+          , Car = Vehicle.discriminator('gh2719PopulationCar', carSchema)
+          , Bus = Vehicle.discriminator('gh2719PopulationBus', busSchema)
+          , Garage = db.model('gh2719PopulationGarage', garageSchema);
+
+        Garage.create({name: 'My', num_of_places: 3}, function(err, garage){
+          assert.ifError(err);
+          Car.create({ speed: 160, garage: garage }, function(err, car) {
+            assert.ifError(err);
+            Bus.create({ speed: 80, garage: garage }, function(err, bus) {
+              assert.ifError(err);
+              Vehicle.find({}).populate('garage').exec(function(err, vehicles){
+                assert.ifError(err);
+
+                vehicles.forEach(function(v){
+                  assert.ok(v.garage instanceof Garage);
+                });
+
+                done();
+              });
+            })
+          });
+        });
+      })
     });
 
     describe('aggregate', function() {