Add All method to mongo.Cursor

GODRIVER-916

Change-Id: I6eca805d21a1950b6cee808248ec6edbd9dcf8bc

Author: Divjot Arora

mongo/cursor.go

@@ -10,6 +10,7 @@ import (
 	"context"
 	"errors"
 	"io"
+	"reflect"
 
 	"go.mongodb.org/mongo-driver/bson"
 	"go.mongodb.org/mongo-driver/bson/bsoncodec"
@@ -128,3 +129,68 @@ func (c *Cursor) Err() error { return c.err }
 
 // Close closes this cursor.
 func (c *Cursor) Close(ctx context.Context) error { return c.bc.Close(ctx) }
+
+// All iterates the cursor and decodes each document into results.
+// The results parameter must be a pointer to a slice. The slice pointed to by results will be completely overwritten.
+// If the cursor has been iterated, any previously iterated documents will not be included in results.
+func (c *Cursor) All(ctx context.Context, results interface{}) error {
+	resultsVal := reflect.ValueOf(results)
+	if resultsVal.Kind() != reflect.Ptr {
+		return errors.New("results argument must be a pointer to a slice")
+	}
+
+	sliceVal := resultsVal.Elem()
+	elementType := sliceVal.Type().Elem()
+	var index int
+	var err error
+
+	batch := c.batch // exhaust the current batch before iterating the batch cursor
+	for {
+		sliceVal, index, err = c.addFromBatch(sliceVal, elementType, batch, index)
+		if err != nil {
+			return err
+		}
+
+		if !c.bc.Next(ctx) {
+			break
+		}
+
+		batch = c.bc.Batch()
+	}
+
+	if err = c.bc.Err(); err != nil {
+		return err
+	}
+
+	resultsVal.Elem().Set(sliceVal.Slice(0, index))
+	return nil
+}
+
+// addFromBatch adds all documents from batch to sliceVal starting at the given index. It returns the new slice value,
+// the next empty index in the slice, and an error if one occurs.
+func (c *Cursor) addFromBatch(sliceVal reflect.Value, elemType reflect.Type, batch *bsoncore.DocumentSequence,
+	index int) (reflect.Value, int, error) {
+
+	docs, err := batch.Documents()
+	if err != nil {
+		return sliceVal, index, err
+	}
+
+	for _, doc := range docs {
+		if sliceVal.Len() == index {
+			// slice is full
+			newElem := reflect.New(elemType)
+			sliceVal = reflect.Append(sliceVal, newElem.Elem())
+			sliceVal = sliceVal.Slice(0, sliceVal.Cap())
+		}
+
+		currElem := sliceVal.Index(index).Addr().Interface()
+		if err = bson.UnmarshalWithRegistry(c.registry, doc, currElem); err != nil {
+			return sliceVal, index, err
+		}
+
+		index++
+	}
+
+	return sliceVal, index, nil
+}

mongo/cursor_test.go

@@ -1,10 +1,137 @@
 package mongo
 
-import "testing"
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/x/bsonx/bsoncore"
+	"go.mongodb.org/mongo-driver/x/mongo/driver/topology"
+)
+
+type testBatchCursor struct {
+	batches []*bsoncore.DocumentSequence
+	batch   *bsoncore.DocumentSequence
+}
+
+func newTestBatchCursor(numBatches, batchSize int) *testBatchCursor {
+	batches := make([]*bsoncore.DocumentSequence, 0, numBatches)
+
+	counter := 0
+	for batch := 0; batch < numBatches; batch++ {
+		var docSequence []byte
+
+		for doc := 0; doc < batchSize; doc++ {
+			var elem []byte
+			elem = bsoncore.AppendInt32Element(elem, "foo", int32(counter))
+			counter++
+
+			var doc []byte
+			doc = bsoncore.BuildDocumentFromElements(doc, elem)
+			docSequence = append(docSequence, doc...)
+		}
+
+		batches = append(batches, &bsoncore.DocumentSequence{
+			Style: bsoncore.SequenceStyle,
+			Data:  docSequence,
+		})
+	}
+
+	return &testBatchCursor{
+		batches: batches,
+	}
+}
+
+func (tbc *testBatchCursor) ID() int64 {
+	if len(tbc.batches) == 0 {
+		return 0 // cursor exhausted
+	}
+
+	return 10
+}
+
+func (tbc *testBatchCursor) Next(context.Context) bool {
+	if len(tbc.batches) == 0 {
+		return false
+	}
+
+	tbc.batch = tbc.batches[0]
+	tbc.batches = tbc.batches[1:]
+	return true
+}
+
+func (tbc *testBatchCursor) Batch() *bsoncore.DocumentSequence {
+	return tbc.batch
+}
+
+func (tbc *testBatchCursor) Server() *topology.Server {
+	return nil
+}
+
+func (tbc *testBatchCursor) Err() error {
+	return nil
+}
+
+func (tbc *testBatchCursor) Close(context.Context) error {
+	return nil
+}
 
 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) {})
 	t.Run("returns false if error occurred", func(t *testing.T) {})
 	t.Run("returns false if ID is zero and no more docs", func(t *testing.T) {})
