Propagate the batch size in ChangeStreamPublisherImpl.withDocumentClass (#664)

Propagate the batch size in `ChangeStreamPublisherImpl.withDocumentClass`

JAVA-4010

Co-authored-by: Ross Lawley <[email protected]>

Author: stIncMale

driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/internal/BatchCursorPublisher.java

@@ -38,8 +38,14 @@ abstract class BatchCursorPublisher<T> implements Publisher<T> {
     private Integer batchSize;
 
     BatchCursorPublisher(@Nullable final ClientSession clientSession, final MongoOperationPublisher<T> mongoOperationPublisher) {
+        this(clientSession, mongoOperationPublisher, null);
+    }
+
+    BatchCursorPublisher(@Nullable final ClientSession clientSession, final MongoOperationPublisher<T> mongoOperationPublisher,
+                         @Nullable final Integer batchSize) {
         this.clientSession = clientSession;
         this.mongoOperationPublisher = notNull("mongoOperationPublisher", mongoOperationPublisher);
+        this.batchSize = batchSize;
     }
 
     abstract AsyncReadOperation<AsyncBatchCursor<T>> asAsyncReadOperation(int initialBatchSize);

driver-reactive-streams/src/main/com/mongodb/reactivestreams/client/internal/ChangeStreamPublisherImpl.java

@@ -112,7 +112,8 @@ public ChangeStreamPublisher<T> collation(@Nullable final Collation collation) {
 
     @Override
     public <TDocument> Publisher<TDocument> withDocumentClass(final Class<TDocument> clazz) {
-        return new BatchCursorPublisher<TDocument>(getClientSession(), getMongoOperationPublisher().withDocumentClass(clazz)) {
+        return new BatchCursorPublisher<TDocument>(getClientSession(), getMongoOperationPublisher().withDocumentClass(clazz),
+                getBatchSize()) {
             @Override
             AsyncReadOperation<AsyncBatchCursor<TDocument>> asAsyncReadOperation(final int initialBatchSize) {
                 return createChangeStreamOperation(getMongoOperationPublisher().getCodecRegistry().get(clazz), initialBatchSize);

driver-reactive-streams/src/test/unit/com/mongodb/reactivestreams/client/internal/ChangeStreamPublisherImplTest.java

@@ -88,13 +88,15 @@ void shouldBuildTheExpectedOperationWhenSettingDocumentClass() {
         List<BsonDocument> pipeline = singletonList(BsonDocument.parse("{'$match': 1}"));
         TestOperationExecutor executor = createOperationExecutor(singletonList(getBatchCursor()));
 
+        int batchSize = 100;
         Publisher<BsonDocument> publisher = new ChangeStreamPublisherImpl<>(null, createMongoOperationPublisher(executor),
                                                                             Document.class, pipeline, ChangeStreamLevel.COLLECTION)
-                        .withDocumentClass(BsonDocument.class);
+                .batchSize(batchSize)
+                .withDocumentClass(BsonDocument.class);
 
         ChangeStreamOperation<BsonDocument> expectedOperation =
                 new ChangeStreamOperation<>(NAMESPACE, FullDocument.DEFAULT, pipeline, getDefaultCodecRegistry().get(BsonDocument.class))
-                        .batchSize(Integer.MAX_VALUE)
+                        .batchSize(batchSize)
                         .retryReads(true);
 
         // default input should be as expected