From ab49369e6c9c9919568beb7fdd15d15f23e32eb2 Mon Sep 17 00:00:00 2001
From: jiangz222 <jiangz0222@gmail.com>
Date: Sat, 9 Apr 2022 13:52:07 +0800
Subject: [PATCH] Uptoofficial (#237)

---
 go.mod        |  2 +-
 go.sum        |  4 ++--
 query.go      | 13 ++++++-------
 query_test.go | 25 +++++++++++++++++++++++++
 4 files changed, 34 insertions(+), 10 deletions(-)

diff --git a/go.mod b/go.mod
index 39e3970..ee2df0a 100644
--- a/go.mod
+++ b/go.mod
@@ -5,5 +5,5 @@ go 1.16
 require (
 	github.com/go-playground/validator/v10 v10.4.1
 	github.com/stretchr/testify v1.6.1
-	go.mongodb.org/mongo-driver v1.8.2
+	go.mongodb.org/mongo-driver v1.9.0
 )
diff --git a/go.sum b/go.sum
index 32121ee..1350751 100644
--- a/go.sum
+++ b/go.sum
@@ -44,8 +44,8 @@ github.com/xdg-go/stringprep v1.0.2 h1:6iq84/ryjjeRmMJwxutI51F2GIPlP5BfTvXHeYjyh
 github.com/xdg-go/stringprep v1.0.2/go.mod h1:8F9zXuvzgwmyT5DUm4GUfZGDdT3W+LCvS6+da4O5kxM=
 github.com/youmark/pkcs8 v0.0.0-20181117223130-1be2e3e5546d h1:splanxYIlg+5LfHAM6xpdFEAYOk8iySO56hMFq6uLyA=
 github.com/youmark/pkcs8 v0.0.0-20181117223130-1be2e3e5546d/go.mod h1:rHwXgn7JulP+udvsHwJoVG1YGAP6VLg4y9I5dyZdqmA=
-go.mongodb.org/mongo-driver v1.8.2 h1:8ssUXufb90ujcIvR6MyE1SchaNj0SFxsakiZgxIyrMk=
-go.mongodb.org/mongo-driver v1.8.2/go.mod h1:0sQWfOeY63QTntERDJJ/0SuKK0T1uVSgKCuAROlKEPY=
+go.mongodb.org/mongo-driver v1.9.0 h1:f3aLGJvQmBl8d9S40IL+jEyBC6hfLPbJjv9t5hEM9ck=
+go.mongodb.org/mongo-driver v1.9.0/go.mod h1:0sQWfOeY63QTntERDJJ/0SuKK0T1uVSgKCuAROlKEPY=
 golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACkg1iLfiJU5Ep61QUkGW8qpdssI0+w=
 golang.org/x/crypto v0.0.0-20200622213623-75b288015ac9/go.mod h1:LzIPMQfyMNhhGPhUkYOs5KpL4U8rLKemX1yGLhDgUto=
 golang.org/x/crypto v0.0.0-20201216223049-8b5274cf687f h1:aZp0e2vLN4MToVqnjNEYEtrEA8RH8U8FN1CU7JgqsPU=
diff --git a/query.go b/query.go
index 79dddf0..ca1f5e1 100644
--- a/query.go
+++ b/query.go
@@ -43,17 +43,18 @@ type Query struct {
 	registry   *bsoncodec.Registry
 }
 
-// Sort is Used to set the sorting rules for the returned results
-// Format: "age" or "+age" means to sort the age field in ascending order, "-age" means in descending order
-// When multiple sort fields are passed in at the same time, they are arranged in the order in which the fields are passed in.
-// For example, {"age", "-name"}, first sort by age in ascending order, then sort by name in descending order
-
+// BatchSize sets the value for the BatchSize field.
+// Means the maximum number of documents to be included in each batch returned by the server.
 func (q *Query) BatchSize(n int64) QueryI {
 	newQ := q
 	newQ.batchSize = &n
 	return newQ
 }
 
+// Sort is Used to set the sorting rules for the returned results
+// Format: "age" or "+age" means to sort the age field in ascending order, "-age" means in descending order
+// When multiple sort fields are passed in at the same time, they are arranged in the order in which the fields are passed in.
+// For example, {"age", "-name"}, first sort by age in ascending order, then sort by name in descending order
 func (q *Query) Sort(fields ...string) QueryI {
 	if len(fields) == 0 {
 		// A nil bson.D will not correctly serialize, but this case is no-op
@@ -155,7 +156,6 @@ func (q *Query) All(result interface{}) error {
 		}
 	}
 	opt := options.Find()
-
 	if q.sort != nil {
 		opt.SetSort(q.sort)
 	}
@@ -171,7 +171,6 @@ func (q *Query) All(result interface{}) error {
 	if q.hint != nil {
 		opt.SetHint(q.hint)
 	}
-
 	if q.batchSize != nil {
 		opt.SetBatchSize(int32(*q.batchSize))
 	}
diff --git a/query_test.go b/query_test.go
index 16e7357..ff817b1 100644
--- a/query_test.go
+++ b/query_test.go
@@ -830,3 +830,28 @@ func TestQuery_Apply(t *testing.T) {
 	ast.Equal("", res4.Name)
 	ast.Equal(0, res4.Age)
 }
+
+func TestQuery_BatchSize(t *testing.T) {
+	ast := require.New(t)
+	cli := initClient("test")
+	defer cli.Close(context.Background())
+	defer cli.DropCollection(context.Background())
+	cli.EnsureIndexes(context.Background(), nil, []string{"name"})
+
+	id1 := primitive.NewObjectID()
+	id2 := primitive.NewObjectID()
+	id3 := primitive.NewObjectID()
+	id4 := primitive.NewObjectID()
+	docs := []interface{}{
+		bson.M{"_id": id1, "name": "Alice", "age": 18},
+		bson.M{"_id": id2, "name": "Alice", "age": 19},
+		bson.M{"_id": id3, "name": "Lucas", "age": 20},
+		bson.M{"_id": id4, "name": "Lucas", "age": 21},
+	}
+	_, _ = cli.InsertMany(context.Background(), docs)
+	var res []QueryTestItem
+
+	err := cli.Find(context.Background(), bson.M{"name": "Alice"}).BatchSize(1).All(&res)
+	ast.NoError(err)
+	ast.Len(res, 2)
+}