// Tests that commands like find, aggregate and update accepts a 'let' parameter which defines
// variables for use in expressions within the command.
// Before 7.2, $rand was evaluated more than once in sharded find (SERVER-75927).
// @tags: [
//   # Requires a batch size greater than one.
//   does_not_support_config_fuzzer,
//   requires_fcv_81,
//   requires_getmore,
//   # Explain will return different plan than expected when a collection becomes a time-series
//   # collection. Also, query shape will be different.
//   exclude_from_timeseries_crud_passthrough,
// ]
//
import {FixtureHelpers} from "jstests/libs/fixture_helpers.js";
import {getPlanStage, getSingleNodeExplain, planHasStage} from "jstests/libs/query/analyze_plan.js";
import {checkSbeFullFeatureFlagEnabled} from "jstests/libs/query/sbe_util.js";

const testDB = db.getSiblingDB("command_let_variables");
const coll = testDB.command_let_variables;
coll.drop();

const isMongos = FixtureHelpers.isMongos(testDB);
const isCollSharded = FixtureHelpers.isSharded(coll);

const testDocs = [
    {
        _id: 1,
        Species: "Blackbird (Turdus merula)",
        population_trends: [
            {term: {start: 1970, end: 2014}, pct_change: -16, annual: -0.38, trend: "no change"},
            {term: {start: 2009, end: 2014}, pct_change: -2, annual: -0.36, trend: "no change"},
        ],
    },
    {
        _id: 2,
        Species: "Bullfinch (Pyrrhula pyrrhula)",
        population_trends: [
            {term: {start: 1970, end: 2014}, pct_change: -39, annual: -1.13, trend: "no change"},
            {term: {start: 2009, end: 2014}, pct_change: 12, annual: 2.38, trend: "weak increase"},
        ],
    },
    {
        _id: 3,
        Species: "Chaffinch (Fringilla coelebs)",
        population_trends: [
            {term: {start: 1970, end: 2014}, pct_change: 27, annual: 0.55, trend: "no change"},
            {term: {start: 2009, end: 2014}, pct_change: -7, annual: -1.49, trend: "weak decline"},
        ],
    },
    {
        _id: 4,
        Species: "Song Thrush (Turdus philomelos)",
        population_trends: [
            {term: {start: 1970, end: 2014}, pct_change: -53, annual: -1.7, trend: "weak decline"},
            {term: {start: 2009, end: 2014}, pct_change: -4, annual: -0.88, trend: "no change"},
        ],
    },
];

assert.commandWorked(coll.insert(testDocs));

// Aggregate tests
const pipeline = [
    {$project: {_id: 0}},
    {$unwind: "$population_trends"},
    {$match: {$expr: {$eq: ["$population_trends.trend", "$$target_trend"]}}},
    {$sort: {Species: 1}},
];
let expectedResults = [
    {
        Species: "Bullfinch (Pyrrhula pyrrhula)",
        population_trends: {term: {start: 2009, end: 2014}, pct_change: 12, annual: 2.38, trend: "weak increase"},
    },
];

assert.eq(coll.aggregate(pipeline, {let: {target_trend: "weak increase"}}).toArray(), expectedResults);

expectedResults = [
    {
        Species: "Chaffinch (Fringilla coelebs)",
        population_trends: {term: {start: 2009, end: 2014}, pct_change: -7, annual: -1.49, trend: "weak decline"},
    },
    {
        Species: "Song Thrush (Turdus philomelos)",
        population_trends: {term: {start: 1970, end: 2014}, pct_change: -53, annual: -1.7, trend: "weak decline"},
    },
];
assert.eq(coll.aggregate(pipeline, {let: {target_trend: "weak decline"}}).toArray(), expectedResults);

// Test that running explain on the agg command works as expected.
let explain = assert.commandWorked(
    testDB.runCommand({
        explain: {aggregate: coll.getName(), pipeline, let: {target_trend: "weak decline"}, cursor: {}},
        verbosity: "executionStats",
    }),
);

