MongoDB - Batch Size On Aggregate Options & Find Options

extensions/mongodb-client/runtime/src/main/java/io/quarkus/mongodb/AggregateOptions.java

@@ -21,6 +21,7 @@ public class AggregateOptions {
     private boolean bypassDocumentValidation;
     private Collation collation;
     private String comment;
+    private int batchSize;
 
     /**
      * Enables writing to temporary files. A null value indicates that it's unspecified.
@@ -114,6 +115,23 @@ public AggregateOptions hint(Bson hint) {
         return this;
     }
 
+    /**
+     * Sets the number of documents to return per batch.
+     *
+     * <p>
+     * Overrides the {@link org.reactivestreams.Subscription#request(long)} value for setting the batch size, allowing for fine
+     * grained
+     * control over the underlying cursor.
+     * </p>
+     *
+     * @param size the batch size
+     * @return this
+     */
+    public AggregateOptions batchSize(int size) {
+        this.batchSize = size;
+        return this;
+    }
+
     public <T> AggregatePublisher<T> apply(AggregatePublisher<T> stream) {
         AggregatePublisher<T> publisher = stream;
 
@@ -134,6 +152,9 @@ public <T> AggregatePublisher<T> apply(AggregatePublisher<T> stream) {
         if (maxTime > 0) {
             publisher.maxAwaitTime(maxTime, maxTimeUnit);
         }
+        if (batchSize > 0) {
+            publisher.batchSize(batchSize);
+        }
         return publisher;
     }
 }

extensions/mongodb-client/runtime/src/main/java/io/quarkus/mongodb/FindOptions.java

@@ -33,6 +33,7 @@ public class FindOptions {
     private boolean showRecordId;
     private long maxAwaitTime;
     private TimeUnit maxAwaitTimeUnit;
+    private int batchSize;
 
     /**
      * Sets the query filter to apply to the query.
@@ -250,6 +251,23 @@ public FindOptions showRecordId(boolean showRecordId) {
         return this;
     }
 
+    /**
+     * Sets the number of documents to return per batch.
+     *
+     * <p>
+     * Overrides the {@link org.reactivestreams.Subscription#request(long)} value for setting the batch size, allowing for fine
+     * grained
+     * control over the underlying cursor.
+     * </p>
+     *
+     * @param size the batch size
+     * @return this
+     */
+    public FindOptions batchSize(int size) {
+        this.batchSize = size;
+        return this;
+    }
+
     public <T> FindPublisher<T> apply(FindPublisher<T> stream) {
         FindPublisher<T> publisher = stream;
         if (filter != null) {
@@ -306,6 +324,9 @@ public <T> FindPublisher<T> apply(FindPublisher<T> stream) {
         if (showRecordId) {
             publisher = publisher.showRecordId(true);
         }
+        if (batchSize > 0) {
+            publisher.batchSize(batchSize);
+        }
         return publisher;
 
     }