Merge pull request #1 from yeeuu/develop

New Storage and QuerySet API

Author: ipfans

examples/main.go

@@ -36,34 +36,24 @@ func main() {
 		panic(err)
 	}
 	store := mongo.NewStorage(sess)
-	err = store.Insert(&userSelector{}, user{Name: "hello"})
+	err = store.Query(&userSelector{}).Insert(user{Name: "hello"})
 	if err != nil {
 		panic(err)
 	}
 	var u user
 	name := "hello"
-	qs, err := store.Query(&userSelector{
+	err = store.Query(&userSelector{
 		Name: &name,
-	})
+	}).One(&u)
 	if err != nil {
 		panic(err)
 	}
-	err = qs.One(&u)
-	if err != nil {
-		panic(err)
-	}
-	qs.Close()
 	fmt.Println(u)
-	err = store.Insert(&userSelector{}, user{Name: "world"})
-	qs, err = store.Query(&userSelector{})
-	if err != nil {
-		panic(err)
-	}
+	store.Query(&userSelector{}).Insert(user{Name: "world"})
 	var users []user
-	err = qs.All(&users)
+	err = store.Query(&userSelector{}).All(&users)
 	if err != nil {
 		panic(err)
 	}
-	qs.Close()
 	fmt.Println(users)
 }

queryset.go

@@ -1,56 +1,166 @@
 package mongo
 