if (!isMongos) {
    if (checkSbeFullFeatureFlagEnabled(testDB)) {
        // $unwind should be pushed down to SBE.
        assert(planHasStage(testDB, explain, "UNWIND"), explain);
    } else {
        explain = getSingleNodeExplain(explain);
        assert(explain.hasOwnProperty("stages"), explain);
        assert.neq(explain.stages.length, 0, explain);
        let lastStage = explain.stages[explain.stages.length - 1];
        assert.eq(lastStage.nReturned, 2, explain);
    }
}

if (!isMongos && !TestData.testingReplicaSetEndpoint) {
    // Test that if runtimeConstants and let are both specified, both will coexist.
    // Runtime constants are not allowed on mongos passthroughs.
    // Must set 'fromRouter: true' as otherwise 'runtimeConstants' is disallowed on mongod.
    let constants = {
        localNow: new Date(),
        clusterTime: new Timestamp(0, 0),
    };

    assert.eq(
        coll
            .aggregate(pipeline, {
                runtimeConstants: constants,
                let: {target_trend: "weak decline"},
                fromRouter: true,
            })
            .toArray(),
        expectedResults,
    );

    // Test that undefined let params in the pipeline fail gracefully.
    assert.commandFailedWithCode(
        testDB.runCommand({
            aggregate: coll.getName(),
            pipeline: pipeline,
            runtimeConstants: constants,
            cursor: {},
            let: {cat: "not_a_bird"},
            fromRouter: true,
        }),
        17276,
    );
    // Test null and empty let parameters
    const pipeline_no_lets = [
        {$project: {_id: 0}},
        {$unwind: "$population_trends"},
        {$match: {$expr: {$eq: ["$population_trends.trend", "weak decline"]}}},
        {$sort: {Species: 1}},
    ];
    assert.eq(
        coll.aggregate(pipeline_no_lets, {runtimeConstants: constants, let: {}, fromRouter: true}).toArray(),
        expectedResults,
    );

    assert.commandWorked(
        testDB.runCommand({
            aggregate: coll.getName(),
            pipeline: pipeline_no_lets,
            runtimeConstants: constants,
            cursor: {},
            let: null,
            fromRouter: true,
        }),
    );

    assert.commandFailedWithCode(
        testDB.runCommand({
            aggregate: coll.getName(),
            pipeline: pipeline_no_lets,
            runtimeConstants: constants,
            cursor: {},
            let: 1,
            fromRouter: true,
        }),
        ErrorCodes.TypeMismatch,
    );
}

// Test that $project stage can use 'let' variables
assert.eq(
    testDB.runCommand({
        aggregate: coll.getName(),
        pipeline: [
            {
                $project: {
                    "var": {
                        $let: {
                            vars: {variable: "INNER"},
                            "in": {
                                $cond: {
                                    "if": {$eq: [{$substr: ["$Species", 0, 1]}, "B"]},
                                    then: "$$variable",
                                    "else": "---",
                                },
                            },
                        },
                    },
                },
            },
            {$match: {$expr: {$eq: ["$var", "INNER"]}}},
        ],
        cursor: {},
    }).cursor.firstBatch.length,
    2,
);

// Test that $project stage can access command-level 'let' variables.
assert.eq(
    testDB.runCommand({
        aggregate: coll.getName(),
        pipeline: [
            {
                $project: {
                    "var": {
                        $cond: {
                            "if": {$eq: [{$substr: ["$Species", 0, 1]}, "B"]},
                            then: "$$variable",
                            "else": "---",
                        },
                    },
                },
            },
            {$match: {$expr: {$eq: ["$var", "OUTER"]}}},
        ],
        cursor: {},
        "let": {variable: "OUTER"},
    }).cursor.firstBatch.length,
    2,
);

