From 115c288a0be75acb4e0aae49bf10e5739273ad83 Mon Sep 17 00:00:00 2001
From: Sebastian Frank <sebastian@webmakers.de>
Date: Wed, 9 Feb 2022 15:49:15 +0100
Subject: [PATCH] encapsulate ErrNotFound

---
 session.go | 35 +++++++++++++++++++++++++++++++----
 1 file changed, 31 insertions(+), 4 deletions(-)

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
+}