Merge pull request #11600 from Automattic/gh-7971

(Gh 7971) Can define discriminators in the schema

lib/schema.js

@@ -540,6 +540,12 @@ Schema.prototype.add = function add(obj, prefix) {
         this.nested[prefix.substring(0, prefix.length - 1)] = true;
       }
       this.path(prefix + key, val);
+      if (val[0] != null && !(val[0].instanceOfSchema) && utils.isPOJO(val[0].discriminators)) {
+        const schemaType = this.path(prefix + key);
+        for (const key in val[0].discriminators) {
+          schemaType.discriminator(key, val[0].discriminators[key]);
+        }
+      }
     } else if (Object.keys(val).length < 1) {
       // Special-case: {} always interpreted as Mixed path so leaf at this node
       if (prefix) {
@@ -569,6 +575,12 @@ Schema.prototype.add = function add(obj, prefix) {
           this.nested[prefix.substring(0, prefix.length - 1)] = true;
         }
         this.path(prefix + key, val);
+        if (val != null && !(val.instanceOfSchema) && utils.isPOJO(val.discriminators)) {
+          const schemaType = this.path(prefix + key);
+          for (const key in val.discriminators) {
+            schemaType.discriminator(key, val.discriminators[key]);
+          }
+        }
       }
     }
   }

test/model.discriminator.test.js

@@ -1930,4 +1930,32 @@ describe('model', function() {
     assert.equal(updatedDoc.sheetOptions.location.id, 'inColumn');
     assert.equal(updatedDoc.sheetOptions.location.timeout, 2000);
   });
+
+  it('allows defining discriminator at the subSchema level in the subschema (gh-7971)', async function() {
+    const eventSchema = new Schema({ message: String },
+      { discriminatorKey: 'kind', _id: false });
+    const clickedSchema = new Schema({
+      element: {
+        type: String,
+        required: true
+      }
+    }, { _id: false });
+    const batchSchema = new Schema({ events: [eventSchema], mainEvent: { type: eventSchema, discriminators: { Clicked: clickedSchema } } });
+    const arraySchema = new Schema({ arrayEvent: [{ type: eventSchema, discriminators: { Clicked: clickedSchema } }] });
+    const Batch = db.model('Batch', batchSchema);
+    const Arrays = db.model('Array', arraySchema);
+
+    const batch = await Batch.create({
+      events: [{ message: 'Hello World' }],
+      mainEvent: { message: 'Goodbye', element: 'The Discriminator', kind: 'Clicked' }
+    });
+
+    assert(batch.mainEvent.element);
+
+    const array = await Arrays.create({
+      arrayEvent: [{ message: 'An array', element: 'with discriminators', kind: 'Clicked' }]
+    });
+
+    assert(array.arrayEvent[0].element);
+  });
 });