GODRIVER-1605 Add SetBatchSize to mongo.Cursor (#1201)

Co-authored-by: Matt Dale <[email protected]>

Author: prestonvasquez

mongo/batch_cursor.go

@@ -34,6 +34,10 @@ type batchCursor interface {
 
 	// Close closes the cursor.
 	Close(context.Context) error
+
+	// The SetBatchSize method is a modifier function used to adjust the
+	// batch size of the cursor that implements it.
+	SetBatchSize(int32)
 }
 
 // changeStreamCursor is the interface implemented by batch cursors that also provide the functionality for retrieving

mongo/cursor.go

@@ -314,6 +314,14 @@ func (c *Cursor) closeImplicitSession() {
 	}
 }
 
+// SetBatchSize sets the number of documents to fetch from the database with
+// each iteration of the cursor's "Next" method. Note that some operations set
+// an initial cursor batch size, so this setting only affects subsequent
+// document batches fetched from the database.
+func (c *Cursor) SetBatchSize(batchSize int32) {
+	c.bc.SetBatchSize(batchSize)
+}
+
 // BatchCursorFromCursor returns a driver.BatchCursor for the given Cursor. If there is no underlying
 // driver.BatchCursor, nil is returned.
 //

mongo/cursor_test.go

@@ -86,6 +86,8 @@ func (tbc *testBatchCursor) Close(context.Context) error {
 	return nil
 }
 
+func (tbc *testBatchCursor) SetBatchSize(int32) {}
+
 func TestCursor(t *testing.T) {
 	t.Run("loops until docs available", func(t *testing.T) {})
 	t.Run("returns false on context cancellation", func(t *testing.T) {})

x/mongo/driver/batch_cursor.go

@@ -308,23 +308,37 @@ func (bc *BatchCursor) KillCursor(ctx context.Context) error {
 	}.Execute(ctx)
 }
 
+// calcGetMoreBatchSize calculates the number of documents to return in the
+// response of a "getMore" operation based on the given limit, batchSize, and
+// number of documents already returned. Returns false if a non-trivial limit is
+// lower than or equal to the number of documents already returned.
+func calcGetMoreBatchSize(bc BatchCursor) (int32, bool) {
+	gmBatchSize := bc.batchSize
+
+	// Account for legacy operations that don't support setting a limit.
+	if bc.limit != 0 && bc.numReturned+bc.batchSize >= bc.limit {
+		gmBatchSize = bc.limit - bc.numReturned
+		if gmBatchSize <= 0 {
+			return gmBatchSize, false
+		}
+	}
+
+	return gmBatchSize, true
+}
+
 func (bc *BatchCursor) getMore(ctx context.Context) {
 	bc.clearBatch()
 	if bc.id == 0 {
 		return
 	}
 
-	// Required for legacy operations which don't support limit.
-	numToReturn := bc.batchSize
-	if bc.limit != 0 && bc.numReturned+bc.batchSize >= bc.limit {
-		numToReturn = bc.limit - bc.numReturned
-		if numToReturn <= 0 {
-			err := bc.Close(ctx)
-			if err != nil {
-				bc.err = err
-			}
-			return
+	numToReturn, ok := calcGetMoreBatchSize(*bc)
+	if !ok {
+		if err := bc.Close(ctx); err != nil {
+			bc.err = err
 		}
+
+		return
 	}
 
 	bc.err = Operation{

x/mongo/driver/batch_cursor_test.go

@@ -13,7 +13,11 @@ import (
 )
 
 func TestBatchCursor(t *testing.T) {
+	t.Parallel()
+
 	t.Run("setBatchSize", func(t *testing.T) {
+		t.Parallel()
+
 		var size int32
 		bc := &BatchCursor{
 			batchSize: size,
@@ -24,4 +28,65 @@ func TestBatchCursor(t *testing.T) {
 		bc.SetBatchSize(size)
 		assert.Equal(t, size, bc.batchSize, "expected batchSize %v, got %v", size, bc.batchSize)
 	})
+
+	t.Run("calcGetMoreBatchSize", func(t *testing.T) {
+		t.Parallel()
+
+		for _, tcase := range []struct {
+			name                               string
+			size, limit, numReturned, expected int32
+			ok                                 bool
+		}{
+			{
+				name:     "empty",
+				expected: 0,
+				ok:       true,
+			},
+			{
+				name:     "batchSize NEQ 0",
+				size:     4,
+				expected: 4,
+				ok:       true,
+			},
+			{
+				name:     "limit NEQ 0",
+				limit:    4,
+				expected: 0,
+				ok:       true,
+			},
+			{
+				name:        "limit NEQ and batchSize + numReturned EQ limit",
+				size:        4,
+				limit:       8,
+				numReturned: 4,
+				expected:    4,
+				ok:          true,
+			},
+			{
+				name:        "limit makes batchSize negative",
+				numReturned: 4,
+				limit:       2,
+				expected:    -2,
+				ok:          false,
+			},
+		} {
+			tcase := tcase
+			t.Run(tcase.name, func(t *testing.T) {
+				t.Parallel()
+
+				bc := &BatchCursor{
+					limit:       tcase.limit,
+					batchSize:   tcase.size,
+					numReturned: tcase.numReturned,
+				}
+
+				bc.SetBatchSize(tcase.size)
+
+				size, ok := calcGetMoreBatchSize(*bc)
+
+				assert.Equal(t, tcase.expected, size, "expected batchSize %v, got %v", tcase.expected, size)
+				assert.Equal(t, tcase.ok, ok, "expected ok %v, got %v", tcase.ok, ok)
+			})
+		}
+	})
 }

x/mongo/driver/list_collections_batch_cursor.go

@@ -127,3 +127,8 @@ func (*ListCollectionsBatchCursor) projectNameElement(rawDoc bsoncore.Document)
 	filteredDoc = bsoncore.BuildDocument(filteredDoc, filteredElems)
 	return filteredDoc, nil
 }
+
+// SetBatchSize sets the batchSize for future getMores.
+func (lcbc *ListCollectionsBatchCursor) SetBatchSize(size int32) {
+	lcbc.bc.SetBatchSize(size)
+}