Add Collation support for calling Count() on a Query (#166)

Author: Cedric

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
 }

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)