// Test that $project stage can use stage-level and command-level 'let' variables in same command.
assert.eq(
    testDB.runCommand({
        aggregate: coll.getName(),
        pipeline: [
            {
                $project: {
                    "var": {
                        $let: {
                            vars: {innerVar: "INNER"},
                            "in": {
                                $cond: {
                                    "if": {$eq: [{$substr: ["$Species", 0, 1]}, "B"]},
                                    then: {$concat: ["$$innerVar", "$$outerVar"]},
                                    "else": "---",
                                },
                            },
                        },
                    },
                },
            },
            {$match: {$expr: {$eq: ["$var", "INNEROUTER"]}}},
        ],
        cursor: {},
        "let": {outerVar: "OUTER"},
    }).cursor.firstBatch.length,
    2,
);

// Test that $project stage follows variable scoping rules with stage-level and command-level 'let'
// variables.
assert.eq(
    testDB.runCommand({
        aggregate: coll.getName(),
        pipeline: [
            {
                $project: {
                    "var": {
                        $let: {
                            vars: {variable: "INNER"},
                            "in": {
                                $cond: {
                                    "if": {$eq: [{$substr: ["$Species", 0, 1]}, "B"]},
                                    then: "$$variable",
                                    "else": "---",
                                },
                            },
                        },
                    },
                },
            },
            {$match: {$expr: {$eq: ["$var", "INNER"]}}},
        ],
        cursor: {},
        "let": {variable: "OUTER"},
    }).cursor.firstBatch.length,
    2,
);

// Test that the find command works correctly with a let parameter argument.
let result = assert.commandWorked(
    testDB.runCommand({
        find: coll.getName(),
        let: {target_species: "Song Thrush (Turdus philomelos)"},
        filter: {$expr: {$eq: ["$Species", "$$target_species"]}},
        projection: {_id: 0},
    }),
).cursor.firstBatch;
expectedResults = {
    Species: "Song Thrush (Turdus philomelos)",
    population_trends: [
        {term: {start: 1970, end: 2014}, pct_change: -53, annual: -1.7, trend: "weak decline"},
        {term: {start: 2009, end: 2014}, pct_change: -4, annual: -0.88, trend: "no change"},
    ],
};
assert.eq(result.length, 1);
assert.eq(expectedResults, result[0]);

// Test that let parameters work as expected when the find is run as an explain.
explain = assert.commandWorked(
    testDB.runCommand({
        explain: {
            find: coll.getName(),
            let: {target_species: "Song Thrush (Turdus philomelos)"},
            filter: {$expr: {$eq: ["$Species", "$$target_species"]}},
            projection: {_id: 0},
        },
        verbosity: "executionStats",
    }),
);
if (!isMongos) {
    assert.eq(explain.executionStats.nReturned, 1, explain);
}

// Delete tests with let params will delete a record, assert that a point-wise find yields an
// empty result, and then restore the collection state for further tests down the line. We can't
// exercise a multi-delete here (limit: 0) because of failures in sharded txn passthrough tests.
assert.commandWorked(
    testDB.runCommand({
        delete: coll.getName(),
        let: {target_species: "Song Thrush (Turdus philomelos)"},
        deletes: [{q: {$and: [{_id: 4}, {$expr: {$eq: ["$Species", "$$target_species"]}}]}, limit: 1}],
    }),
);

result = assert.commandWorked(testDB.runCommand({find: coll.getName(), filter: {$expr: {$eq: ["$_id", "4"]}}})).cursor
    .firstBatch;
assert.eq(result.length, 0);

assert.commandWorked(coll.insert({_id: 4, Species: "bird_to_remove"}));

// Test that explain of a delete command works as expected with 'let' parameters.
explain = assert.commandWorked(
    testDB.runCommand({
        explain: {
            delete: coll.getName(),
            let: {target_species: "bird_to_remove"},
            deletes: [{q: {$and: [{_id: 4}, {$expr: {$eq: ["$Species", "$$target_species"]}}]}, limit: 1}],
        },
        verbosity: "executionStats",
    }),
);
if (!isMongos) {
    let deleteStage = getPlanStage(explain.executionStats.executionStages, "DELETE");
    assert.eq(deleteStage.nWouldDelete, 1, explain);
}

