Skip to content

Commit 3ffdf76

Browse files
benjirewisBenjamin Rewis
authored andcommitted
GODRIVER-1831 Avoid accidental extra getMore when using limit on Find (#562)
1 parent c1a6bf1 commit 3ffdf76

File tree

2 files changed

+98
-1
lines changed

2 files changed

+98
-1
lines changed

mongo/integration/collection_test.go

Lines changed: 97 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -796,6 +796,103 @@ func TestCollection(t *testing.T) {
796796
_, ok := err.(mongo.CommandError)
797797
assert.True(mt, ok, "expected error type %v, got %v", mongo.CommandError{}, err)
798798
})
799+
mt.Run("limit and batch size and skip", func(mt *mtest.T) {
800+
testCases := []struct {
801+
limit int64
802+
batchSize int32
803+
skip int64
804+
name string
805+
}{
806+
{
807+
99, 100, 10,
808+
"case 1",
809+
},
810+
{
811+
100, 100, 20,
812+
"case 2",
813+
},
814+
{
815+
80, 20, 90,
816+
"case 3",
817+
},
818+
{
819+
201, 201, 0,
820+
"case 4",
821+
},
822+
{
823+
100, 200, 120,
824+
"case 5",
825+
},
826+
}
827+
for _, tc := range testCases {
828+
mt.Run(tc.name, func(mt *mtest.T) {
829+
var insertDocs []interface{}
830+
for i := 1; i <= 201; i++ {
831+
insertDocs = append(insertDocs, bson.D{{"x", int32(i)}})
832+
}
833+
834+
_, err := mt.Coll.InsertMany(mtest.Background, insertDocs)
835+
assert.Nil(mt, err, "InsertMany error for initial data: %v", err)
836+
837+
findOptions := options.Find().SetLimit(tc.limit).SetBatchSize(tc.batchSize).
838+
SetSkip(tc.skip)
839+
cursor, err := mt.Coll.Find(mtest.Background, bson.D{}, findOptions)
840+
assert.Nil(mt, err, "Find error: %v", err)
841+
842+
var docs []interface{}
843+
err = cursor.All(mtest.Background, &docs)
844+
assert.Nil(mt, err, "All error: %v", err)
845+
if (201 - tc.skip) < tc.limit {
846+
assert.Equal(mt, int(201-tc.skip), len(docs), "expected number of docs to be %v, got %v", int(201-tc.skip), len(docs))
847+
} else {
848+
assert.Equal(mt, int(tc.limit), len(docs), "expected number of docs to be %v, got %v", tc.limit, len(docs))
849+
}
850+
})
851+
}
852+
})
853+
mt.Run("unset batch size does not surpass limit", func(mt *mtest.T) {
854+
testCases := []struct {
855+
limit int64
856+
name string
857+
}{
858+
{
859+
99,
860+
"99",
861+
},
862+
{
863+
100,
864+
"100",
865+
},
866+
{
867+
101,
868+
"101",
869+
},
870+
{
871+
200,
872+
"200",
873+
},
874+
}
875+
for _, tc := range testCases {
876+
mt.Run(tc.name, func(mt *mtest.T) {
877+
var insertDocs []interface{}
878+
for i := 1; i <= 201; i++ {
879+
insertDocs = append(insertDocs, bson.D{{"x", int32(i)}})
880+
}
881+
882+
_, err := mt.Coll.InsertMany(mtest.Background, insertDocs)
883+
assert.Nil(mt, err, "InsertMany error for initial data: %v", err)
884+
opts := options.Find().SetSkip(0).SetLimit(tc.limit)
885+
cursor, err := mt.Coll.Find(mtest.Background, bson.D{}, opts)
886+
assert.Nil(mt, err, "Find error with limit %v: %v", tc.limit, err)
887+
888+
var docs []interface{}
889+
err = cursor.All(mtest.Background, &docs)
890+
assert.Nil(mt, err, "All error with limit %v: %v", tc.limit, err)
891+
892+
assert.Equal(mt, int(tc.limit), len(docs), "expected number of docs to be %v, got %v", tc.limit, len(docs))
893+
})
894+
}
895+
})
799896
})
800897
mt.RunOpts("find one", noClientOpts, func(mt *mtest.T) {
801898
mt.Run("limit", func(mt *mtest.T) {

x/mongo/driver/batch_cursor.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -241,7 +241,7 @@ func (bc *BatchCursor) getMore(ctx context.Context) {
241241

242242
// Required for legacy operations which don't support limit.
243243
numToReturn := bc.batchSize
244-
if bc.limit != 0 && bc.numReturned+bc.batchSize > bc.limit {
244+
if bc.limit != 0 && bc.numReturned+bc.batchSize >= bc.limit {
245245
numToReturn = bc.limit - bc.numReturned
246246
if numToReturn <= 0 {
247247
err := bc.Close(ctx)

0 commit comments

Comments
 (0)