Add improved Bulk Write API for Java Reactive Driver

driver-core/src/main/com/mongodb/internal/async/AsyncBatchCursor.java

@@ -19,8 +19,11 @@
 import com.mongodb.internal.operation.BatchCursor;
 
 import java.io.Closeable;
+import java.util.ArrayList;
 import java.util.List;
 
+import static com.mongodb.internal.async.AsyncRunnable.beginAsync;
+
 /**
  * MongoDB returns query results as batches, and this interface provides an asynchronous iterator over those batches.  The first call to
  * the {@code next} method will return the first batch, and subsequent calls will trigger an asynchronous request to get the next batch
@@ -72,4 +75,22 @@ public interface AsyncBatchCursor<T> extends Closeable {
      */
     @Override
     void close();
+
+    default void exhaustCursor(final SingleResultCallback<List<List<T>>> finalCallback) {
+        List<List<T>> results = new ArrayList<>();
+
+        beginAsync().thenRunDoWhileLoop(iterationCallback -> {
+                    beginAsync().
+                            thenSupply(this::next)
+                            .thenConsume((batch, callback) -> {
+                                if (batch != null && !batch.isEmpty()) {
+                                    results.add(batch);
+                                }
+                                callback.complete(callback);
+                            }).finish(iterationCallback);
+                }, () -> !this.isClosed())
+                .<List<List<T>>>thenSupply(callback -> {
+                    callback.complete(results);
+                }).finish(finalCallback);
+     }
 }

driver-core/src/main/com/mongodb/internal/operation/BatchCursor.java

@@ -25,6 +25,12 @@
 import java.util.Iterator;
 import java.util.List;
 
