/**
 * Tests the behavior of a $lookup when a shard contains incorrect routing information for the
 * local and/or foreign collections. This includes when the shard thinks the collection is sharded
 * when it's not, and likewise when it thinks the collection is unsharded but is actually sharded.
 *
 * We restart a mongod to cause it to forget that a collection was sharded. When restarted, we
 * expect it to still have all the previous data.
 *
 * @tags: [
 *   requires_persistence,
 *    # TODO (SERVER-88125): Re-enable this test or add an explanation why it is incompatible.
 *    embedded_router_incompatible,
 * ]
 *
 */
import {ShardingTest} from "jstests/libs/shardingtest.js";

// Restarts the primary shard and ensures that it believes both collections are unsharded.
function restartPrimaryShard(rs, ...expectedCollections) {
    // Returns true if the shard is aware that the collection is sharded.
    function hasRoutingInfoForNs(shardConn, coll) {
        const res = shardConn.adminCommand({getShardVersion: coll, fullMetadata: true});
        assert.commandWorked(res);
        return res.metadata.collVersion != undefined;
    }

    rs.restart(0);
    rs.awaitSecondaryNodes();

    expectedCollections.forEach(function(coll) {
        assert(!hasRoutingInfoForNs(rs.getPrimary(), coll.getFullName()),
               'Shard role not cleared for ' + coll.getFullName());
    });
}

// Disable checking for index consistency to ensure that the config server doesn't trigger a
// StaleShardVersion exception on shard0 and cause it to refresh its sharding metadata.
const nodeOptions = {
    setParameter: {enableShardedIndexConsistencyCheck: false}
};

const testName = "lookup_stale_mongod";
const st =
    new ShardingTest({shards: 2, mongos: 2, rs: {nodes: 1}, other: {configOptions: nodeOptions}});

const mongos0DB = st.s0.getDB(testName);
const mongos0LocalColl = mongos0DB[testName + "_local"];
const mongos0ForeignColl = mongos0DB[testName + "_foreign"];

const mongos1DB = st.s1.getDB(testName);
const mongos1LocalColl = mongos1DB[testName + "_local"];
const mongos1ForeignColl = mongos1DB[testName + "_foreign"];

const pipeline = [
    {$lookup: {localField: "a", foreignField: "b", from: mongos0ForeignColl.getName(), as: "same"}},
    // Unwind the results of the $lookup, so we can sort by them to get a consistent ordering
    // for the query results.
    {$unwind: "$same"},
    {$sort: {_id: 1, "same._id": 1}}
];

// The results are expected to be correct if the $lookup stage is executed on the mongos which
// is aware that the collection is sharded.
const expectedResults = [
    {_id: 0, a: 1, "same": {_id: 0, b: 1}},
    {_id: 1, a: null, "same": {_id: 1, b: null}},
    {_id: 1, a: null, "same": {_id: 2}},
    {_id: 2, "same": {_id: 1, b: null}},
    {_id: 2, "same": {_id: 2}}
];

// Ensure that shard0 is the primary shard.
assert.commandWorked(mongos0DB.adminCommand(
    {enableSharding: mongos0DB.getName(), primaryShard: st.shard0.shardName}));

assert.commandWorked(mongos0LocalColl.insert({_id: 0, a: 1}));
assert.commandWorked(mongos0LocalColl.insert({_id: 1, a: null}));

assert.commandWorked(mongos0ForeignColl.insert({_id: 0, b: 1}));
assert.commandWorked(mongos0ForeignColl.insert({_id: 1, b: null}));

// Send writes through mongos1 such that it's aware of the collections and believes they are
// unsharded.
assert.commandWorked(mongos1LocalColl.insert({_id: 2}));
assert.commandWorked(mongos1ForeignColl.insert({_id: 2}));

//
// Test unsharded local and sharded foreign collections, with the primary shard unaware that
// the foreign collection is sharded.
//

// Shard the foreign collection.
assert.commandWorked(
    mongos0DB.adminCommand({shardCollection: mongos0ForeignColl.getFullName(), key: {_id: 1}}));

// Split the collection into 2 chunks: [MinKey, 1), [1, MaxKey).
assert.commandWorked(
    mongos0DB.adminCommand({split: mongos0ForeignColl.getFullName(), middle: {_id: 1}}));

// Move the [minKey, 1) chunk to shard1.
assert.commandWorked(mongos0DB.adminCommand({
    moveChunk: mongos0ForeignColl.getFullName(),
    find: {_id: 0},
    to: st.shard1.shardName,
    _waitForDelete: true
}));