// Test that the .remove() shell helper supports let parameters.
result = assert.commandWorked(
    coll.remove(
        {$and: [{_id: 4}, {$expr: {$eq: ["$Species", "$$target_species"]}}]},
        {justOne: true, let: {target_species: "bird_to_remove"}},
    ),
);
assert.eq(result.nRemoved, 1);

result = assert.commandWorked(testDB.runCommand({find: coll.getName(), filter: {$expr: {$eq: ["$_id", "4"]}}})).cursor
    .firstBatch;
assert.eq(result.length, 0);

// Test that reserved names are not allowed as let variable names.
[
    {name: "Reserved", errorCode: ErrorCodes.FailedToParse},
    {name: "REMOVE", errorCode: ErrorCodes.FailedToParse},
    {name: "ROOT", errorCode: ErrorCodes.FailedToParse},
    {name: "NOW", errorCode: 4738901},
    {name: "CLUSTER_TIME", errorCode: 4738901},
    {name: "IS_MR", errorCode: 4738901},
    {name: "JS_SCOPE", errorCode: 4738901},
].forEach((caseInfo) => {
    const reservedName = caseInfo.name;
    const expectedError = caseInfo.errorCode;
    assert.commandFailedWithCode(
        testDB.runCommand({aggregate: coll.getName(), pipeline: [], cursor: {}, let: {[reservedName]: "failure"}}),
        expectedError,
        `Expected an aggregate with the variable ${reservedName} to fail.`,
    );
    assert.commandFailedWithCode(
        testDB.runCommand({update: coll.getName(), updates: [{q: {}, u: []}], let: {[reservedName]: "failure"}}),
        expectedError,
        `Expected a pipeline style update with the variable ${reservedName} to fail.`,
    );
});

// Test that let variables can be used within views.
assert.commandWorked(
    testDB.runCommand({
        create: "core-viewColl",
        viewOn: coll.getName(),
        pipeline: [{$match: {Species: "Song Thrush (Turdus philomelos)"}}],
    }),
);
assert.commandWorked(
    testDB.runCommand({
        aggregate: "core-viewColl",
        pipeline: [{$addFields: {var: "$$variable"}}],
        let: {variable: "Song Thrush"},
        cursor: {},
    }),
);

assert.commandWorked(coll.insert({_id: 5, Species: "spy_bird"}));

// Test that explain of findAndModify works correctly with let parameters.
explain = assert.commandWorked(
    testDB.runCommand({
        explain: {
            findAndModify: coll.getName(),
            let: {target_species: "spy_bird"},
            // Querying on _id field for sharded collection passthroughs.
            query: {$and: [{_id: 5}, {$expr: {$eq: ["$Species", "$$target_species"]}}]},
            update: {Species: "questionable_bird"},
            new: true,
        },
        verbosity: "executionStats",
    }),
);
if (!isMongos) {
    let updateStage = getPlanStage(explain.executionStats.executionStages, "UPDATE");
    assert.eq(updateStage.nMatched, 1, explain);
    assert.eq(updateStage.nWouldModify, 1, explain);
}

// Test that findAndModify works correctly with let parameter arguments.
result = assert.commandWorked(
    testDB.runCommand({
        findAndModify: coll.getName(),
        let: {target_species: "spy_bird"},
        // Querying on _id field for sharded collection passthroughs.
        query: {$and: [{_id: 5}, {$expr: {$eq: ["$Species", "$$target_species"]}}]},
        update: {Species: "questionable_bird"},
        new: true,
    }),
);
expectedResults = {
    _id: 5,
    Species: "questionable_bird",
};
assert.eq(expectedResults, result.value, result);

result = assert.commandWorked(
    testDB.runCommand({
        findAndModify: coll.getName(),
        let: {species_name: "not_a_bird", realSpecies: "dino"},
        // Querying on _id field for sharded collection passthroughs.
        query: {$and: [{_id: 5}, {$expr: {$eq: ["$Species", "questionable_bird"]}}]},
        update: [{$project: {Species: "$$species_name"}}, {$addFields: {suspect: "$$realSpecies"}}],
        new: true,
    }),
);
expectedResults = {
    _id: 5,
    Species: "not_a_bird",
    suspect: "dino",
};
assert.eq(expectedResults, result.value, result);

