Merge pull request #14500 from Automattic/vkarpov15/gh-14444

fix(model): make `Model.recompileSchema()` also re-apply discriminators

Author: vkarpov15

lib/helpers/discriminator/applyEmbeddedDiscriminators.js

@@ -2,7 +2,7 @@
 
 module.exports = applyEmbeddedDiscriminators;
 
-function applyEmbeddedDiscriminators(schema, seen = new WeakSet()) {
+function applyEmbeddedDiscriminators(schema, seen = new WeakSet(), overwriteExisting = false) {
   if (seen.has(schema)) {
     return;
   }
@@ -16,13 +16,17 @@ function applyEmbeddedDiscriminators(schema, seen = new WeakSet()) {
     if (!schemaType.schema._applyDiscriminators) {
       continue;
     }
-    if (schemaType._appliedDiscriminators) {
+    if (schemaType._appliedDiscriminators && !overwriteExisting) {
       continue;
     }
     for (const discriminatorKey of schemaType.schema._applyDiscriminators.keys()) {
       const discriminatorSchema = schemaType.schema._applyDiscriminators.get(discriminatorKey);
       applyEmbeddedDiscriminators(discriminatorSchema, seen);
-      schemaType.discriminator(discriminatorKey, discriminatorSchema);
+      schemaType.discriminator(
+        discriminatorKey,
+        discriminatorSchema,
+        overwriteExisting ? { overwriteExisting: true } : null
+      );
     }
     schemaType._appliedDiscriminators = true;
   }

lib/helpers/model/discriminator.js

@@ -21,7 +21,7 @@ const CUSTOMIZABLE_DISCRIMINATOR_OPTIONS = {
  * ignore
  */
 
-module.exports = function discriminator(model, name, schema, tiedValue, applyPlugins, mergeHooks) {
+module.exports = function discriminator(model, name, schema, tiedValue, applyPlugins, mergeHooks, overwriteExisting) {
   if (!(schema && schema.instanceOfSchema)) {
     throw new Error('You must pass a valid discriminator Schema');
   }
@@ -205,7 +205,7 @@ module.exports = function discriminator(model, name, schema, tiedValue, applyPlu
 
   model.schema.discriminators[name] = schema;
 
-  if (model.discriminators[name] && !schema.options.overwriteModels) {
+  if (model.discriminators[name] && !schema.options.overwriteModels && !overwriteExisting) {
     throw new Error('Discriminator with name "' + name + '" already exists');
   }
 

lib/model.js

@@ -23,6 +23,7 @@ const VersionError = require('./error/version');
 const ParallelSaveError = require('./error/parallelSave');
 const applyDefaultsHelper = require('./helpers/document/applyDefaults');
 const applyDefaultsToPOJO = require('./helpers/model/applyDefaultsToPOJO');
+const applyEmbeddedDiscriminators = require('./helpers/discriminator/applyEmbeddedDiscriminators');
 const applyHooks = require('./helpers/model/applyHooks');
 const applyMethods = require('./helpers/model/applyMethods');
 const applyProjection = require('./helpers/projection/applyProjection');
@@ -5031,6 +5032,14 @@ Model.__subclass = function subclass(conn, schema, collection) {
 
 Model.recompileSchema = function recompileSchema() {
   this.prototype.$__setSchema(this.schema);
+
+  if (this.schema._applyDiscriminators != null) {
+    for (const disc of this.schema._applyDiscriminators.keys()) {
+      this.discriminator(disc, this.schema._applyDiscriminators.get(disc));
+    }
+  }
+
+  applyEmbeddedDiscriminators(this.schema, new WeakSet(), true);
 };
 
 /**

lib/schema/subdocument.js

@@ -325,7 +325,7 @@ SchemaSubdocument.prototype.discriminator = function(name, schema, options) {
     schema = schema.clone();
   }
 
-  schema = discriminator(this.caster, name, schema, value);
+  schema = discriminator(this.caster, name, schema, value, null, null, options.overwriteExisting);
 
   this.caster.discriminators[name] = _createConstructor(schema, this.caster);
 

lib/schemaType.js

@@ -878,7 +878,7 @@ SchemaType.prototype.validateAll = function(validators) {
  *
  *     schema.path('name').validate({
  *       validator: function() { throw new Error('Oops!'); },
- *       // `errors['name']` will be "Oops!"
+ *       // `errors['name'].message` will be "Oops!"
  *       message: function(props) { return props.reason.message; }
  *     });
  *

test/model.test.js

@@ -7419,6 +7419,50 @@ describe('Model', function() {
     assert.equal(doc.myVirtual, 'Hello from myVirtual');
   });
 
+  it('supports recompiling model with new discriminators (gh-14444) (gh-14296)', function() {
+    // Define discriminated schema
+    const decoratorSchema = new Schema({
+      type: { type: String, required: true }
+    }, { discriminatorKey: 'type' });
+
+    class Decorator {
+      whoAmI() { return 'I am BaseDeco'; }
+    }
+    decoratorSchema.loadClass(Decorator);
+
+    // Define discriminated class before model is compiled
+    class Deco1 extends Decorator { whoAmI() { return 'I am Test1'; }}
+    const deco1Schema = new Schema({});
+    deco1Schema.loadClass(Deco1);
+    decoratorSchema.discriminator('Test1', deco1Schema);
+
+    // Define model that uses discriminated schema
+    const shopSchema = new Schema({
+      item: { type: decoratorSchema, required: true }
+    });
+    const shopModel = db.model('Test', shopSchema);
+
+    // Define another discriminated class after the model is compiled
+    class Deco2 extends Decorator { whoAmI() { return 'I am Test2'; }}
+    const deco2Schema = new Schema({});
+    deco2Schema.loadClass(Deco2);
+    decoratorSchema.discriminator('Test2', deco2Schema);
+
+    let instance = new shopModel({ item: { type: 'Test1' } });
+    assert.equal(instance.item.whoAmI(), 'I am Test1');
+
+    instance = new shopModel({ item: { type: 'Test2' } });
+    assert.equal(instance.item.whoAmI(), 'I am BaseDeco');
+
+    shopModel.recompileSchema();
+
+    instance = new shopModel({ item: { type: 'Test1' } });
+    assert.equal(instance.item.whoAmI(), 'I am Test1');
+
+    instance = new shopModel({ item: { type: 'Test2' } });
+    assert.equal(instance.item.whoAmI(), 'I am Test2');
+  });
+
   it('inserts versionKey even if schema has `toObject.versionKey` set to false (gh-14344)', async function() {
     const schema = new mongoose.Schema(
       { name: String },