feat(query): support for casting array operators in $expr

Automattic · Jan 27, 2022 · 94671bc · 94671bc
1 parent 29b323c
commit 94671bc
Show file tree

Hide file tree

Showing 2 changed files with 78 additions and 4 deletions.
diff --git a/lib/helpers/query/cast$expr.js b/lib/helpers/query/cast$expr.js
@@ -7,6 +7,9 @@ const castNumber = require('../../cast/number');
 const booleanComparison = new Set(['$and', '$or', '$not']);
 const comparisonOperator = new Set(['$cmp', '$eq', '$lt', '$lte', '$gt', '$gte']);
 const arithmeticOperatorArray = new Set([
+  // avoid casting '$add' or '$subtract', because expressions can be either number or date,
+  // and we don't have a good way of inferring which arguments should be numbers and which should
+  // be dates.
   '$multiply',
   '$divide',
   '$log',
@@ -44,6 +47,11 @@ const arithmeticOperatorNumber = new Set([
   '$degreesToRadians',
   '$radiansToDegrees'
 ]);
+const arrayElementOperators = new Set([
+  '$arrayElemAt',
+  '$first',
+  '$last'
+]);
 const dateOperators = new Set([
   '$year',
   '$month',
@@ -52,7 +60,11 @@ const dateOperators = new Set([
   '$dayOfYear',
   '$hour',
   '$minute',
-  '$second'
+  '$second',
+  '$isoDayOfWeek',
+  '$isoWeekYear',
+  '$isoWeek',
+  '$millisecond'
 ]);
 
 module.exports = function cast$expr(val, schema, strictQuery) {
@@ -93,10 +105,17 @@ function _castExpression(val, schema, strictQuery) {
     } else if (arithmeticOperatorArray.has(key)) {
       val[key] = castArithmetic(val[key], schema, strictQuery);
     } else if (arithmeticOperatorNumber.has(key)) {
-      val[key] = castArithmeticSingle(val[key], schema, strictQuery);
+      val[key] = castNumberOperator(val[key], schema, strictQuery);
     }
   }
 
+  if (val.$in) {
+    val.$in = castIn(val.$in, schema, strictQuery);
+  }
+  if (val.$size) {
+    val.$size = castNumberOperator(val.$size, schema, strictQuery);
+  }
+
   _omitUndefined(val);
 
   return val;
@@ -112,7 +131,7 @@ function _omitUndefined(val) {
 }
 
 // { $op: <number> }
-function castArithmeticSingle(val) {
+function castNumberOperator(val) {
   if (!isLiteral(val)) {
     return val;
   }
@@ -124,6 +143,37 @@ function castArithmeticSingle(val) {
   }
 }
 
+function castIn(val, schema, strictQuery) {
+  let search = val[0];
+  let path = val[1];
+  if (!isPath(path)) {
+    return val;
+  }
+
+  path = path.slice(1);
+  const schematype = schema.path(path);
+  if (schematype == null) {
+    if (strictQuery === false) {
+      return val;
+    } else if (strictQuery === 'throw') {
+      throw new StrictModeError('$in');
+    }
+
+    return void 0;
+  }
+
+  if (!schematype.$isMongooseArray) {
+    throw new Error('Path must be an array for $in');
+  }
+
+  if (schematype.$isMongooseDocumentArray) {
+    search = schematype.$embeddedSchemaType.cast(search);
+  } else {
+    search = schematype.caster.cast(search);
+  }
+  return [search, val[1]];
+}
+
 // { $op: [<number>, <number>] }
 function castArithmetic(val) {
   if (!Array.isArray(val)) {
@@ -170,6 +220,16 @@ function castComparison(val, schema, strictQuery) {
         if (dateOperators.has(key) && isPath(lhs[key])) {
           path = lhs[key].slice(1) + '.' + key;
           caster = castNumber;
+        } else if (arrayElementOperators.has(key) && isPath(lhs[key])) {
+          path = lhs[key].slice(1) + '.' + key;
+          schematype = schema.path(lhs[key].slice(1));
+          if (schematype != null) {
+            if (schematype.$isMongooseDocumentArray) {
+              schematype = schematype.$embeddedSchemaType;
+            } else if (schematype.$isMongooseArray) {
+              schematype = schematype.caster;
+            }
+          }
         }
       }
     }

diff --git a/test/helpers/query.cast$expr.test.js b/test/helpers/query.cast$expr.test.js
@@ -6,7 +6,7 @@ const cast$expr = require('../../lib/helpers/query/cast$expr');
 
 describe('castexpr', function() {
   it('casts comparisons', function() {
-    const testSchema = new Schema({ date: Date, spent: Number, budget: Number });
+    const testSchema = new Schema({ date: Date, spent: Number, budget: Number, nums: [Number] });
 
     let res = cast$expr({ $eq: ['$date', '2021-06-01'] }, testSchema);
     assert.deepEqual(res, { $eq: ['$date', new Date('2021-06-01')] });
@@ -25,6 +25,9 @@ describe('castexpr', function() {
 
     res = cast$expr({ $gt: ['$spent', '$budget'] }, testSchema);
     assert.deepStrictEqual(res, { $gt: ['$spent', '$budget'] });
+
+    res = cast$expr({ $gt: [{ $last: '$nums' }, '42'] }, testSchema);
+    assert.deepStrictEqual(res, { $gt: [{ $last: '$nums' }, 42] });
   });
 
   it('casts conditions', function() {
@@ -87,4 +90,15 @@ describe('castexpr', function() {
       cast$expr({ $eq: [{ $year: '$date' }, 'not a number'] }, testSchema);
     }, /Cast to Number failed/);
   });
+
+  it('casts $in', function() {
+    const testSchema = new Schema({ nums: [Number], docs: [new Schema({ prop: Number }, { _id: false })] });
+
+    let res = cast$expr({ $in: ['42', '$nums'] }, testSchema);
+    assert.deepStrictEqual(res, { $in: [42, '$nums'] });
+
+    res = cast$expr({ $in: [{ prop: '42' }, '$docs'] }, testSchema);
+    res.$in[0] = res.$in[0].toBSON(); // So `deepStrictEqual()` doesn't complain about subdoc internals
+    assert.deepStrictEqual(res, { $in: [{ prop: 42 }, '$docs'] });
+  });
 });