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/internal" "github.com/tal-tech/go-zero/core/stores/mongo" "github.com/tal-tech/go-zero/core/stores/redis" ) type Model struct { *mongo.Model cache internal.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 := internal.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 := internal.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 internal.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 }