diff --git a/session.go b/session.go index a06bdee..254b973 100644 --- a/session.go +++ b/session.go @@ -109,7 +109,11 @@ func (q *Query) Select(selector interface{}) *Query { } func (q *Query) One(result interface{}) error { - return q.query.One(result) + err := q.query.One(result) + if err == mgo.ErrNotFound { + err = &ErrNotFound{msg: err.Error()} + } + return err } func (q *Query) Sort(fields ...string) *Query { @@ -128,11 +132,19 @@ func (q *Query) Limit(n int) *Query { } func (q *Query) All(result interface{}) error { - return q.query.All(result) + err := q.query.All(result) + if err == mgo.ErrNotFound { + err = &ErrNotFound{msg: err.Error()} + } + return err } func (q *Query) Count() (int, error) { - return q.query.Count() + c, err := q.query.Count() + if err == mgo.ErrNotFound { + err = &ErrNotFound{msg: err.Error()} + } + return c, err } type Index struct { @@ -155,9 +167,24 @@ type Pipe struct { } func (p *Pipe) All(result interface{}) error { - return p.pipe.All(result) + err := p.pipe.All(result) + if err == mgo.ErrNotFound { + err = &ErrNotFound{msg: err.Error()} + } + return err } func (c *Collection) Pipe(pipeline interface{}) *Pipe { return &Pipe{pipe: c.collection.Pipe(pipeline).AllowDiskUse()} } + +type ErrNotFound struct { + msg string +} + +func (e *ErrNotFound) Error() string { + if e.msg == "" { + return "not found" + } + return e.msg +}