From 2b58f9aac9d32018fa6bfc03940f7c9d2c796741 Mon Sep 17 00:00:00 2001
From: wanghao-bianjie <53526401+wanghao-bianjie@users.noreply.github.com>
Date: Sun, 28 Aug 2022 10:59:27 +0800
Subject: [PATCH] add find collation option (#254)

---
 interface.go |  3 +++
 query.go     | 13 +++++++++++++
 2 files changed, 16 insertions(+)

diff --git a/interface.go b/interface.go
index f04e358..61ed7e8 100644
--- a/interface.go
+++ b/interface.go
@@ -13,6 +13,8 @@
 
 package qmgo
 
+import "go.mongodb.org/mongo-driver/mongo/options"
+
 // CollectionI
 // 集合操作接口
 //type CollectionI interface {
@@ -48,6 +50,7 @@ type CursorI interface {
 
 // QueryI Query interface
 type QueryI interface {
+	Collation(collation *options.Collation) QueryI
 	Sort(fields ...string) QueryI
 	Select(selector interface{}) QueryI
 	Skip(n int64) QueryI
diff --git a/query.go b/query.go
index 53c4bb9..eb0833d 100644
--- a/query.go
+++ b/query.go
@@ -37,6 +37,7 @@ type Query struct {
 	skip            *int64
 	batchSize       *int64
 	noCursorTimeout *bool
+	collation       *options.Collation
 
 	ctx        context.Context
 	collection *mongo.Collection
@@ -44,6 +45,12 @@ type Query struct {
 	registry   *bsoncodec.Registry
 }
 
+func (q *Query) Collation(collation *options.Collation) QueryI {
+	newQ := q
+	newQ.collation = collation
+	return newQ
+}
+
 func (q *Query) NoCursorTimeout(n bool) QueryI {
 	newQ := q
 	newQ.noCursorTimeout = &n
@@ -128,6 +135,9 @@ func (q *Query) One(result interface{}) error {
 	}
 	opt := options.FindOne()
 
+	if q.collation != nil {
+		opt.SetCollation(q.collation)
+	}
 	if q.sort != nil {
 		opt.SetSort(q.sort)
 	}
@@ -163,6 +173,9 @@ func (q *Query) All(result interface{}) error {
 		}
 	}
 	opt := options.Find()
+	if q.collation != nil {
+		opt.SetCollation(q.collation)
+	}
 	if q.sort != nil {
 		opt.SetSort(q.sort)
 	}