SERVER-57462 Change the shell to use OP_MSG for exhaust queries instead of OP_QUERY

3 files changed, 76 insertions, 14 deletions

diff --git a/jstests/noPassthrough/exhaust_compression.js b/jstests/noPassthrough/exhaust_compression.js
index 3ede60b4aa3..f0dc57666b0 100644
--- a/jstests/noPassthrough/exhaust_compression.js
+++ b/jstests/noPassthrough/exhaust_compression.js
@@ -14,15 +14,24 @@ var runTest = function(compressor) {
             assert.commandWorked(db.runCommand({insert: collName, documents: [{a: i}]}));
         }
 
+        const kBatchSize = 2;
         const preRes =
             assert.commandWorked(db.adminCommand({configureFailPoint: fp, mode: "alwaysOn"}));
 
-        db.exhaustCollection.find({}).batchSize(2).addOption(DBQuery.Option.exhaust).toArray();
+        db.exhaustCollection.find({})
+            .batchSize(kBatchSize)
+            .addOption(DBQuery.Option.exhaust)
+            .toArray();
 
         const postRes =
             assert.commandWorked(db.adminCommand({configureFailPoint: fp, mode: "off"}));
 
-        assert.eq(preRes.count + 1, postRes.count, "Exhaust messages are not compressed");
+        // The initial response for find command has kBatchSize docs and the remaining docs comes
+        // in batches of kBatchSize in response to the getMore command with the exhaustAllowed bit
+        // set.
+        const kExpectedDelta = Math.floor((kDocumentCount - kBatchSize) / kBatchSize);
+        assert.eq(
+            postRes.count - preRes.count, kExpectedDelta, "Exhaust messages are not compressed");
     }, mongo.port, false, "--networkMessageCompressors", compressor);
 
     shell();
diff --git a/jstests/noPassthrough/exhaust_cursors.js b/jstests/noPassthrough/exhaust_cursors.js
new file mode 100644
index 00000000000..e1cb7ee9a3e
--- /dev/null
+++ b/jstests/noPassthrough/exhaust_cursors.js
@@ -0,0 +1,65 @@
+/**
+ * Verifies mongos and mongod's behavior for exhaust queries.
+ *
+ * 'ShardingTest' requires replication.
+ * @tags: [requires_replication]
+ */
+(function() {
+"use strict";
+
+const docs = [{a: 1}, {a: 2}, {a: 3}, {a: 4}, {a: 5}, {a: 6}, {a: 7}, {a: 8}];
+const kBatchSize = 3;
+const kNoOfDocs = docs.length;
+
+[{
+    "setUp": () => {
+        const conn = MongoRunner.runMongod();
+        const db = conn.getDB(jsTestName());
+        return {"env": conn, "db": db};
+    },
+    "tearDown": (env) => MongoRunner.stopMongod(env),
+    "verifyThis": (cursor, docIdx, doc) => {
+        // Because the first batch is returned from a find command without exhaustAllowed bit,
+        // moreToCome bit is not set in reply message.
+        const isFirstBatch = docIdx < kBatchSize;
+
+        // The last batch which does not contain the full batch size is returned without moreToCome
+        // bit set.
+        const isLastBatch = docIdx >= kNoOfDocs - (kNoOfDocs % kBatchSize);
+
+        if (isFirstBatch || isLastBatch) {
+            assert(!cursor._hasMoreToCome(), `${docIdx} doc: ${doc}`);
+        } else {
+            assert(cursor._hasMoreToCome(), `${docIdx} doc: ${doc}`);
+        }
+    }
+},
+ {
+     "setUp": () => {
+         const st = new ShardingTest({shards: 1, config: 1});
+         const db = st.s0.getDB(jsTestName());
+         return {"env": st, "db": db};
+     },
+     "tearDown": (env) => env.stop(),
+     "verifyThis": (cursor, docIdx, doc) => {
+         // Mongos does not support exhaust queries, not by returning an error but by sending reply
+         // without moreToCome bit set. So, _hasMoreToCome() is always false.
+         assert(!cursor._hasMoreToCome(), `${docIdx} doc: ${doc}`);
+     }
+ }].forEach(({setUp, tearDown, verifyThis}) => {
+    const {env, db} = setUp();
+
+    db.coll.drop();
+
+    assert.commandWorked(db.coll.insert(docs));
+
+    let cursor = db.coll.find().batchSize(kBatchSize).addOption(DBQuery.Option.exhaust);
+    let docIdx = 0;
+    cursor.forEach(doc => {
+        verifyThis(cursor, docIdx, doc);
+        ++docIdx;
+    });
+
+    tearDown(env);
+});
+}());
diff --git a/jstests/sharding/query_sharded.js b/jstests/sharding/query_sharded.js
index b79cff55cb1..a043f257d6c 100644
--- a/jstests/sharding/query_sharded.js
+++ b/jstests/sharding/query_sharded.js
@@ -9,18 +9,6 @@ var coll = mongos.getCollection("foo.bar");
 
 //
 //
-// Ensure we can't use exhaust option through mongos
-coll.remove({});
-assert.commandWorked(coll.insert({a: 'b'}));
-var query = coll.find({});
-assert.neq(null, query.next());
-query = coll.find({}).addOption(DBQuery.Option.exhaust);
-assert.throws(function() {
-    query.next();
-});
-
-//
-//
 // Ensure we can't trick mongos by inserting exhaust option on a command through mongos
 coll.remove({});
 assert.commandWorked(coll.insert({a: 'b'}));