From aa90acddf5e8b5aca31f593d3f8b157ec339ab21 Mon Sep 17 00:00:00 2001
From: Valeri Karpov <val@karpov.io>
Date: Thu, 14 Sep 2023 15:57:46 -0400
Subject: [PATCH] fix(schema): handle number discriminator keys when using
 `Schema.prototype.discriminator()`

Fix #13788
---
 lib/index.js        |  4 ++--
 lib/schema.js       | 21 +++++++++++----------
 test/schema.test.js | 38 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 51 insertions(+), 12 deletions(-)

diff --git a/lib/index.js b/lib/index.js
index 6c551ae5dfa..8affae77804 100644
--- 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
index 261211c1fb6..b608961b753 100644
--- 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
index 1f7a028f1b8..701ba9d9fd3 100644
--- a/test/schema.test.js
+++ b/test/schema.test.js
@@ -2873,6 +2873,44 @@ 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