+
+	t.Run("TestAll", func(t *testing.T) {
+		t.Run("errors if argument is not pointer to slice", func(t *testing.T) {
+			cursor, err := newCursor(newTestBatchCursor(1, 5), nil)
+			require.Nil(t, err)
+			err = cursor.All(context.Background(), []bson.D{})
+			require.NotNil(t, err)
+		})
+
+		t.Run("fills slice with all documents", func(t *testing.T) {
+			cursor, err := newCursor(newTestBatchCursor(1, 5), nil)
+			require.Nil(t, err)
+
+			var docs []bson.D
+			err = cursor.All(context.Background(), &docs)
+			require.Nil(t, err)
+			require.Equal(t, 5, len(docs))
+
+			for index, doc := range docs {
+				require.Equal(t, doc, bson.D{{"foo", int32(index)}})
+			}
+		})
+
+		t.Run("decodes each document into slice type", func(t *testing.T) {
+			cursor, err := newCursor(newTestBatchCursor(1, 5), nil)
+			require.Nil(t, err)
+
+			type Document struct {
+				Foo int32 `bson:"foo"`
+			}
+			var docs []Document
+			err = cursor.All(context.Background(), &docs)
+			require.Nil(t, err)
+			require.Equal(t, 5, len(docs))
+
+			for index, doc := range docs {
+				require.Equal(t, doc, Document{Foo: int32(index)})
+			}
+		})
+
+		t.Run("multiple batches are included", func(t *testing.T) {
+			cursor, err := newCursor(newTestBatchCursor(2, 5), nil)
+			var docs []bson.D
+			err = cursor.All(context.Background(), &docs)
+			require.Nil(t, err)
+			require.Equal(t, 10, len(docs))
+
+			for index, doc := range docs {
+				require.Equal(t, doc, bson.D{{"foo", int32(index)}})
+			}
+		})
+	})
 }

x/bsonx/bsoncore/document_sequence.go

@@ -83,9 +83,9 @@ func (ds *DocumentSequence) ResetIterator() {
 	ds.Pos = 0
 }
 
-// documents returns a slice of the documents. If nil either the Data field is also nil or could not
+// Documents returns a slice of the documents. If nil either the Data field is also nil or could not
 // be properly read.
-func (ds *DocumentSequence) documents() ([]Document, error) {
+func (ds *DocumentSequence) Documents() ([]Document, error) {
 	if ds == nil {
 		return nil, nil
 	}

x/bsonx/bsoncore/document_sequence_test.go

@@ -103,7 +103,7 @@ func TestDocumentSequence(t *testing.T) {
 					Style: tc.style,
 					Data:  tc.data,
 				}
-				documents, err := ds.documents()
+				documents, err := ds.Documents()
 				if !cmp.Equal(documents, tc.documents) {
 					t.Errorf("Documents do not match. got %v; want %v", documents, tc.documents)
 				}
@@ -252,7 +252,7 @@ func TestDocumentSequence(t *testing.T) {
 					Style: tc.style,
 					Data:  tc.data,
 				}
-				docs, err := ds.documents()
+				docs, err := ds.Documents()
 				if err != nil {
 					t.Fatalf("Unexpected error: %v", err)
 				}
@@ -376,7 +376,7 @@ func TestDocumentSequence(t *testing.T) {
 		t.Run("Documents", func(t *testing.T) {
 			defer capturePanic()
 			var ds *DocumentSequence
-			_, _ = ds.documents()
+			_, _ = ds.Documents()
 		})
 		t.Run("Next", func(t *testing.T) {
 			defer capturePanic()