-import "github.com/globalsign/mgo"
+import (
+	"github.com/globalsign/mgo"
+	"github.com/globalsign/mgo/bson"
+)
 
 // QuerySet likes mgo.Query.
 type QuerySet interface {
+	Insert(docs ...interface{}) error
+	Update(update interface{}) error
+	UpdateAll(update interface{}) (info *mgo.ChangeInfo, err error)
+	Upsert(update interface{}) (info *mgo.ChangeInfo, err error)
+	Remove() error
+	RemoveAll() (info *mgo.ChangeInfo, err error)
 	All(result interface{}) error
 	One(result interface{}) error
 	Count() (int, error)
 	Limit(n int) QuerySet
 	Skip(n int) QuerySet
 	Prefetch(p float64) QuerySet
 	Sort(fields ...string) QuerySet
-	Close()
 }
 
 type mongoQuerySet struct {
-	sess  *mgo.Session
-	query *mgo.Query
+	sess       *mgo.Session
+	database   string
+	collection string
+	query      bson.M
+	err        error
+	limit      *int
+	skip       *int
+	prefetch   *float64
+	sort       []string
+}
+
+func (qs *mongoQuerySet) getSession() (*mgo.Collection, func()) {
+	sess := qs.sess.Copy()
+	return sess.DB(qs.database).C(qs.collection), func() {
+		sess.Close()
+	}
+}
+
+func (qs *mongoQuerySet) getQuery(c *mgo.Collection) *mgo.Query {
+	q := c.Find(qs.query)
+	if qs.limit != nil {
+		q = q.Limit(*qs.limit)
+	}
+	if qs.skip != nil {
+		q = q.Skip(*qs.skip)
+	}
+	if qs.prefetch != nil {
+		q = q.Prefetch(*qs.prefetch)
+	}
+	if len(qs.sort) != 0 {
+		q = q.Sort(qs.sort...)
+	}
+	return q
+}
+
+func (qs *mongoQuerySet) Insert(docs ...interface{}) error {
+	c, close := qs.getSession()
+	defer close()
+	err := c.Insert(docs...)
+	return err
+}
+
+func (qs *mongoQuerySet) Update(update interface{}) error {
+	if qs.err != nil {
+		return qs.err
+	}
+	c, close := qs.getSession()
+	defer close()
+	err := c.Update(qs.query, update)
+	return err
+}
+
+func (qs *mongoQuerySet) UpdateAll(update interface{}) (info *mgo.ChangeInfo, err error) {
+	if qs.err != nil {
+		err = qs.err
+		return
+	}
+	c, close := qs.getSession()
+	defer close()
+	info, err = c.UpdateAll(qs.query, update)
+	return
+}
+
+func (qs *mongoQuerySet) Upsert(update interface{}) (info *mgo.ChangeInfo, err error) {
+	if qs.err != nil {
+		err = qs.err
+		return
+	}
+	c, close := qs.getSession()
+	defer close()
+	info, err = c.Upsert(qs.query, update)
+	return
+}
+
+func (qs *mongoQuerySet) Remove() error {
+	if qs.err != nil {
+		return qs.err
+	}
+	c, close := qs.getSession()
+	defer close()
+	err := c.Remove(qs.query)
+	return err
+}
+
+func (qs *mongoQuerySet) RemoveAll() (info *mgo.ChangeInfo, err error) {
+	if qs.err != nil {
+		err = qs.err
+		return
+	}
+	c, close := qs.getSession()
+	defer close()
+	info, err = c.RemoveAll(qs.query)
+	return
 }
 
 func (qs *mongoQuerySet) All(result interface{}) error {
-	return qs.query.All(result)
+	if qs.err != nil {
+		return qs.err
+	}
+	c, close := qs.getSession()
+	defer close()
+	return qs.getQuery(c).All(result)
 }
 
 func (qs *mongoQuerySet) One(result interface{}) error {
-	return qs.query.One(result)
+	if qs.err != nil {
+		return qs.err
+	}
+	c, close := qs.getSession()
+	defer close()
+	return qs.getQuery(c).One(result)
 }
 
 func (qs *mongoQuerySet) Count() (int, error) {
-	return qs.query.Count()
+	if qs.err != nil {
+		return 0, qs.err
+	}
+	c, close := qs.getSession()
+	defer close()
+	return qs.getQuery(c).Count()
 }
 
 func (qs *mongoQuerySet) Limit(n int) QuerySet {
-	qs.query = qs.query.Limit(n)
+	qs.limit = &n
 	return qs
 }
 
 func (qs *mongoQuerySet) Skip(n int) QuerySet {
-	qs.query = qs.query.Skip(n)
+	qs.skip = &n
 	return qs
 }
 
 func (qs *mongoQuerySet) Prefetch(p float64) QuerySet {
-	qs.query = qs.query.Prefetch(p)
+	qs.prefetch = &p
 	return qs
 }
 
 func (qs *mongoQuerySet) Sort(fields ...string) QuerySet {
-	qs.query = qs.query.Sort(fields...)
+	qs.sort = fields
 	return qs
 }
-
-func (qs *mongoQuerySet) Close() {
-	qs.sess.Close()
-}

storage.go

@@ -2,19 +2,20 @@ package mongo
 
 import (
 	"github.com/globalsign/mgo"
+	"github.com/globalsign/mgo/bson"
 	"github.com/pkg/errors"
 )
 
 // Storage supply a interface for make Queryset.
 type Storage interface {
-	Query(interface{}) (QuerySet, error)
-	Raw(interface{}, interface{}) (QuerySet, error)
-	Insert(base interface{}, docs ...interface{}) error
-	Update(filter, update interface{}) error
-	UpdateAll(filter, update interface{}) (info *mgo.ChangeInfo, err error)
-	Upsert(filter, update interface{}) (info *mgo.ChangeInfo, err error)
-	Remove(filter interface{}) error
-	RemoveAll(filter interface{}) (info *mgo.ChangeInfo, err error)
+	Query(Collection) QuerySet
+	Raw(Collection, bson.M) QuerySet
+	// Insert(c Collection, docs ...interface{}) error
+	// Update(filter, update interface{}) error
+	// UpdateAll(filter, update interface{}) (info *mgo.ChangeInfo, err error)
+	// Upsert(filter, update interface{}) (info *mgo.ChangeInfo, err error)
+	// Remove(filter interface{}) error
+	// RemoveAll(filter interface{}) (info *mgo.ChangeInfo, err error)
 }
 
 type mongoStorage struct {
@@ -28,99 +29,24 @@ func NewStorage(sess *mgo.Session) Storage {
 	}
 }
 
