From ccc74cbe566038f883eb4e608c93e8096241099e Mon Sep 17 00:00:00 2001
From: Cedric Cordenier <ccordenier@utilitywarehouse.co.uk>
Date: Tue, 15 May 2018 18:43:28 +0100
Subject: [PATCH] Add Collation support for calling Count() on a Query

---
 session.go      | 11 ++++++-----
 session_test.go | 32 ++++++++++++++++++++++++++++++++
 2 files changed, 38 insertions(+), 5 deletions(-)

diff --git a/session.go b/session.go
index 61a9cb22d..5bc918730 100644
--- a/session.go
+++ b/session.go
@@ -4564,10 +4564,11 @@ func (iter *Iter) getMoreCmd() *queryOp {
 type countCmd struct {
 	Count     string
 	Query     interface{}
-	Limit     int32  `bson:",omitempty"`
-	Skip      int32  `bson:",omitempty"`
-	Hint      bson.D `bson:"hint,omitempty"`
-	MaxTimeMS int    `bson:"maxTimeMS,omitempty"`
+	Limit     int32      `bson:",omitempty"`
+	Skip      int32      `bson:",omitempty"`
+	Hint      bson.D     `bson:"hint,omitempty"`
+	MaxTimeMS int        `bson:"maxTimeMS,omitempty"`
+	Collation *Collation `bson:"collation,omitempty"`
 }
 
 // Count returns the total number of documents in the result set.
@@ -4593,7 +4594,7 @@ func (q *Query) Count() (n int, err error) {
 	// simply want a Zero bson.D
 	hint, _ := q.op.options.Hint.(bson.D)
 	result := struct{ N int }{}
-	err = session.DB(dbname).Run(countCmd{cname, query, limit, op.skip, hint, op.options.MaxTimeMS}, &result)
+	err = session.DB(dbname).Run(countCmd{cname, query, limit, op.skip, hint, op.options.MaxTimeMS, op.options.Collation}, &result)
 
 	return result.N, err
 }
diff --git a/session_test.go b/session_test.go
index 23396434d..f09018c70 100644
--- a/session_test.go
+++ b/session_test.go
@@ -1560,6 +1560,38 @@ func (s *S) TestCountQuery(c *C) {
 	c.Assert(n, Equals, 2)
 }
 
+func (s *S) TestCountQueryWithCollation(c *C) {
+	if !s.versionAtLeast(3, 4) {
+		c.Skip("depends on mongodb 3.4+")
+	}
+	session, err := mgo.Dial("localhost:40001")
+	c.Assert(err, IsNil)
+	defer session.Close()
+
+	coll := session.DB("mydb").C("mycoll")
+	c.Assert(err, IsNil)
+
+	collation := &mgo.Collation{
+		Locale:   "en",
+		Strength: 2,
+	}
+	err = coll.EnsureIndex(mgo.Index{
+		Key:       []string{"n"},
+		Collation: collation,
+	})
+	c.Assert(err, IsNil)
+
+	ns := []string{"hello", "Hello", "hEllO"}
+	for _, n := range ns {
+		err := coll.Insert(M{"n": n})
+		c.Assert(err, IsNil)
+	}
+
+	n, err := coll.Find(M{"n": "hello"}).Collation(collation).Count()
+	c.Assert(err, IsNil)
+	c.Assert(n, Equals, 3)
+}
+
 func (s *S) TestCountQuerySorted(c *C) {
 	session, err := mgo.Dial("localhost:40001")
 	c.Assert(err, IsNil)