Merge pull request #13496 from Automattic/vkarpov15/gh-13492

fix(query): handle non-string discriminator key values in query
Automattic · Jun 12, 2023 · 99e5f48 · 99e5f48
2 parents c0087c6 + fb0539a
commit 99e5f48
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 3 deletions.
diff --git a/lib/cast.js b/lib/cast.js
@@ -118,9 +118,17 @@ module.exports = function cast(schema, obj, options, context) {
             discriminatorKey != null &&
             pathLastHalf !== discriminatorKey) {
             const discriminatorVal = get(obj, pathFirstHalf + '.' + discriminatorKey);
-            if (discriminatorVal != null) {
-              schematype = _schematype.schema.discriminators[discriminatorVal].
-                path(pathLastHalf);
+            const discriminators = _schematype.schema.discriminators;
+            if (typeof discriminatorVal === 'string' && discriminators[discriminatorVal] != null) {
+
+              schematype = discriminators[discriminatorVal].path(pathLastHalf);
+            } else if (discriminatorVal != null &&
+              Object.keys(discriminatorVal).length === 1 &&
+              Array.isArray(discriminatorVal.$in) &&
+              discriminatorVal.$in.length === 1 &&
+              typeof discriminatorVal.$in[0] === 'string' &&
+              discriminators[discriminatorVal.$in[0]] != null) {
+              schematype = discriminators[discriminatorVal.$in[0]].path(pathLastHalf);
             }
           }
         }

diff --git a/test/cast.test.js b/test/cast.test.js
@@ -193,4 +193,34 @@ describe('cast: ', function() {
       name: 'foo'
     });
   });
+
+  it('handles $in with discriminators if $in has exactly 1 element (gh-13492)', function() {
+    const itemStateFooSchema = new Schema({
+      fieldFoo: String
+    });
+
+    const itemSchema = new Schema({
+      state: new Schema(
+        {
+          field: String
+        },
+        {
+          discriminatorKey: 'type'
+        }
+      )
+    });
+
+    itemSchema.path('state').discriminator('FOO', itemStateFooSchema);
+
+    const res = cast(itemSchema, {
+      'state.type': {
+        $in: ['FOO']
+      },
+      'state.fieldFoo': 44
+    });
+    assert.deepStrictEqual(res, {
+      'state.type': { $in: ['FOO'] },
+      'state.fieldFoo': '44'
+    });
+  });
 });