SERVER-25235: singleBatch special case in QueryRequest::asAggregationCommand

When converting a QueryRequest into an aggregation command, drop the
'singleBatch' option if 'limit' is present and set to 1.

2 files changed, 27 insertions, 5 deletions

diff --git a/src/mongo/db/query/query_request.cpp b/src/mongo/db/query/query_request.cpp
index 9d76e502ec3..96129eb0d52 100644
--- a/src/mongo/db/query/query_request.cpp
+++ b/src/mongo/db/query/query_request.cpp
@@ -922,11 +922,6 @@ StatusWith<BSONObj> QueryRequest::asAggregationCommand() const {
         return {ErrorCodes::InvalidPipelineOperator,
                 str::stream() << "Option " << kMaxField << " not supported in aggregation."};
     }
-    if (!_wantMore) {
-        return {ErrorCodes::InvalidPipelineOperator,
-                str::stream() << "Option " << kSingleBatchField
-                              << " not supported in aggregation."};
-    }
     if (_maxScan != 0) {
         return {ErrorCodes::InvalidPipelineOperator,
                 str::stream() << "Option " << kMaxScanField << " not supported in aggregation."};
@@ -979,6 +974,13 @@ StatusWith<BSONObj> QueryRequest::asAggregationCommand() const {
         return {ErrorCodes::BadValue,
                 str::stream() << "Cannot convert to an aggregation if ntoreturn is set."};
     }
+    // The aggregation command normally does not support the 'singleBatch' option, but we make a
+    // special exception if 'limit' is set to 1.
+    if (!_wantMore && _limit.value_or(0) != 1LL) {
+        return {ErrorCodes::InvalidPipelineOperator,
+                str::stream() << "Option " << kSingleBatchField
+                              << " not supported in aggregation."};
+    }
 
     // Now that we've successfully validated this QR, begin building the aggregation command.
     aggregationBuilder.append("aggregate", _nss.coll());
diff --git a/src/mongo/db/query/query_request_test.cpp b/src/mongo/db/query/query_request_test.cpp
index 384c723a5fd..6fa047efdca 100644
--- a/src/mongo/db/query/query_request_test.cpp
+++ b/src/mongo/db/query/query_request_test.cpp
@@ -1093,5 +1093,25 @@ TEST(QueryRequestTest, ParseMaxTimeMSPositiveInRangeSucceeds) {
     ASSERT_EQ(maxTime.getValue(), 300);
 }
 
+TEST(QueryRequestTest, ConvertToAggregationFailsWithNoWantMore) {
+    QueryRequest qr(testns);
+    qr.setWantMore(false);
+    ASSERT_NOT_OK(qr.asAggregationCommand());
+}
+
+TEST(QueryRequestTest, ConvertToAggregationFailsWithNoWantMoreAndLimit) {
+    QueryRequest qr(testns);
+    qr.setWantMore(false);
+    qr.setLimit(7);
+    ASSERT_NOT_OK(qr.asAggregationCommand());
+}
+
+TEST(QueryRequestTest, ConvertToAggregationSucceedsWithNoWantMoreLimitOne) {
+    QueryRequest qr(testns);
+    qr.setWantMore(false);
+    qr.setLimit(1);
+    ASSERT_OK(qr.asAggregationCommand());
+}
+
 }  // namespace mongo
 }  // namespace