// Test that explain of update works correctly with let parameters.
explain = assert.commandWorked(
    testDB.runCommand({
        explain: {
            update: coll.getName(),
            updates: [
                {
                    q: {_id: 3, $expr: {$eq: ["$Species", "$$target_species"]}},
                    u: [{$set: {Species: "$$new_name"}}],
                },
            ],
            let: {target_species: "Chaffinch (Fringilla coelebs)", new_name: "Chaffinch"},
        },
        verbosity: "executionStats",
    }),
);
if (!isMongos) {
    let updateStage = getPlanStage(explain.executionStats.executionStages, "UPDATE");
    assert.eq(updateStage.nMatched, 1, explain);
    assert.eq(updateStage.nWouldModify, 1, explain);
}

// Test that update respects different parameters in both the query and update part.
result = assert.commandWorked(
    testDB.runCommand({
        update: coll.getName(),
        updates: [
            {
                q: {_id: 3, $expr: {$eq: ["$Species", "$$target_species"]}},
                u: [{$set: {Species: "$$new_name"}}],
            },
        ],
        let: {target_species: "Chaffinch (Fringilla coelebs)", new_name: "Chaffinch"},
    }),
);
assert.eq(result.n, 1);
assert.eq(result.nModified, 1);

result = assert.commandWorked(
    testDB.runCommand({find: coll.getName(), filter: {$expr: {$eq: ["$Species", "Chaffinch (Fringilla coelebs)"]}}}),
);
assert.eq(result.cursor.firstBatch.length, 0);

result = assert.commandWorked(
    testDB.runCommand({find: coll.getName(), filter: {$expr: {$eq: ["$Species", "Chaffinch"]}}}),
);
assert.eq(result.cursor.firstBatch.length, 1);

// Test that update respects runtime constants and parameters.
result = assert.commandWorked(
    testDB.runCommand({
        update: coll.getName(),
        updates: [
            {
                q: {_id: 3, $expr: {$eq: ["$Species", "$$target_species"]}},
                u: [{$set: {Timestamp: "$$NOW"}}, {$set: {Species: "$$new_name"}}],
            },
        ],
        let: {target_species: "Chaffinch", new_name: "Pied Piper"},
    }),
);
assert.eq(result.n, 1);
assert.eq(result.nModified, 1);

result = assert.commandWorked(
    testDB.runCommand({find: coll.getName(), filter: {$expr: {$eq: ["$Species", "Chaffinch"]}}}),
);
assert.eq(result.cursor.firstBatch.length, 0, result);

result = assert.commandWorked(
    testDB.runCommand({find: coll.getName(), filter: {$expr: {$eq: ["$Species", "Pied Piper"]}}}),
);
assert.eq(result.cursor.firstBatch.length, 1, result);

// This forces a multi-statement transaction to commit if this test is running in one of the
// multi-statement transaction passthrough suites. We need to do this to ensure the updates
// above commit before running an update that will fail, as the failed update aborts the entire
// transaction and rolls back the updates above.
assert.commandWorked(testDB.runCommand({ping: 1}));

// Test that undefined let params in the update's query part fail gracefully.
assert.commandFailedWithCode(
    testDB.runCommand({
        update: coll.getName(),
        updates: [
            {
                q: {$expr: {$eq: ["$Species", "$$target_species"]}},
                u: [{$set: {Species: "Homo Erectus"}}],
            },
        ],
        let: {cat: "not_a_bird"},
    }),
    17276,
);

// Test that undefined let params in the update's update part fail gracefully.
assert.commandFailedWithCode(
    testDB.runCommand({
        update: coll.getName(),
        updates: [
            {
                q: {_id: 3, $expr: {$eq: ["$Species", "Chaffinch (Fringilla coelebs)"]}},
                u: [{$set: {Species: "$$new_name"}}],
            },
        ],
        let: {cat: "not_a_bird"},
    }),
    17276,
);

