fix(document): handle embedded recursive discriminators on nested pat…

…h defined using Schema.prototype.discriminator Fix #14245
Automattic · Jan 12, 2024 · d3e22af · d3e22af
1 parent ac9af5b
commit d3e22af
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 5 deletions.
diff --git a/lib/helpers/discriminator/applyEmbeddedDiscriminators.js b/lib/helpers/discriminator/applyEmbeddedDiscriminators.js
@@ -19,11 +19,10 @@ function applyEmbeddedDiscriminators(schema, seen = new WeakSet()) {
     if (schemaType._appliedDiscriminators) {
       continue;
     }
-    for (const disc of schemaType.schema._applyDiscriminators.keys()) {
-      schemaType.discriminator(
-        disc,
-        schemaType.schema._applyDiscriminators.get(disc)
-      );
+    for (const discriminatorKey of schemaType.schema._applyDiscriminators.keys()) {
+      const discriminatorSchema = schemaType.schema._applyDiscriminators.get(discriminatorKey);
+      applyEmbeddedDiscriminators(discriminatorSchema, seen);
+      schemaType.discriminator(discriminatorKey, discriminatorSchema);
     }
     schemaType._appliedDiscriminators = true;
   }

diff --git a/test/document.test.js b/test/document.test.js
@@ -12787,6 +12787,54 @@ describe('document', function() {
     await doc2.save();
   });
 
+  it('handles embedded recursive discriminators on nested path defined using Schema.prototype.discriminator (gh-14245)', async function() {
+    const baseSchema = new Schema({
+      type: { type: Number, required: true }
+    }, { discriminatorKey: 'type' });
+
+    class Base {
+      whoAmI() { return 'I am Base'; }
+    }
+
+    baseSchema.loadClass(Base);
+
+    class NumberTyped extends Base {
+      whoAmI() { return 'I am NumberTyped'; }
+    }
+
+    class StringTyped extends Base {
+      whoAmI() { return 'I am StringTyped'; }
+    }
+
+    const selfRefSchema = new Schema({
+      self: { type: [baseSchema], required: true }
+    });
+
+    class SelfReferenceTyped extends Base {
+      whoAmI() { return 'I am SelfReferenceTyped'; }
+    }
+
+    selfRefSchema.loadClass(SelfReferenceTyped);
+    baseSchema.discriminator(5, selfRefSchema);
+
+    const numberTypedSchema = new Schema({}).loadClass(NumberTyped);
+    const stringTypedSchema = new Schema({}).loadClass(StringTyped);
+    baseSchema.discriminator(1, numberTypedSchema);
+    baseSchema.discriminator(3, stringTypedSchema);
+    const containerSchema = new Schema({ items: [baseSchema] });
+    const containerModel = db.model('Test', containerSchema);
+
+    const instance = await containerModel.create({
+      items: [{ type: 5, self: [{ type: 1 }, { type: 3 }] }]
+    });
+
+    assert.equal(instance.items[0].whoAmI(), 'I am SelfReferenceTyped');
+    assert.deepStrictEqual(instance.items[0].self.map(item => item.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);