SERVER-72416 Find and findAndModify commands' ExpressionContext should inherit collection level collation

3 files changed, 276 insertions, 45 deletions

diff --git a/jstests/core/collation.js b/jstests/core/collation.js
index 12843b40864..1bbdc2cd47a 100644
--- a/jstests/core/collation.js
+++ b/jstests/core/collation.js
@@ -27,7 +27,6 @@ load("jstests/libs/index_catalog_helpers.js");
 load("jstests/concurrency/fsm_workload_helpers/server_types.js");
 // For isReplSet
 load("jstests/libs/fixture_helpers.js");
-load("jstests/libs/sbe_explain_helpers.js");  // For engineSpecificAssertion.
 // For areAllCollectionsClustered.
 load("jstests/libs/clustered_collections/clustered_collection_util.js");
 
@@ -758,9 +757,7 @@ if (!isClustered) {
     assert.commandWorked(testDb.createCollection(coll.getName(), {collation: {locale: "en_US"}}));
     explainRes = coll.explain("executionStats").find({_id: "foo"}).finish();
     assert.commandWorked(explainRes);
-    let classicAssert = null !== getPlanStage(getWinningPlan(explainRes.queryPlanner), "IDHACK");
-    let sbeAssert = null !== getPlanStage(getWinningPlan(explainRes.queryPlanner), "IXSCAN");
-    engineSpecificAssertion(classicAssert, sbeAssert, testDb, explainRes);
+    assert.neq(null, getPlanStage(explainRes.executionStats.executionStages, "IDHACK"), explainRes);
 
     // Find on _id should use idhack stage when explicitly given query collation matches
     // collection default.
diff --git a/jstests/core/or_to_in.js b/jstests/core/or_to_in.js
index 3c9cacfbcdf..5d3c745dc95 100644
--- a/jstests/core/or_to_in.js
+++ b/jstests/core/or_to_in.js
@@ -25,19 +25,25 @@ function compareValues(v1, v2) {
 }
 
 // Check that 'expectedQuery' and 'actualQuery' have the same plans, and produce the same result.
-function assertEquivPlanAndResult(expectedQuery, actualQuery) {
+function assertEquivPlanAndResult(expectedQuery, actualQuery, supportWithCollation) {
     const expectedExplain = coll.find(expectedQuery).explain("queryPlanner");
     const actualExplain = coll.find(actualQuery).explain("queryPlanner");
 
     // The queries must be rewritten into the same form.
-    assert.docEq(expectedExplain.parsedQuery, actualExplain.parsedQuery);
+    assert.docEq(expectedExplain.queryPlanner.parsedQuery, actualExplain.queryPlanner.parsedQuery);
 
-    // Check if the test queries produce the same plans with collations
-    const expectedExplainColln =
+    // We are always running these queries to ensure a server crash is not triggered.
+    // TODO SERVER-72450: Add appropriate assertions for the output.
+    const expectedExplainCollation =
         coll.find(expectedQuery).sort({f1: 1}).collation({locale: 'en_US'}).explain("queryPlanner");
-    const actualExplainColln =
+    const actualExplainCollation =
         coll.find(actualQuery).sort({f1: 1}).collation({locale: 'en_US'}).explain("queryPlanner");
-    assert.docEq(expectedExplainColln.parsedQuery, actualExplainColln.parsedQuery);
+
+    if (supportWithCollation) {
+        // Check if the test queries produce the same plans with collations.
+        assert.docEq(expectedExplainCollation.queryPlanner.parsedQuery,
+                     actualExplainCollation.queryPlanner.parsedQuery);
+    }
 
     // Make sure both queries have the same access plan.
     const expectedPlan = getWinningPlan(expectedExplain.queryPlanner);
@@ -51,12 +57,13 @@ function assertEquivPlanAndResult(expectedQuery, actualQuery) {
     const actualRes = coll.find(actualQuery).toArray();
     assert(arrayEq(expectedRes, actualRes, false, compareValues),
            `expected=${expectedRes}, actual=${actualRes}`);
+
     // also with collation
-    const expectedResColln =
+    const expectedResCollation =
         coll.find(expectedQuery).sort({f1: 1}).collation({locale: 'en_US'}).toArray();
-    const actualResColln =
+    const actualResCollation =
         coll.find(actualQuery).sort({f1: 1}).collation({locale: 'en_US'}).toArray();
-    assert(arrayEq(expectedResColln, actualResColln, false, compareValues),
+    assert(arrayEq(expectedResCollation, actualResCollation, false, compareValues),
            `expected=${expectedRes}, actual=${actualRes}`);
 }
 
@@ -94,32 +101,71 @@ assert.commandWorked(coll.insert(data));
 
 // Pairs of queries where the first one is expressed via OR (which is supposed to be
 // rewritten as IN), and the second one is an equivalent query using IN.
+//
+// The third element of the array is optional, if present, implies that the rewrite is not
+// supported when there is a collation involved.
+//
+// TODO SERVER-72450: Remove or update this logic related to collation, and enforce stronger
+// assertions.
 const positiveTestQueries = [
-    [{$or: [{f1: 5}, {f1: 3}, {f1: 7}]}, {f1: {$in: [7, 3, 5]}}],
-    [{$or: [{f1: {$eq: 5}}, {f1: {$eq: 3}}, {f1: {$eq: 7}}]}, {f1: {$in: [7, 3, 5]}}],
-    [{$or: [{f1: 42}, {f1: NaN}, {f1: 99}]}, {f1: {$in: [42, NaN, 99]}}],
-    [{$or: [{f1: /^x/}, {f1: "ab"}]}, {f1: {$in: [/^x/, "ab"]}}],
-    [{$or: [{f1: /^x/}, {f1: "^a"}]}, {f1: {$in: [/^x/, "^a"]}}],
-    [{$or: [{f1: 42}, {f1: null}, {f1: 99}]}, {f1: {$in: [42, 99, null]}}],
-    [{$or: [{f1: 1}, {f2: 9}, {f1: 99}]}, {$or: [{f2: 9}, {f1: {$in: [1, 99]}}]}],
-    [{$or: [{f1: {$regex: /^x/}}, {f1: {$regex: /ab/}}]}, {f1: {$in: [/^x/, /ab/]}}],
-    [
-        {$and: [{$or: [{f1: 7}, {f1: 3}, {f1: 5}]}, {$or: [{f1: 1}, {f1: 2}, {f1: 3}]}]},
-        {$and: [{f1: {$in: [7, 3, 5]}}, {f1: {$in: [1, 2, 3]}}]}
-    ],
-    [
-        {$or: [{$or: [{f1: 7}, {f1: 3}, {f1: 5}]}, {$or: [{f1: 1}, {f1: 2}, {f1: 3}]}]},
-        {$or: [{f1: {$in: [7, 3, 5]}}, {f1: {$in: [1, 2, 3]}}]}
-    ],
-    [
-        {$or: [{$and: [{f1: 7}, {f2: 7}, {f1: 5}]}, {$or: [{f1: 1}, {f1: 2}, {f1: 3}]}]},
-        {$or: [{$and: [{f1: 7}, {f2: 7}, {f1: 5}]}, {f1: {$in: [1, 2, 3]}}]},
-    ],
-    [{$or: [{f2: [32, 52]}, {f2: [42, [13, 11]]}]}, {f2: {$in: [[32, 52], [42, [13, 11]]]}}],
-    [{$or: [{f2: 52}, {f2: 13}]}, {f2: {$in: [52, 13]}}],
-    [{$or: [{f2: [11]}, {f2: [23]}]}, {f2: {$in: [[11], [23]]}}],
-    [{$or: [{f1: 42}, {f1: null}]}, {f1: {$in: [42, null]}}],
-    [{$or: [{f1: "a"}, {f1: "b"}, {f1: /c/}]}, {f1: {$in: ["a", "b", /c/]}}],
+    {actualQuery: {$or: [{f1: 5}, {f1: 3}, {f1: 7}]}, expectedQuery: {f1: {$in: [7, 3, 5]}}},
+    {
+        actualQuery: {$or: [{f1: {$eq: 5}}, {f1: {$eq: 3}}, {f1: {$eq: 7}}]},
+        expectedQuery: {f1: {$in: [7, 3, 5]}}
+    },
+    {
+        actualQuery: {$or: [{f1: 42}, {f1: NaN}, {f1: 99}]},
+        expectedQuery: {f1: {$in: [42, NaN, 99]}}
+    },
+    {
+        actualQuery: {$or: [{f1: /^x/}, {f1: "ab"}]},
+        expectedQuery: {f1: {$in: [/^x/, "ab"]}},
+        cannotRewriteWithCollation: true
+    },
+    {
+        actualQuery: {$or: [{f1: /^x/}, {f1: "^a"}]},
+        expectedQuery: {f1: {$in: [/^x/, "^a"]}},
+        cannotRewriteWithCollation: true
+    },
+    {
+        actualQuery: {$or: [{f1: 42}, {f1: null}, {f1: 99}]},
+        expectedQuery: {f1: {$in: [42, 99, null]}}
+    },
+    {
+        actualQuery: {$or: [{f1: 1}, {f2: 9}, {f1: 99}]},
+        expectedQuery: {$or: [{f2: 9}, {f1: {$in: [1, 99]}}]}
+    },
+    {
+        actualQuery: {$or: [{f1: {$regex: /^x/}}, {f1: {$regex: /ab/}}]},
+        expectedQuery: {f1: {$in: [/^x/, /ab/]}}
+    },
+    {
+        actualQuery:
+            {$and: [{$or: [{f1: 7}, {f1: 3}, {f1: 5}]}, {$or: [{f1: 1}, {f1: 2}, {f1: 3}]}]},
+        expectedQuery: {$and: [{f1: {$in: [7, 3, 5]}}, {f1: {$in: [1, 2, 3]}}]}
+    },
+    {
+        actualQuery:
+            {$or: [{$or: [{f1: 7}, {f1: 3}, {f1: 5}]}, {$or: [{f1: 1}, {f1: 2}, {f1: 3}]}]},
+        expectedQuery: {$or: [{f1: {$in: [7, 3, 5]}}, {f1: {$in: [1, 2, 3]}}]}
+    },
+    {
+        actualQuery:
+            {$or: [{$and: [{f1: 7}, {f2: 7}, {f1: 5}]}, {$or: [{f1: 1}, {f1: 2}, {f1: 3}]}]},
+        expectedQuery: {$or: [{$and: [{f1: 7}, {f2: 7}, {f1: 5}]}, {f1: {$in: [1, 2, 3]}}]},
+    },
+    {
+        actualQuery: {$or: [{f2: [32, 52]}, {f2: [42, [13, 11]]}]},
+        expectedQuery: {f2: {$in: [[32, 52], [42, [13, 11]]]}}
+    },
+    {actualQuery: {$or: [{f2: 52}, {f2: 13}]}, expectedQuery: {f2: {$in: [52, 13]}}},
+    {actualQuery: {$or: [{f2: [11]}, {f2: [23]}]}, expectedQuery: {f2: {$in: [[11], [23]]}}},
+    {actualQuery: {$or: [{f1: 42}, {f1: null}]}, expectedQuery: {f1: {$in: [42, null]}}},
+    {
+        actualQuery: {$or: [{f1: "a"}, {f1: "b"}, {f1: /c/}]},
+        expectedQuery: {f1: {$in: ["a", "b", /c/]}},
+        cannotRewriteWithCollation: true
+    },
 ];
 
 // These $or queries should not be rewritten into $in because of different semantics.
@@ -133,22 +179,27 @@ for (const query of negativeTestQueries) {
     assertOrNotRewrittenToIn(query);
 }
 
-function testOrToIn(queries) {
+function testOrToIn(queries, usesCollation) {
     for (const queryPair of queries) {
-        assertEquivPlanAndResult(queryPair[0], queryPair[1]);
+        if (usesCollation && queryPair.cannotRewriteWithCollation) {
+            continue;
+        }
+        assertEquivPlanAndResult(
+            queryPair.actualQuery, queryPair.expectedQuery, !queryPair.cannotRewriteWithCollation);
     }
 }
 
-testOrToIn(positiveTestQueries);  // test without indexes
+testOrToIn(positiveTestQueries, false /* usesCollation */);  // test without indexes
 
 assert.commandWorked(coll.createIndex({f1: 1}));
 
-testOrToIn(positiveTestQueries);  // single index
+testOrToIn(positiveTestQueries, false /* usesCollation */);  // single index
 
 assert.commandWorked(coll.createIndex({f2: 1}));
 assert.commandWorked(coll.createIndex({f1: 1, f2: 1}));
 
-testOrToIn(positiveTestQueries);  // three indexes, requires multiplanning
+testOrToIn(positiveTestQueries,
+           false /* usesCollation */);  // three indexes, requires multiplanning
 
 // Test with a collection that has a collation, and that collation is the same as the query
 // collation
@@ -156,12 +207,12 @@ coll.drop();
 assert.commandWorked(db.createCollection("orToIn", {collation: {locale: 'en_US'}}));
 coll = db.orToIn;
 assert.commandWorked(coll.insert(data));
-testOrToIn(positiveTestQueries);
+testOrToIn(positiveTestQueries, true /* usesCollation */);
 // Test with a collection that has a collation, and that collation is different from the query
 // collation
 coll.drop();
 assert.commandWorked(db.createCollection("orToIn", {collation: {locale: 'de'}}));
 coll = db.orToIn;
 assert.commandWorked(coll.insert(data));
-testOrToIn(positiveTestQueries);
+testOrToIn(positiveTestQueries, true /* usesCollation */);
 }());
diff --git a/jstests/core/project_with_collation.js b/jstests/core/project_with_collation.js
new file mode 100644
index 00000000000..3f16716532f
--- /dev/null
+++ b/jstests/core/project_with_collation.js
@@ -0,0 +1,183 @@
+//  Tests to verify the behavior of find command's project in the presence of collation.
+//
+// @tags: [
+//   assumes_no_implicit_collection_creation_after_drop,
+// ]
+
+(function() {
+'use strict';
+
+const collation = {
+    locale: "en_US",
+    strength: 2
+};
+const withCollationCollName = jsTestName() + "_collation";
+const noCollationCollName = jsTestName() + "_noCollation";
+
+function setupCollection(withCollation) {
+    const insertCollName = withCollation ? withCollationCollName : noCollationCollName;
+    db[insertCollName].drop();
+    if (withCollation) {
+        assert.commandWorked(db.createCollection(insertCollName, {collation: withCollation}));
+    }
+
+    const insertColl = db[insertCollName];
+    assert.commandWorked(insertColl.insert(
+        {_id: 0, str: "a", array: [{str: "b"}, {str: "A"}, {str: "B"}, {str: "a"}]}));
+    assert.commandWorked(insertColl.insert({_id: 1, str: "a", elemMatch: [{str: "A"}, "ignored"]}));
+    assert.commandWorked(insertColl.insert({_id: 2, str: "A", elemMatch: ["ignored", {str: "a"}]}));
+    assert.commandWorked(insertColl.insert({_id: 3, str: "B"}));
+
+    return insertColl;
+}
+
+function runQueryWithCollation(testColl, collationToUse) {
+    let findCmd =
+        testColl.find({str: 'A', elemMatch: {$elemMatch: {str: "a"}}}, {_id: 1, 'elemMatch.$': 1});
+    if (collationToUse) {
+        findCmd = findCmd.collation(collationToUse);
+    }
+    const elemMatchOutput = findCmd.toArray();
+    assert.sameMembers(elemMatchOutput,
+                       [{_id: 1, elemMatch: [{str: "A"}]}, {_id: 2, elemMatch: [{str: "a"}]}]);
+
+    findCmd =
+        testColl.find({str: 'A'}, {sortedArray: {$sortArray: {input: "$array", sortBy: {str: 1}}}});
+    if (collationToUse) {
+        findCmd = findCmd.collation(collationToUse);
+    }
+    const sortArrayOutput = findCmd.toArray();
+    assert.sameMembers(sortArrayOutput, [
+        {_id: 0, sortedArray: [{str: "A"}, {str: "a"}, {str: "b"}, {str: "B"}]},
+        {_id: 1, sortedArray: null},
+        {_id: 2, sortedArray: null}
+    ]);
+
+    const findAndUpdateOutput = testColl.findAndModify({
+        query: {_id: 1, str: 'A', elemMatch: {$elemMatch: {str: "a"}}},
+        fields: {_id: 1, 'elemMatch.$': 1, updated: 1},
+        update: {$set: {updated: true}},
+        collation: collationToUse
+    });
+    assert.docEq(findAndUpdateOutput, {_id: 1, elemMatch: [{str: "A"}]});
+
+    const findAndUpdateWithSortArrayOutput = testColl.findAndModify({
+        query: {_id: 0, str: 'A'},
+        fields: {_id: 1, str: 1, sortedArray: {$sortArray: {input: "$array", sortBy: {str: 1}}}},
+        update: {$set: {updated: true}},
+        collation: collationToUse,
+        new: true
+    });
+    assert.docEq(findAndUpdateWithSortArrayOutput,
+                 {_id: 0, str: "a", sortedArray: [{str: "A"}, {str: "a"}, {str: "b"}, {str: "B"}]});
+
+    const findAndRemoveOutput = testColl.findAndModify({
+        query: {_id: 1, str: 'A', elemMatch: {$elemMatch: {str: "a"}}},
+        fields: {_id: 1, 'elemMatch.$': 1, updated: 1},
+        remove: true,
+        collation: collationToUse,
+    });
+    assert.docEq(findAndRemoveOutput, {_id: 1, elemMatch: [{str: "A"}], updated: true});
+
+    const findAndRemoveWithSortArrayOutput = testColl.findAndModify({
+        query: {_id: 0, str: 'A'},
+        fields: {_id: 1, str: 1, sortedArray: {$sortArray: {input: "$array", sortBy: {str: 1}}}},
+        remove: true,
+        collation: collationToUse
+    });
+    assert.docEq(findAndRemoveWithSortArrayOutput,
+                 {_id: 0, str: "a", sortedArray: [{str: "A"}, {str: "a"}, {str: "b"}, {str: "B"}]});
+}
+
+// The output for the below two tests should not depend on the collection level collation.
+let collWithCollation = setupCollection({locale: "en_US"});
+runQueryWithCollation(collWithCollation, collation);
+
+let noCollationColl = setupCollection(false);
+runQueryWithCollation(noCollationColl, collation);
+
+// Tests to verify that the projection code inherits collection level collation in the absence of
+// query level collation.
+collWithCollation = setupCollection(collation);
+runQueryWithCollation(collWithCollation, null);
+
+// The output of this should not depend on the collection level collation and simple collation
+// should be applied always.
+function queryWithSimpleCollation(testColl) {
+    const elemMatchOutput =
+        testColl.find({str: 'A', elemMatch: {$elemMatch: {str: "a"}}}, {_id: 1, 'elemMatch.$': 1})
+            .collation({locale: "simple"})
+            .toArray();
+    assert.sameMembers(elemMatchOutput, [{_id: 2, elemMatch: [{str: "a"}]}]);
+
+    const sortArrayOutput =
+        testColl.find({str: 'a'}, {sortedArray: {$sortArray: {input: "$array", sortBy: {str: 1}}}})
+            .collation({locale: "simple"})
+            .toArray();
+    assert.sameMembers(sortArrayOutput, [
+        {_id: 0, sortedArray: [{str: "A"}, {str: "B"}, {str: "a"}, {str: "b"}]},
+        {_id: 1, sortedArray: null}
+    ]);
+
+    // Test findAndModify command with 'update'. Ensure that simple collation is always honored.
+    const findAndUpdateOutput = testColl.findAndModify({
+        query: {str: 'A', elemMatch: {$elemMatch: {str: "a"}}},
+        fields: {_id: 1, 'elemMatch.$': 1, updated: 1},
+        update: {$set: {updated: true}},
+        collation: {locale: "simple"}
+    });
+    assert.docEq(findAndUpdateOutput, {_id: 2, elemMatch: [{str: "a"}]});
+
+    const findAndUpdateWithSortArrayOutput = testColl.findAndModify({
+        query: {_id: 0, str: 'a'},
+        fields: {_id: 1, str: 1, sortedArray: {$sortArray: {input: "$array", sortBy: {str: 1}}}},
+        update: {$set: {updated: true}},
+        collation: {locale: "simple"},
+        new: true
+    });
+    assert.docEq(findAndUpdateWithSortArrayOutput,
+                 {_id: 0, str: "a", sortedArray: [{str: "A"}, {str: "B"}, {str: "a"}, {str: "b"}]});
+
+    // Test findAndModify command with remove:true. Ensure that simple collation is always honored.
+    const findAndRemoveOutput = testColl.findAndModify({
+        query: {_id: 2, str: 'A', elemMatch: {$elemMatch: {str: "a"}}},
+        fields: {_id: 1, 'elemMatch.$': 1, updated: 1},
+        remove: true,
+        collation: {locale: "simple"},
+    });
+    assert.docEq(findAndRemoveOutput, {_id: 2, elemMatch: [{str: "a"}], updated: true});
+
+    const findAndRemoveWithSortArrayOutput = testColl.findAndModify({
+        query: {_id: 0, str: 'a'},
+        fields: {_id: 1, str: 1, sortedArray: {$sortArray: {input: "$array", sortBy: {str: 1}}}},
+        remove: true,
+        collation: {locale: "simple"}
+    });
+    assert.docEq(findAndRemoveWithSortArrayOutput,
+                 {_id: 0, str: "a", sortedArray: [{str: "A"}, {str: "B"}, {str: "a"}, {str: "b"}]});
+}
+
+noCollationColl = setupCollection(false);
+queryWithSimpleCollation(noCollationColl);
+
+collWithCollation = setupCollection(collation);
+queryWithSimpleCollation(collWithCollation);
+
+// Test with views.
+(function viewWithCollation() {
+    collWithCollation = setupCollection(collation);
+    db[jsTestName() + "_view"].drop();
+    assert.commandWorked(
+        db.createView(jsTestName() + "_view", withCollationCollName, [], {collation: collation}));
+    const viewColl = db[jsTestName() + "_view"];
+
+    const sortArrayOutput =
+        viewColl.find({str: 'A'}, {sortedArray: {$sortArray: {input: "$array", sortBy: {str: 1}}}})
+            .toArray();
+    assert.sameMembers(sortArrayOutput, [
+        {_id: 0, sortedArray: [{str: "A"}, {str: "a"}, {str: "b"}, {str: "B"}]},
+        {_id: 1, sortedArray: null},
+        {_id: 2, sortedArray: null}
+    ]);
+})();
+})();