SERVER-31897 Fix exhaust queries in the shell

(cherry picked from commit 502a72fb94bf05d2b6a24d6d2975a9ef14a5abef)

2 files changed, 29 insertions, 1 deletions

diff --git a/jstests/core/exhaust.js b/jstests/core/exhaust.js
new file mode 100644
index 00000000000..cef3b32fcd2
--- /dev/null
+++ b/jstests/core/exhaust.js
@@ -0,0 +1,24 @@
+(function() {
+    'use strict';
+
+    var c = db.exhaustColl;
+    c.drop();
+
+    const docCount = 4;
+    for (var i = 0; i < docCount; i++) {
+        assert.writeOK(c.insert({a: i}));
+    }
+
+    // Check that the query works without exhaust set
+    assert.eq(c.find().batchSize(1).itcount(), docCount);
+
+    // Now try to run the same query with exhaust
+    try {
+        assert.eq(c.find().batchSize(1).addOption(DBQuery.Option.exhaust).itcount(), docCount);
+    } catch (e) {
+        // The exhaust option is not valid against mongos, ensure that this query throws the right
+        // code
+        assert.eq(e.code, 18526);
+    }
+
+}());
diff --git a/src/mongo/client/dbclientcursor.cpp b/src/mongo/client/dbclientcursor.cpp
index 0a439f610fb..6de2fc26500 100644
--- a/src/mongo/client/dbclientcursor.cpp
+++ b/src/mongo/client/dbclientcursor.cpp
@@ -204,6 +204,10 @@ bool DBClientCursor::initLazyFinish(bool& retry) {
 }
 
 void DBClientCursor::requestMore() {
+    if (opts & QueryOption_Exhaust) {
+        return exhaustReceiveMore();
+    }
+
     invariant(!_connectionHasPendingReplies);
     verify(cursorId && batch.pos == batch.objs.size());
 
@@ -238,7 +242,7 @@ void DBClientCursor::requestMore() {
 /** with QueryOption_Exhaust, the server just blasts data at us (marked at end with cursorid==0). */
 void DBClientCursor::exhaustReceiveMore() {
     verify(cursorId && batch.pos == batch.objs.size());
-    verify(!haveLimit);
+    uassert(50657, "Cannot have limit for exhaust query", !haveLimit);
     Message response;
     verify(_client);
     if (!_client->recv(response, _lastRequestId)) {