Merge pull request #14131 from Automattic/vkarpov15/gh-14109

fix: allow adding discriminators using `Schema.prototype.discriminator()` to subdocuments after defining parent schema
Automattic · Dec 1, 2023 · 42bf516 · 42bf516
2 parents e1426ca + 16e5308
commit 42bf516
Show file tree

Hide file tree

Showing 5 changed files with 124 additions and 12 deletions.
diff --git a/lib/helpers/discriminator/applyEmbeddedDiscriminators.js b/lib/helpers/discriminator/applyEmbeddedDiscriminators.js
@@ -0,0 +1,23 @@
+'use strict';
+
+module.exports = applyEmbeddedDiscriminators;
+
+function applyEmbeddedDiscriminators(schema, seen = new WeakSet()) {
+  if (seen.has(schema)) {
+    return;
+  }
+  seen.add(schema);
+  for (const path of Object.keys(schema.paths)) {
+    const schemaType = schema.paths[path];
+    if (!schemaType.schema) {
+      continue;
+    }
+    applyEmbeddedDiscriminators(schemaType.schema, seen);
+    if (!schemaType.schema._applyDiscriminators) {
+      continue;
+    }
+    for (const disc of schemaType.schema._applyDiscriminators.keys()) {
+      schemaType.discriminator(disc, schemaType.schema._applyDiscriminators.get(disc));
+    }
+  }
+}
diff --git a/lib/index.js b/lib/index.js
@@ -32,6 +32,7 @@ const sanitizeFilter = require('./helpers/query/sanitizeFilter');
 const isBsonType = require('./helpers/isBsonType');
 const MongooseError = require('./error/mongooseError');
 const SetOptionError = require('./error/setOptionError');
+const applyEmbeddedDiscriminators = require('./helpers/discriminator/applyEmbeddedDiscriminators');
 
 const defaultMongooseSymbol = Symbol.for('mongoose:default');
 
@@ -629,6 +630,8 @@ Mongoose.prototype._model = function(name, schema, collection, options) {
     }
   }
 
+  applyEmbeddedDiscriminators(schema);
+
   return model;
 };
 

diff --git a/lib/schema/SubdocumentPath.js b/lib/schema/SubdocumentPath.js
@@ -55,11 +55,6 @@ function SubdocumentPath(schema, path, options) {
   this.$isSingleNested = true;
   this.base = schema.base;
   SchemaType.call(this, path, options, 'Embedded');
-  if (schema._applyDiscriminators != null && !options?._skipApplyDiscriminators) {
-    for (const disc of schema._applyDiscriminators.keys()) {
-      this.discriminator(disc, schema._applyDiscriminators.get(disc));
-    }
-  }
 }
 
 /*!

diff --git a/lib/schema/documentarray.js b/lib/schema/documentarray.js
@@ -88,12 +88,6 @@ function DocumentArrayPath(key, schema, options, schemaOptions) {
 
   this.$embeddedSchemaType.caster = this.Constructor;
   this.$embeddedSchemaType.schema = this.schema;
-
-  if (schema._applyDiscriminators != null && !options?._skipApplyDiscriminators) {
-    for (const disc of schema._applyDiscriminators.keys()) {
-      this.discriminator(disc, schema._applyDiscriminators.get(disc));
-    }
-  }
 }
 
 /**
@@ -528,7 +522,7 @@ DocumentArrayPath.prototype.cast = function(value, doc, init, prev, options) {
 
 DocumentArrayPath.prototype.clone = function() {
   const options = Object.assign({}, this.options);
-  const schematype = new this.constructor(this.path, this.schema, { ...options, _skipApplyDiscriminators: true }, this.schemaOptions);
+  const schematype = new this.constructor(this.path, this.schema, options, this.schemaOptions);
   schematype.validators = this.validators.slice();
   if (this.requiredValidator !== undefined) {
     schematype.requiredValidator = this.requiredValidator;

diff --git a/test/document.test.js b/test/document.test.js
@@ -12657,6 +12657,103 @@ describe('document', function() {
     );
   });
 
+  it('handles embedded discriminators defined using Schema.prototype.discriminator after defining schema (gh-14109) (gh-13898)', async function() {
+    const baseNestedDiscriminated = new Schema({
+      type: { type: Number, required: true }
+    }, { discriminatorKey: 'type' });
+
+    class BaseClass {
+      whoAmI() {
+        return 'I am baseNestedDiscriminated';
+      }
+    }
+
+    baseNestedDiscriminated.loadClass(BaseClass);
+
+    class NumberTyped extends BaseClass {
+      whoAmI() {
+        return 'I am NumberTyped';
+      }
+    }
+
+    class StringTyped extends BaseClass {
+      whoAmI() {
+        return 'I am StringTyped';
+      }
+    }
+
+    const containsNestedSchema = new Schema({
+      nestedDiscriminatedTypes: { type: [baseNestedDiscriminated], required: true }
+    });
+
+    // After `containsNestedSchema`, in #13898 test these were before `containsNestedSchema`
+    baseNestedDiscriminated.discriminator(1, new Schema({}).loadClass(NumberTyped));
+    baseNestedDiscriminated.discriminator('3', new Schema({}).loadClass(StringTyped));
+
+    class ContainsNested {
+      whoAmI() {
+        return 'I am ContainsNested';
+      }
+    }
+    containsNestedSchema.loadClass(ContainsNested);
+
+    const Test = db.model('Test', containsNestedSchema);
+    const instance = await Test.create({ type: 1, nestedDiscriminatedTypes: [{ type: 1 }, { type: '3' }] });
+    assert.deepStrictEqual(
+      instance.nestedDiscriminatedTypes.map(i => i.whoAmI()),
+      ['I am NumberTyped', 'I am StringTyped']
+    );
+  });
+
+  it('handles embedded discriminators on nested path defined using Schema.prototype.discriminator (gh-14109) (gh-13898)', async function() {
+    const baseNestedDiscriminated = new Schema({
+      type: { type: Number, required: true }
+    }, { discriminatorKey: 'type' });
+
+    class BaseClass {
+      whoAmI() {
+        return 'I am baseNestedDiscriminated';
+      }
+    }
+
+    baseNestedDiscriminated.loadClass(BaseClass);
+
+    class NumberTyped extends BaseClass {
+      whoAmI() {
+        return 'I am NumberTyped';
+      }
+    }
+
+    class StringTyped extends BaseClass {
+      whoAmI() {
+        return 'I am StringTyped';
+      }
+    }
+
+    const containsNestedSchema = new Schema({
+      nestedDiscriminatedTypes: { type: [baseNestedDiscriminated], required: true }
+    });
+
+    baseNestedDiscriminated.discriminator(1, new Schema({}).loadClass(NumberTyped));
+    baseNestedDiscriminated.discriminator('3', new Schema({}).loadClass(StringTyped));
+
+    const l2Schema = new Schema({ l3: containsNestedSchema });
+    const l1Schema = new Schema({ l2: l2Schema });
+
+    const Test = db.model('Test', l1Schema);
+    const instance = await Test.create({
+      l2: {
+        l3: {
+          nestedDiscriminatedTypes: [{ type: 1 }, { type: '3' }]
+        }
+      }
+    });
+    assert.deepStrictEqual(
+      instance.l2.l3.nestedDiscriminatedTypes.map(i => i.whoAmI()),
+      ['I am NumberTyped', 'I am StringTyped']
+    );
+  });
+
   it('can use `collection` as schema name (gh-13956)', async function() {
     const schema = new mongoose.Schema({ name: String, collection: String });
     const Test = db.model('Test', schema);