Skip to content

Commit

Permalink
Add Collation support for calling Count() on a Query (#166)
Browse files Browse the repository at this point in the history
  • Loading branch information
Cedric authored and domodwyer committed May 17, 2018
1 parent 2e9fa92 commit bd62d93
Show file tree
Hide file tree
Showing 2 changed files with 38 additions and 5 deletions.
11 changes: 6 additions & 5 deletions session.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -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
}
Expand Down
32 changes: 32 additions & 0 deletions session_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down

0 comments on commit bd62d93

Please sign in to comment.