-func (ms *mongoStorage) realAction(action string, base interface{}, docs ...interface{}) (info *mgo.ChangeInfo, err error) {
-	sess := ms.sess.Copy()
-	defer sess.Close()
-	c, err := getCollection(sess, base)
-	if err != nil {
-		err = errors.Wrap(err, "Get collection info failed")
-		return
-	}
-	realQuery, err := structToBsonM(base)
-	if err != nil {
-		err = errors.Wrap(err, "Convert struct to query failed")
-		return
-	}
-	switch action {
-	case "insert":
-		err = c.Insert(docs...)
-	case "update":
-		err = c.Update(realQuery, docs[0])
-	case "updateAll":
-		info, err = c.UpdateAll(realQuery, docs[0])
-	case "upsert":
-		info, err = c.Upsert(realQuery, docs[0])
-	case "remove":
-		err = c.Remove(realQuery)
-	case "removeAll":
-		info, err = c.RemoveAll(realQuery)
-	}
-	return
-}
-
-func (ms *mongoStorage) Insert(base interface{}, docs ...interface{}) error {
-	_, err := ms.realAction("insert", base, docs...)
-	return err
-}
-
-func (ms *mongoStorage) Update(filter, update interface{}) error {
-	_, err := ms.realAction("update", filter, update)
-	return err
-}
-
-func (ms *mongoStorage) UpdateAll(filter, update interface{}) (info *mgo.ChangeInfo, err error) {
-	info, err = ms.realAction("updateAll", filter, update)
-	return
-}
-
-func (ms *mongoStorage) Upsert(filter, update interface{}) (info *mgo.ChangeInfo, err error) {
-	info, err = ms.realAction("upsert", filter, update)
-	return
-}
-
-func (ms *mongoStorage) Remove(filter interface{}) error {
-	_, err := ms.realAction("remove", filter, nil)
-	return err
-}
-
-func (ms *mongoStorage) RemoveAll(filter interface{}) (info *mgo.ChangeInfo, err error) {
-	info, err = ms.realAction("removeAll", filter, nil)
-	return
-}
-
-func (ms *mongoStorage) Query(query interface{}) (QuerySet, error) {
-	var err error
-	var c *mgo.Collection
-	sess := ms.sess.Copy()
+func (ms *mongoStorage) Query(query Collection) QuerySet {
 	qs := &mongoQuerySet{
-		sess: sess,
-	}
-	c, err = getCollection(sess, query)
-	if err != nil {
-		err = errors.Wrap(err, "Get collection info failed")
-		return qs, err
+		sess:       ms.sess,
+		database:   query.Database(),
+		collection: query.Collection(),
 	}
 	realQuery, err := structToBsonM(query)
-	if err != nil {
-		err = errors.Wrap(err, "Convert struct to query failed")
-		return qs, err
-	}
-	qs.query = c.Find(realQuery)
-	return qs, err
+	qs.err = errors.Wrap(err, "convert to bson.M failed")
+	qs.query = realQuery
+	return qs
 }
 
-func (ms *mongoStorage) Raw(base, query interface{}) (QuerySet, error) {
-	var err error
-	var c *mgo.Collection
-	sess := ms.sess.Copy()
+func (ms *mongoStorage) Raw(base Collection, query bson.M) QuerySet {
 	qs := &mongoQuerySet{
-		sess: sess,
-	}
-	c, err = getCollection(sess, query)
-	if err != nil {
-		err = errors.Wrap(err, "Get collection info failed")
-		return qs, err
+		sess:       ms.sess,
+		database:   base.Database(),
+		collection: base.Collection(),
+		query:      query,
 	}
-	qs.query = c.Find(query)
-	return qs, err
+	return qs
 }