+import static java.util.Spliterator.IMMUTABLE;
+import static java.util.Spliterator.ORDERED;
+import static java.util.Spliterators.spliteratorUnknownSize;
+import static java.util.stream.Collectors.toList;
+import static java.util.stream.StreamSupport.stream;
+
 /**
  * MongoDB returns query results as batches, and this interface provideds an iterator over those batches.  The first call to
  * the {@code next} method will return the first batch, and subsequent calls will trigger a  request to get the next batch
@@ -98,4 +104,9 @@ public interface BatchCursor<T> extends Iterator<List<T>>, Closeable {
     ServerCursor getServerCursor();
 
     ServerAddress getServerAddress();
+
+    default List<List<T>> exhaustCursor() {
+        return stream(spliteratorUnknownSize(this, ORDERED | IMMUTABLE), false)
+                .collect(toList());
+    }
 }

driver-core/src/main/com/mongodb/internal/operation/ClientBulkWriteOperation.java

@@ -132,7 +132,6 @@
 import static com.mongodb.internal.operation.CommandOperationHelper.initialRetryState;
 import static com.mongodb.internal.operation.CommandOperationHelper.shouldAttemptToRetryWriteAndAddRetryableLabel;
 import static com.mongodb.internal.operation.CommandOperationHelper.transformWriteException;
-import static com.mongodb.internal.operation.CursorHelper.exhaustCursorAsync;
 import static com.mongodb.internal.operation.OperationHelper.isRetryableWrite;
 import static com.mongodb.internal.operation.SyncOperationHelper.cursorDocumentToBatchCursor;
 import static com.mongodb.internal.operation.SyncOperationHelper.decorateWriteWithRetries;
@@ -141,12 +140,8 @@
 import static java.util.Collections.emptyMap;
 import static java.util.Collections.singletonList;
 import static java.util.Optional.ofNullable;
-import static java.util.Spliterator.IMMUTABLE;
-import static java.util.Spliterator.ORDERED;
-import static java.util.Spliterators.spliteratorUnknownSize;
 import static java.util.stream.Collectors.toList;
 import static java.util.stream.Collectors.toSet;
-import static java.util.stream.StreamSupport.stream;
 
 /**
  * This class is not part of the public API and may be removed or changed at any time.
@@ -544,7 +539,7 @@ private List<List<BsonDocument>> exhaustBulkWriteCommandOkResponseCursor(
                 options.getComment().orElse(null),
                 connectionSource,
                 connection)) {
-            return stream(spliteratorUnknownSize(cursor, ORDERED | IMMUTABLE), false).collect(toList());
+           return cursor.exhaustCursor();
         }
     }
 
@@ -562,7 +557,7 @@ private void exhaustBulkWriteCommandOkResponseCursorAsync(final AsyncConnectionS
                 connection);
 
         beginAsync().<List<List<BsonDocument>>>thenSupply(callback -> {
-            exhaustCursorAsync(cursor, callback);
+             cursor.exhaustCursor(callback);
         }).thenAlwaysRunAndFinish(() -> {
             if (!cursor.isClosed()) {
                 cursor.close();

driver-core/src/main/com/mongodb/internal/operation/CursorHelper.java

@@ -16,41 +16,16 @@
 
 package com.mongodb.internal.operation;
 
-import com.mongodb.internal.async.AsyncBatchCursor;
-import com.mongodb.internal.async.SingleResultCallback;
 import com.mongodb.lang.Nullable;
 import org.bson.BsonDocument;
 import org.bson.BsonInt32;
 
-import java.util.ArrayList;
-import java.util.List;
-
-import static com.mongodb.internal.async.AsyncRunnable.beginAsync;
-
 final class CursorHelper {
 
     static BsonDocument getCursorDocumentFromBatchSize(@Nullable final Integer batchSize) {
         return batchSize == null ? new BsonDocument() : new BsonDocument("batchSize", new BsonInt32(batchSize));
     }
 
-    public static <T> void exhaustCursorAsync(final AsyncBatchCursor<T> cursor, final SingleResultCallback<List<List<T>>> finalCallback) {
-        List<List<T>> results = new ArrayList<>();
-
-        beginAsync().thenRunDoWhileLoop(iterationCallback -> {
-                    beginAsync().
-                            thenSupply(cursor::next)
-                            .thenConsume((batch, callback) -> {
-                                if (batch != null && !batch.isEmpty()) {
-                                    results.add(batch);
-                                }
-                                callback.complete(callback);
-                            }).finish(iterationCallback);
-                }, () -> !cursor.isClosed())
-                .<List<List<T>>>thenSupply(callback -> {
-                    callback.complete(results);
-                }).finish(finalCallback);
-    }
-
     private CursorHelper() {
     }
 }

driver-core/src/test/functional/com/mongodb/client/test/CollectionHelper.java

@@ -363,6 +363,13 @@ public void deleteOne(final Bson filter) {
                 .execute(getBinding());
     }
 
+    public void deleteMany(final Bson filter) {
+        new MixedBulkWriteOperation(namespace,
+                singletonList(new DeleteRequest(filter.toBsonDocument(Document.class, registry)).multi(true)),
+                true, WriteConcern.ACKNOWLEDGED, false)
+                .execute(getBinding());
+    }
+
     public List<T> find(final Bson filter) {
         return find(filter, null);
     }

driver-core/src/test/functional/com/mongodb/internal/operation/AsyncCommandBatchCursorFunctionalTest.java

@@ -23,6 +23,7 @@
 import com.mongodb.ServerCursor;
 import com.mongodb.client.cursor.TimeoutMode;
 import com.mongodb.client.model.CreateCollectionOptions;
+import com.mongodb.client.model.Filters;
 import com.mongodb.client.model.OperationTest;
 import com.mongodb.internal.binding.AsyncConnectionSource;
 import com.mongodb.internal.connection.AsyncConnection;
@@ -46,7 +47,9 @@
 
 import java.util.ArrayList;
 import java.util.List;
+import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Collectors;
@@ -103,6 +106,95 @@ void cleanup() {
         });
     }
 
+    @Test
+    @DisplayName("should exhaust cursor with multiple batches")
+    void shouldExhaustCursorAsyncWithMultipleBatches() {
+        // given
+        BsonDocument commandResult = executeFindCommand(0, 3); // Fetch in batches of size 3
+        cursor = new AsyncCommandBatchCursor<>(TimeoutMode.CURSOR_LIFETIME, commandResult, 3, 0, DOCUMENT_DECODER,
+                null, connectionSource, connection);
+
+        CompletableFuture<List<List<Document>>> futureResult = new CompletableFuture<>();
+
+       // when
+        cursor.exhaustCursor((result, throwable) -> {
+            if (throwable != null) {
+                futureResult.completeExceptionally(throwable);
+            } else {
+                futureResult.complete(result);
+            }
+        });
+
+        // then
+        assertDoesNotThrow(() -> {
+            List<List<Document>> resultBatches = futureResult.get(5, TimeUnit.SECONDS);
+
+            assertEquals(4, resultBatches.size(), "Expected 4 batches for 10 documents with batch size of 3.");
+
+            int totalDocuments = resultBatches.stream().mapToInt(List::size).sum();
+            assertEquals(10, totalDocuments, "Expected a total of 10 documents.");
+        });
+    }
+
+    @Test
+    @DisplayName("should exhaust cursor with closed cursor")
+    void shouldExhaustCursorAsyncWithClosedCursor() {
+        // given
+        BsonDocument commandResult = executeFindCommand(0, 3);
+        cursor = new AsyncCommandBatchCursor<>(TimeoutMode.CURSOR_LIFETIME, commandResult, 3, 0, DOCUMENT_DECODER,
+                null, connectionSource, connection);
+
+        cursor.close();
+
+        CompletableFuture<List<List<Document>>> futureResult = new CompletableFuture<>();
+
+        // when
+        cursor.exhaustCursor((result, throwable) -> {
+            if (throwable != null) {
+                futureResult.completeExceptionally(throwable);
+            } else {
+                futureResult.complete(result);
+            }
+        });
+
+        //then
+        ExecutionException executionException = assertThrows(ExecutionException.class, () -> {
+            futureResult.get(5, TimeUnit.SECONDS);
+        }, "Expected an exception when operating on a closed cursor.");
+
+        IllegalStateException illegalStateException = (IllegalStateException) executionException.getCause();
+        assertEquals("Cursor has been closed", illegalStateException.getMessage());
+    }
+
+    @Test
+    @DisplayName("should exhaust cursor with empty cursor")
+    void shouldExhaustCursorAsyncWithEmptyCursor() {
+        // given
+        getCollectionHelper().deleteMany(Filters.empty());
+
+        BsonDocument commandResult = executeFindCommand(0, 3); // No documents to fetch
+        cursor = new AsyncCommandBatchCursor<>(TimeoutMode.CURSOR_LIFETIME, commandResult, 3, 0, DOCUMENT_DECODER,
+                null, connectionSource, connection);
+
+        CompletableFuture<List<List<Document>>> futureResult = new CompletableFuture<>();
+
+        // when
+        cursor.exhaustCursor((result, throwable) -> {
+            if (throwable != null) {
+                futureResult.completeExceptionally(throwable);
+            } else {
+                futureResult.complete(result);
+            }
+        });
+
+        // then
+        assertDoesNotThrow(() -> {
+            List<List<Document>> resultBatches = futureResult.get(5, TimeUnit.SECONDS);
+
+            assertTrue(resultBatches.isEmpty(), "Expected no batches for an empty cursor.");
+        });
+    }
+
     @Test
     @DisplayName("server cursor should not be null")
     void theServerCursorShouldNotBeNull() {

driver-core/src/test/functional/com/mongodb/internal/operation/CommandBatchCursorFunctionalTest.java

@@ -22,6 +22,7 @@
 import com.mongodb.ServerCursor;
 import com.mongodb.client.cursor.TimeoutMode;
 import com.mongodb.client.model.CreateCollectionOptions;
+import com.mongodb.client.model.Filters;
 import com.mongodb.client.model.OperationTest;
 import com.mongodb.internal.binding.ConnectionSource;
 import com.mongodb.internal.connection.Connection;
@@ -101,6 +102,55 @@ void cleanup() {
         });
     }
 
+    @Test
+    @DisplayName("should exhaust cursor with multiple batches")
+    void shouldExhaustCursorAsyncWithMultipleBatches() {
+        // given
+        BsonDocument commandResult = executeFindCommand(0, 3); // Fetch in batches of size 3
+        cursor = new CommandBatchCursor<>(TimeoutMode.CURSOR_LIFETIME, commandResult, 3, 0, DOCUMENT_DECODER,
+                null, connectionSource, connection);
+
+        // when
+        List<List<Document>> result = cursor.exhaustCursor();
+
+        // then
+        assertEquals(4, result.size(), "Expected 4 batches for 10 documents with batch size of 3.");
+
+        int totalDocuments = result.stream().mapToInt(List::size).sum();
+        assertEquals(10, totalDocuments, "Expected a total of 10 documents.");
+    }
+
+    @Test
+    @DisplayName("should exhaust cursor with closed cursor")
+    void shouldExhaustCursorAsyncWithClosedCursor() {
+        // given
+        BsonDocument commandResult = executeFindCommand(0, 3);
+        cursor = new CommandBatchCursor<>(TimeoutMode.CURSOR_LIFETIME, commandResult, 3, 0, DOCUMENT_DECODER,
+                null, connectionSource, connection);
+        cursor.close();
+
+        // when & then
+        IllegalStateException illegalStateException = assertThrows(IllegalStateException.class, cursor::exhaustCursor);
+        assertEquals("Cursor has been closed", illegalStateException.getMessage());
+    }
+
+    @Test
+    @DisplayName("should exhaust cursor async with empty cursor")
+    void shouldExhaustCursorAsyncWithEmptyCursor() {
+        // given
+        getCollectionHelper().deleteMany(Filters.empty());
+
+        BsonDocument commandResult = executeFindCommand(0, 3); // No documents to fetch
+        cursor = new CommandBatchCursor<>(TimeoutMode.CURSOR_LIFETIME, commandResult, 3, 0, DOCUMENT_DECODER,
+                null, connectionSource, connection);
+
+        // when
+        List<List<Document>> result = cursor.exhaustCursor();
+
+        // then
+        assertTrue(result.isEmpty(), "Expected no batches for an empty cursor.");
+    }
+
     @Test
     @DisplayName("server cursor should not be null")
     void theServerCursorShouldNotBeNull() {