GODRIVER-1900 add SetBatchSize to driver.BatchCursor (#699)

Author: iwysiu

mongo/integration/cursor_test.go

@@ -218,6 +218,34 @@ func TestCursor(t *testing.T) {
 			assert.Equal(mt, failpointData.ErrorCode, mongoErr.Code, "expected code %v, got: %v", failpointData.ErrorCode, mongoErr.Code)
 		})
 	})
+	// For versions < 3.2, the first find will get all the documents
+	mt.RunOpts("set batchSize", mtest.NewOptions().MinServerVersion("3.2"), func(mt *mtest.T) {
+		initCollection(mt, mt.Coll)
+		mt.ClearEvents()
+
+		// create cursor with batchSize 0
+		cursor, err := mt.Coll.Find(mtest.Background, bson.D{}, options.Find().SetBatchSize(0))
+		assert.Nil(mt, err, "Find error: %v", err)
+		defer cursor.Close(mtest.Background)
+		evt := mt.GetStartedEvent()
+		assert.Equal(mt, "find", evt.CommandName, "expected 'find' event, got '%v'", evt.CommandName)
+		sizeVal, err := evt.Command.LookupErr("batchSize")
+		assert.Nil(mt, err, "expected find command to have batchSize")
+		batchSize := sizeVal.Int32()
+		assert.Equal(mt, int32(0), batchSize, "expected batchSize 0, got %v", batchSize)
+
+		// make sure that the getMore sends the new batchSize
+		batchCursor := mongo.BatchCursorFromCursor(cursor)
+		batchCursor.SetBatchSize(4)
+		assert.True(mt, cursor.Next(mtest.Background), "expected Next true, got false")
+		evt = mt.GetStartedEvent()
+		assert.NotNil(mt, evt, "expected getMore event, got nil")
+		assert.Equal(mt, "getMore", evt.CommandName, "expected 'getMore' event, got '%v'", evt.CommandName)
+		sizeVal, err = evt.Command.LookupErr("batchSize")
+		assert.Nil(mt, err, "expected getMore command to have batchSize")
+		batchSize = sizeVal.Int32()
+		assert.Equal(mt, int32(4), batchSize, "expected batchSize 4, got %v", batchSize)
+	})
 }
 
 type tryNextCursor interface {

x/mongo/driver/batch_cursor.go

@@ -388,6 +388,11 @@ func (bc *BatchCursor) PostBatchResumeToken() bsoncore.Document {
 	return bc.postBatchResumeToken
 }
 
+// SetBatchSize sets the batchSize for future getMores.
+func (bc *BatchCursor) SetBatchSize(size int32) {
+	bc.batchSize = size
+}
+
 func (bc *BatchCursor) getOperationDeployment() Deployment {
 	if bc.connection != nil {
 		return &loadBalancedCursorDeployment{

x/mongo/driver/batch_cursor_test.go

@@ -0,0 +1,21 @@
+package driver
+
+import (
+	"testing"
+
+	"go.mongodb.org/mongo-driver/internal/testutil/assert"
+)
+
+func TestBatchCursor(t *testing.T) {
+	t.Run("setBatchSize", func(t *testing.T) {
+		var size int32
+		bc := &BatchCursor{
+			batchSize: size,
+		}
+		assert.Equal(t, size, bc.batchSize, "expected batchSize %v, got %v", size, bc.batchSize)
+
+		size = int32(4)
+		bc.SetBatchSize(size)
+		assert.Equal(t, size, bc.batchSize, "expected batchSize %v, got %v", size, bc.batchSize)
+	})
+}