// Test that the .update() shell helper supports let parameters.
result = assert.commandWorked(
    coll.update({_id: 3, $expr: {$eq: ["$Species", "$$target_species"]}}, [{$set: {Species: "$$new_name"}}], {
        let: {target_species: "Pied Piper", new_name: "Chaffinch"},
    }),
);
assert.eq(result.nMatched, 1);
assert.eq(result.nModified, 1);

result = assert.commandWorked(
    testDB.runCommand({find: coll.getName(), filter: {$expr: {$eq: ["$Species", "Pied Piper"]}}}),
);
assert.eq(result.cursor.firstBatch.length, 0, result);

result = assert.commandWorked(
    testDB.runCommand({find: coll.getName(), filter: {$expr: {$eq: ["$Species", "Chaffinch"]}}}),
);
assert.eq(result.cursor.firstBatch.length, 1, result);

// Test that let variables can be initialized with an expression.
result = assert.commandWorked(
    testDB.runCommand({
        find: coll.getName(),
        filter: {},
        projection: {_id: "$$a"},
        limit: 1,
        let: {a: {$add: [2, 3]}},
    }),
).cursor.firstBatch;
assert.eq(result, [{_id: 5}]);

// Test that the expression cannot refer to any fields or variables.
assert.commandFailedWithCode(
    testDB.runCommand({
        find: coll.getName(),
        filter: {},
        let: {a: "$_id"},
    }),
    4890500,
);
assert.commandFailedWithCode(
    testDB.runCommand({
        find: coll.getName(),
        filter: {},
        let: {a: "$$oops_undefined_var"},
    }),
    17276,
);

// Test that each expression can refer to previously defined let variables.
result = assert.commandWorked(
    testDB.runCommand({
        find: coll.getName(),
        filter: {},
        projection: {_id: "$$c"},
        limit: 1,
        let: {
            a: 2,
            b: {$add: ["$$a", 3]},
            c: {$multiply: ["$$a", "$$b"]},
        },
    }),
).cursor.firstBatch;
assert.eq([{_id: (2 + 3) * 2}], result);

// Test that $rand is allowed in the initializer.
result = assert.commandWorked(
    testDB.runCommand({
        find: coll.getName(),
        filter: {},
        projection: {_id: "$$r"},
        limit: 1,
        let: {r: {$rand: {}}},
    }),
).cursor.firstBatch[0]._id;
assert.between(0, result, 1);

// Test that each initializer expression is evaluated separately.
{
    const values = assert.commandWorked(
        testDB.runCommand({
            find: coll.getName(),
            filter: {},
            projection: {_id: ["$$a", "$$b", "$$c", "$$d"]},
            limit: 1,
            let: {
                a: {$rand: {}},
                b: {$rand: {}},
                c: {$rand: {}},
                d: {$rand: {}},
            },
        }),
    ).cursor.firstBatch[0]._id;

    const deduped = [...new Set(values)];
    assert.eq(values.length, deduped.length, `Expected all distinct values: ${values}`);
}

// Test that the expressions are evaluated once up front.
const values = assert
    .commandWorked(
        testDB.runCommand({
            find: coll.getName(),
            filter: {},
            projection: {_id: "$$r"},
            let: {r: {$rand: {}}},
        }),
    )
    .cursor.firstBatch.map((doc) => doc._id);
assert.gt(values.length, 1);

const deduped = [...new Set(values)];
assert.eq(1, deduped.length, `Expected all identical values: ${deduped}`);

// Test that expressions wrapped with $literal are serialized correctly when run in sharded cluster
// environments.
result = coll
    .aggregate([{$match: {$expr: {$eq: ["$_id", 2]}}}, {$project: {a: "$$b"}}], {
        let: {b: {$literal: "$notAFieldPath"}},
    })
    .toArray();
assert.eq(result, [{_id: 2, a: "$notAFieldPath"}]);