Merge pull request #13858 from Automattic/vkarpov15/gh-13788

fix(schema): handle number discriminator keys when using `Schema.prototype.discriminator()`
Automattic · Sep 15, 2023 · f9bfae0 · f9bfae0
2 parents 87619d0 + 4c38056
commit f9bfae0
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 12 deletions.
diff --git a/lib/index.js b/lib/index.js
@@ -624,8 +624,8 @@ Mongoose.prototype._model = function(name, schema, collection, options) {
   connection.emit('model', model);
 
   if (schema._applyDiscriminators != null) {
-    for (const disc of Object.keys(schema._applyDiscriminators)) {
-      model.discriminator(disc, schema._applyDiscriminators[disc]);
+    for (const disc of schema._applyDiscriminators.keys()) {
+      model.discriminator(disc, schema._applyDiscriminators.get(disc));
     }
   }
 

diff --git a/lib/schema.js b/lib/schema.js
@@ -449,7 +449,7 @@ Schema.prototype._clone = function _clone(Constructor) {
     s.discriminators = Object.assign({}, this.discriminators);
   }
   if (this._applyDiscriminators != null) {
-    s._applyDiscriminators = Object.assign({}, this._applyDiscriminators);
+    s._applyDiscriminators = new Map(this._applyDiscriminators);
   }
 
   s.aliases = Object.assign({}, this.aliases);
@@ -621,7 +621,8 @@ Schema.prototype.defaultOptions = function(options) {
  * @api public
  */
 Schema.prototype.discriminator = function(name, schema) {
-  this._applyDiscriminators = Object.assign(this._applyDiscriminators || {}, { [name]: schema });
+  this._applyDiscriminators = this._applyDiscriminators || new Map();
+  this._applyDiscriminators.set(name, schema);
 
   return this;
 };
@@ -722,18 +723,18 @@ Schema.prototype.add = function add(obj, prefix) {
         for (const key in val[0].discriminators) {
           schemaType.discriminator(key, val[0].discriminators[key]);
         }
-      } else if (val[0] != null && val[0].instanceOfSchema && utils.isPOJO(val[0]._applyDiscriminators)) {
-        const applyDiscriminators = val[0]._applyDiscriminators || [];
+      } else if (val[0] != null && val[0].instanceOfSchema && val[0]._applyDiscriminators instanceof Map) {
+        const applyDiscriminators = val[0]._applyDiscriminators;
         const schemaType = this.path(prefix + key);
-        for (const disc in applyDiscriminators) {
-          schemaType.discriminator(disc, applyDiscriminators[disc]);
+        for (const disc of applyDiscriminators.keys()) {
+          schemaType.discriminator(disc, applyDiscriminators.get(disc));
         }
       }
-      else if (val != null && val.instanceOfSchema && utils.isPOJO(val._applyDiscriminators)) {
-        const applyDiscriminators = val._applyDiscriminators || [];
+      else if (val != null && val.instanceOfSchema && val._applyDiscriminators instanceof Map) {
+        const applyDiscriminators = val._applyDiscriminators;
         const schemaType = this.path(prefix + key);
-        for (const disc in applyDiscriminators) {
-          schemaType.discriminator(disc, applyDiscriminators[disc]);
+        for (const disc of applyDiscriminators.keys()) {
+          schemaType.discriminator(disc, applyDiscriminators.get(disc));
         }
       }
     } else if (Object.keys(val).length < 1) {

diff --git a/test/schema.test.js b/test/schema.test.js
@@ -2873,6 +2873,43 @@ describe('schema', function() {
     assert(batch.message);
   });
 
+  it('supports numbers with Schema.discriminator() (gh-13788)', async() => {
+    const baseClassSchema = new Schema({
+      type: { type: Number, required: true }
+    }, { discriminatorKey: 'type' });
+
+    class BaseClass {
+      whoAmI() {
+        return 'I am base';
+      }
+    }
+    BaseClass.type = 1;
+
+    baseClassSchema.loadClass(BaseClass);
+
+    class NumberTyped extends BaseClass {
+      whoAmI() {
+        return 'I am NumberTyped';
+      }
+    }
+    NumberTyped.type = 2;
+
+    class StringTyped extends BaseClass {
+      whoAmI() {
+        return 'I am StringTyped';
+      }
+    }
+    StringTyped.type = '3';
+
+    baseClassSchema.discriminator(2, new Schema({}).loadClass(NumberTyped));
+    baseClassSchema.discriminator('3', new Schema({}).loadClass(StringTyped));
+    const Test = db.model('Test', { item: baseClassSchema });
+    let doc = await Test.create({ item: { type: 2 } });
+    assert.equal(doc.item.whoAmI(), 'I am NumberTyped');
+    doc = await Test.create({ item: { type: '3' } });
+    assert.equal(doc.item.whoAmI(), 'I am StringTyped');
+  });
+
   it('can use on as a schema property (gh-11580)', async() => {
     const testSchema = new mongoose.Schema({
       on: String