fix(populate): handle populating paths on documents with discriminato…

…r keys that point to non-existent discriminators Fix #10082
Automattic · Apr 16, 2021 · 7b2b1bf · 7b2b1bf
1 parent 854374b
commit 7b2b1bf
Show file tree

Hide file tree

Showing 3 changed files with 55 additions and 21 deletions.
diff --git a/lib/helpers/populate/getModelsMapForPopulate.js b/lib/helpers/populate/getModelsMapForPopulate.js
@@ -30,7 +30,6 @@ module.exports = function getModelsMapForPopulate(model, docs, options) {
   let currentOptions;
   let modelNames;
   let modelName;
-  let modelForFindSchema;
 
   const originalModel = options.model;
   let isVirtual = false;
@@ -405,28 +404,26 @@ module.exports = function getModelsMapForPopulate(model, docs, options) {
     } else {
       let modelForCurrentDoc = model;
       let schemaForCurrentDoc;
-
-      if (!schema && discriminatorKey) {
-        modelForFindSchema = utils.getValue(discriminatorKey, doc);
-        if (modelForFindSchema) {
-          // `modelForFindSchema` is the discriminator value, so we might need
-          // find the discriminated model name
-          const discriminatorModel = getDiscriminatorByValue(model, modelForFindSchema);
-          if (discriminatorModel != null) {
-            modelForCurrentDoc = discriminatorModel;
-          } else {
-            try {
-              modelForCurrentDoc = model.db.model(modelForFindSchema);
-            } catch (error) {
-              return error;
-            }
+      let discriminatorValue;
+
+      if (!schema && discriminatorKey && (discriminatorValue = utils.getValue(discriminatorKey, doc))) {
+        // `modelNameForFind` is the discriminator value, so we might need
+        // find the discriminated model name
+        const discriminatorModel = getDiscriminatorByValue(model, discriminatorValue) || model;
+        if (discriminatorModel != null) {
+          modelForCurrentDoc = discriminatorModel;
+        } else {
+          try {
+            modelForCurrentDoc = model.db.model(discriminatorValue);
+          } catch (error) {
+            return error;
           }
+        }
 
-          schemaForCurrentDoc = modelForCurrentDoc.schema._getSchema(options.path);
+        schemaForCurrentDoc = modelForCurrentDoc.schema._getSchema(options.path);
 
-          if (schemaForCurrentDoc && schemaForCurrentDoc.caster) {
-            schemaForCurrentDoc = schemaForCurrentDoc.caster;
-          }
+        if (schemaForCurrentDoc && schemaForCurrentDoc.caster) {
+          schemaForCurrentDoc = schemaForCurrentDoc.caster;
         }
       } else {
         schemaForCurrentDoc = schema;

diff --git a/lib/helpers/populate/getSchemaTypes.js b/lib/helpers/populate/getSchemaTypes.js
@@ -139,7 +139,7 @@ module.exports = function getSchemaTypes(schema, doc, path) {
           if (type.schema.discriminators) {
             const discriminatorPaths = [];
             for (const discriminatorName of Object.keys(type.schema.discriminators)) {
-              const _schema = type.schema.discriminators[discriminatorName];
+              const _schema = type.schema.discriminators[discriminatorName] || type.schema;
               const ret = search(parts.slice(p), _schema, null, nestedPath.concat(parts.slice(0, p)));
               if (ret != null) {
                 discriminatorPaths.push(ret);

diff --git a/test/model.discriminator.querying.test.js b/test/model.discriminator.querying.test.js
@@ -871,6 +871,43 @@ describe('model', function() {
           assert.ok(events[0].surveys[0] instanceof Survey);
         });
       });
+
+      it('correctly populates doc with nonexistent discriminator key (gh-10082)', function() {
+        const foodSchema = Schema({ name: String, animal: String });
+        const Food = db.model('Food', foodSchema);
+
+        const animalSchema = Schema({ type: String }, {
+          discriminatorKey: 'type',
+          toJSON: { virtuals: true },
+          toObject: { virtuals: true }
+        });
+        const catSchema = Schema({ catYears: Number });
+        animalSchema.virtual('foods', {
+          ref: 'Food',
+          localField: 'type',
+          foreignField: 'animal',
+          justOne: false
+        });
+        const Animal = db.model('Animal', animalSchema);
+        Animal.discriminator('cat', catSchema);
+        // const Rabbit = Animal.discriminator('rabbit', rabbitSchema);
+
+        return co(function*() {
+          yield Promise.all([
+            new Food({ name: 'Cat Food', animal: 'cat' }).save(),
+            new Food({ name: 'Rabbit Food', animal: 'rabbit' }).save()
+          ]);
+          yield Animal.collection.insertOne({ type: 'cat', catYears: 4 });
+          yield Animal.collection.insertOne({ type: 'rabbit' }); // <-- "rabbit" has no discriminator
+
+          const cat = yield Animal.findOne({ type: 'cat' }).populate('foods');
+          const rabbit = yield Animal.findOne({ type: 'rabbit' }).populate('foods');
+          assert.equal(cat.foods.length, 1);
+          assert.equal(cat.foods[0].name, 'Cat Food');
+          assert.equal(rabbit.foods.length, 1);
+          assert.equal(rabbit.foods[0].name, 'Rabbit Food');
+        });
+      });
     });
 
     describe('deleteOne and deleteMany (gh-8471)', function() {