123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241 |
- package mongoc
- import (
- "log"
- "github.com/globalsign/mgo"
- "github.com/tal-tech/go-zero/core/stores/cache"
- "github.com/tal-tech/go-zero/core/stores/mongo"
- "github.com/tal-tech/go-zero/core/stores/redis"
- )
- type Model struct {
- *mongo.Model
- cache cache.Cache
- generateCollection func(*mgo.Session) *cachedCollection
- }
- func MustNewNodeModel(url, collection string, rds *redis.Redis, opts ...cache.Option) *Model {
- model, err := NewNodeModel(url, collection, rds, opts...)
- if err != nil {
- log.Fatal(err)
- }
- return model
- }
- func MustNewModel(url, collection string, c cache.CacheConf, opts ...cache.Option) *Model {
- model, err := NewModel(url, collection, c, opts...)
- if err != nil {
- log.Fatal(err)
- }
- return model
- }
- func NewNodeModel(url, collection string, rds *redis.Redis, opts ...cache.Option) (*Model, error) {
- c := cache.NewCacheNode(rds, sharedCalls, stats, mgo.ErrNotFound, opts...)
- return createModel(url, collection, c, func(collection mongo.Collection) *cachedCollection {
- return newCollection(collection, c)
- })
- }
- func NewModel(url, collection string, conf cache.CacheConf, opts ...cache.Option) (*Model, error) {
- c := cache.NewCache(conf, sharedCalls, stats, mgo.ErrNotFound, opts...)
- return createModel(url, collection, c, func(collection mongo.Collection) *cachedCollection {
- return newCollection(collection, c)
- })
- }
- func (mm *Model) Count(query interface{}) (int, error) {
- return mm.executeInt(func(c *cachedCollection) (int, error) {
- return c.Count(query)
- })
- }
- func (mm *Model) DelCache(keys ...string) error {
- return mm.cache.DelCache(keys...)
- }
- func (mm *Model) GetCache(key string, v interface{}) error {
- return mm.cache.GetCache(key, v)
- }
- func (mm *Model) GetCollection(session *mgo.Session) *cachedCollection {
- return mm.generateCollection(session)
- }
- func (mm *Model) FindAllNoCache(v interface{}, query interface{}, opts ...QueryOption) error {
- return mm.execute(func(c *cachedCollection) error {
- return c.FindAllNoCache(v, query, opts...)
- })
- }
- func (mm *Model) FindOne(v interface{}, key string, query interface{}) error {
- return mm.execute(func(c *cachedCollection) error {
- return c.FindOne(v, key, query)
- })
- }
- func (mm *Model) FindOneNoCache(v interface{}, query interface{}) error {
- return mm.execute(func(c *cachedCollection) error {
- return c.FindOneNoCache(v, query)
- })
- }
- func (mm *Model) FindOneId(v interface{}, key string, id interface{}) error {
- return mm.execute(func(c *cachedCollection) error {
- return c.FindOneId(v, key, id)
- })
- }
- func (mm *Model) FindOneIdNoCache(v interface{}, id interface{}) error {
- return mm.execute(func(c *cachedCollection) error {
- return c.FindOneIdNoCache(v, id)
- })
- }
- func (mm *Model) Insert(docs ...interface{}) error {
- return mm.execute(func(c *cachedCollection) error {
- return c.Insert(docs...)
- })
- }
- func (mm *Model) Pipe(pipeline interface{}) (mongo.Pipe, error) {
- return mm.pipe(func(c *cachedCollection) mongo.Pipe {
- return c.Pipe(pipeline)
- })
- }
- func (mm *Model) Remove(selector interface{}, keys ...string) error {
- return mm.execute(func(c *cachedCollection) error {
- return c.Remove(selector, keys...)
- })
- }
- func (mm *Model) RemoveNoCache(selector interface{}) error {
- return mm.execute(func(c *cachedCollection) error {
- return c.RemoveNoCache(selector)
- })
- }
- func (mm *Model) RemoveAll(selector interface{}, keys ...string) (*mgo.ChangeInfo, error) {
- return mm.change(func(c *cachedCollection) (*mgo.ChangeInfo, error) {
- return c.RemoveAll(selector, keys...)
- })
- }
- func (mm *Model) RemoveAllNoCache(selector interface{}) (*mgo.ChangeInfo, error) {
- return mm.change(func(c *cachedCollection) (*mgo.ChangeInfo, error) {
- return c.RemoveAllNoCache(selector)
- })
- }
- func (mm *Model) RemoveId(id interface{}, keys ...string) error {
- return mm.execute(func(c *cachedCollection) error {
- return c.RemoveId(id, keys...)
- })
- }
- func (mm *Model) RemoveIdNoCache(id interface{}) error {
- return mm.execute(func(c *cachedCollection) error {
- return c.RemoveIdNoCache(id)
- })
- }
- func (mm *Model) SetCache(key string, v interface{}) error {
- return mm.cache.SetCache(key, v)
- }
- func (mm *Model) Update(selector, update interface{}, keys ...string) error {
- return mm.execute(func(c *cachedCollection) error {
- return c.Update(selector, update, keys...)
- })
- }
- func (mm *Model) UpdateNoCache(selector, update interface{}) error {
- return mm.execute(func(c *cachedCollection) error {
- return c.UpdateNoCache(selector, update)
- })
- }
- func (mm *Model) UpdateId(id, update interface{}, keys ...string) error {
- return mm.execute(func(c *cachedCollection) error {
- return c.UpdateId(id, update, keys...)
- })
- }
- func (mm *Model) UpdateIdNoCache(id, update interface{}) error {
- return mm.execute(func(c *cachedCollection) error {
- return c.UpdateIdNoCache(id, update)
- })
- }
- func (mm *Model) Upsert(selector, update interface{}, keys ...string) (*mgo.ChangeInfo, error) {
- return mm.change(func(c *cachedCollection) (*mgo.ChangeInfo, error) {
- return c.Upsert(selector, update, keys...)
- })
- }
- func (mm *Model) UpsertNoCache(selector, update interface{}) (*mgo.ChangeInfo, error) {
- return mm.change(func(c *cachedCollection) (*mgo.ChangeInfo, error) {
- return c.UpsertNoCache(selector, update)
- })
- }
- func (mm *Model) change(fn func(c *cachedCollection) (*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 *cachedCollection) error) error {
- session, err := mm.TakeSession()
- if err != nil {
- return err
- }
- defer mm.PutSession(session)
- return fn(mm.GetCollection(session))
- }
- func (mm *Model) executeInt(fn func(c *cachedCollection) (int, error)) (int, error) {
- session, err := mm.TakeSession()
- if err != nil {
- return 0, err
- }
- defer mm.PutSession(session)
- return fn(mm.GetCollection(session))
- }
- func (mm *Model) pipe(fn func(c *cachedCollection) mongo.Pipe) (mongo.Pipe, error) {
- session, err := mm.TakeSession()
- if err != nil {
- return nil, err
- }
- defer mm.PutSession(session)
- return fn(mm.GetCollection(session)), nil
- }
- func createModel(url, collection string, c cache.Cache,
- create func(mongo.Collection) *cachedCollection) (*Model, error) {
- model, err := mongo.NewModel(url, collection)
- if err != nil {
- return nil, err
- }
- return &Model{
- Model: model,
- cache: c,
- generateCollection: func(session *mgo.Session) *cachedCollection {
- collection := model.GetCollection(session)
- return create(collection)
- },
- }, nil
- }
|