// Verify $lookup results through the fresh mongos.
restartPrimaryShard(st.rs0, mongos0LocalColl, mongos0ForeignColl);
assert.eq(mongos0LocalColl.aggregate(pipeline).toArray(), expectedResults);

// Verify $lookup results through mongos1, which is not aware that the foreign collection is
// sharded. In this case the results will be correct since the entire pipeline will be run on a
// shard, which will do a refresh before executing the foreign pipeline.
restartPrimaryShard(st.rs0, mongos0LocalColl, mongos0ForeignColl);
assert.eq(mongos1LocalColl.aggregate(pipeline).toArray(), expectedResults);

//
// Test sharded local and sharded foreign collections, with the primary shard unaware that
// either collection is sharded.
//

// Shard the local collection.
assert.commandWorked(
    mongos0DB.adminCommand({shardCollection: mongos0LocalColl.getFullName(), key: {_id: 1}}));

// Split the collection into 2 chunks: [MinKey, 1), [1, MaxKey).
assert.commandWorked(
    mongos0DB.adminCommand({split: mongos0LocalColl.getFullName(), middle: {_id: 1}}));

// Move the [minKey, 1) chunk to shard1.
assert.commandWorked(mongos0DB.adminCommand({
    moveChunk: mongos0LocalColl.getFullName(),
    find: {_id: 0},
    to: st.shard1.shardName,
    _waitForDelete: true
}));

// Verify $lookup results through the fresh mongos.
restartPrimaryShard(st.rs0, mongos0LocalColl, mongos0ForeignColl);
assert.eq(mongos0LocalColl.aggregate(pipeline).toArray(), expectedResults);

// Verify $lookup results through the stale mongos.
restartPrimaryShard(st.rs0, mongos0LocalColl, mongos0ForeignColl);
assert.eq(mongos1LocalColl.aggregate(pipeline).toArray(), expectedResults);

//
// Test sharded local and unsharded foreign collections, with the primary shard unaware that
// the local collection is sharded.
//

// Recreate the foreign collection as unsharded.
mongos0ForeignColl.drop();
assert.commandWorked(mongos0ForeignColl.insert({_id: 0, b: 1}));
assert.commandWorked(mongos0ForeignColl.insert({_id: 1, b: null}));
assert.commandWorked(mongos0ForeignColl.insert({_id: 2}));

// Verify $lookup results through the fresh mongos.
restartPrimaryShard(st.rs0, mongos0LocalColl, mongos0ForeignColl);
assert.eq(mongos0LocalColl.aggregate(pipeline).toArray(), expectedResults);

// Verify $lookup results through the stale mongos.
restartPrimaryShard(st.rs0, mongos0LocalColl, mongos0ForeignColl);
assert.eq(mongos1LocalColl.aggregate(pipeline).toArray(), expectedResults);

//
// Test two-level $lookup with a stale shard handles the shard role recovery case.
//
jsTest.log("Running two-level $lookup with a shard that needs recovery");

assert.commandWorked(st.s0.adminCommand({enableSharding: 'D', primaryShard: st.shard0.shardName}));

const D = st.s0.getDB('D');

assert.commandWorked(D.A.insert({Key: 1, Value: 1}));
assert.commandWorked(D.B.insert({Key: 1, Value: 1}));
assert.commandWorked(D.C.insert({Key: 1, Value: 1}));
assert.commandWorked(D.D.insert({Key: 1, Value: 1}));

const aggPipeline = [{
    $lookup: {
        from: 'B',
        localField: 'Key',
        foreignField: 'Value',
        as: 'Joined',
        pipeline: [
            {
                $lookup: {
                    from: 'C',
                    localField: 'Key',
                    foreignField: 'Value',
                    as: 'Joined',
                }
            },
            {
                $lookup: {
                    from: 'D',
                    localField: 'Key',
                    foreignField: 'Value',
                    as: 'Joined',
                }
            },
        ],
    }
}];

const resultBefore = D.A.aggregate(aggPipeline).toArray();

// Restarting the shard primary in order for the shard role's cache to be cleared
restartPrimaryShard(st.rs0, D.A, D.B, D.C, D.D);

const resultAfter = D.A.aggregate(aggPipeline).toArray();
assert.eq(resultBefore, resultAfter, "Before and after results do not match");

st.stop();
