initial import

2020-07-26 17:09:05 +08:00
commit 7e3a369a8f
647 changed files with 54754 additions and 0 deletions
--- a/core/stores/mongo/model.go
+++ b/core/stores/mongo/model.go
@@ -0,0 +1,164 @@
+package mongo
+
+import (
+	"log"
+	"time"
+
+	"github.com/globalsign/mgo"
+)
+
+type (
+	options struct {
+		timeout time.Duration
+	}
+
+	Option func(opts *options)
+
+	Model struct {
+		session    *concurrentSession
+		db         *mgo.Database
+		collection string
+		opts       []Option
+	}
+)
+
+func MustNewModel(url, database, collection string, opts ...Option) *Model {
+	model, err := NewModel(url, database, collection, opts...)
+	if err != nil {
+		log.Fatal(err)
+	}
+
+	return model
+}
+
+func NewModel(url, database, collection string, opts ...Option) (*Model, error) {
+	session, err := getConcurrentSession(url)
+	if err != nil {
+		return nil, err
+	}
+
+	return &Model{
+		session:    session,
+		db:         session.DB(database),
+		collection: collection,
+		opts:       opts,
+	}, nil
+}
+
+func (mm *Model) Find(query interface{}) (Query, error) {
+	return mm.query(func(c Collection) Query {
+		return c.Find(query)
+	})
+}
+
+func (mm *Model) FindId(id interface{}) (Query, error) {
+	return mm.query(func(c Collection) Query {
+		return c.FindId(id)
+	})
+}
+
+func (mm *Model) GetCollection(session *mgo.Session) Collection {
+	return newCollection(mm.db.C(mm.collection).With(session))
+}
+
+func (mm *Model) Insert(docs ...interface{}) error {
+	return mm.execute(func(c Collection) error {
+		return c.Insert(docs...)
+	})
+}
+
+func (mm *Model) Pipe(pipeline interface{}) (Pipe, error) {
+	return mm.pipe(func(c Collection) Pipe {
+		return c.Pipe(pipeline)
+	})
+}
+
+func (mm *Model) PutSession(session *mgo.Session) {
+	mm.session.putSession(session)
+}
+
+func (mm *Model) Remove(selector interface{}) error {
+	return mm.execute(func(c Collection) error {
+		return c.Remove(selector)
+	})
+}
+
+func (mm *Model) RemoveAll(selector interface{}) (*mgo.ChangeInfo, error) {
+	return mm.change(func(c Collection) (*mgo.ChangeInfo, error) {
+		return c.RemoveAll(selector)
+	})
+}
+
+func (mm *Model) RemoveId(id interface{}) error {
+	return mm.execute(func(c Collection) error {
+		return c.RemoveId(id)
+	})
+}
+
+func (mm *Model) TakeSession() (*mgo.Session, error) {
+	return mm.session.takeSession(mm.opts...)
+}
+
+func (mm *Model) Update(selector, update interface{}) error {
+	return mm.execute(func(c Collection) error {
+		return c.Update(selector, update)
+	})
+}
+
+func (mm *Model) UpdateId(id, update interface{}) error {
+	return mm.execute(func(c Collection) error {
+		return c.UpdateId(id, update)
+	})
+}
+
+func (mm *Model) Upsert(selector, update interface{}) (*mgo.ChangeInfo, error) {
+	return mm.change(func(c Collection) (*mgo.ChangeInfo, error) {
+		return c.Upsert(selector, update)
+	})
+}
+
+func (mm *Model) change(fn func(c Collection) (*mgo.ChangeInfo, error)) (*mgo.ChangeInfo, error) {
+	session, err := mm.TakeSession()
+	if err != nil {
+		return nil, err
+	}
+	defer mm.PutSession(session)
+
+	return fn(mm.GetCollection(session))
+}
+
+func (mm *Model) execute(fn func(c Collection) error) error {
+	session, err := mm.TakeSession()
+	if err != nil {
+		return err
+	}
+	defer mm.PutSession(session)
+
+	return fn(mm.GetCollection(session))
+}
+
+func (mm *Model) pipe(fn func(c Collection) Pipe) (Pipe, error) {
+	session, err := mm.TakeSession()
+	if err != nil {
+		return nil, err
+	}
+	defer mm.PutSession(session)
+
+	return fn(mm.GetCollection(session)), nil
+}
+
+func (mm *Model) query(fn func(c Collection) Query) (Query, error) {
+	session, err := mm.TakeSession()
+	if err != nil {
+		return nil, err
+	}
+	defer mm.PutSession(session)
+
+	return fn(mm.GetCollection(session)), nil
+}
+
+func WithTimeout(timeout time.Duration) Option {
+	return func(opts *options) {
+		opts.timeout = timeout
+	}
+}