cachedmodel.go 6.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241
  1. package mongoc
  2. import (
  3. "log"
  4. "github.com/globalsign/mgo"
  5. "github.com/tal-tech/go-zero/core/stores/cache"
  6. "github.com/tal-tech/go-zero/core/stores/mongo"
  7. "github.com/tal-tech/go-zero/core/stores/redis"
  8. )
  9. type Model struct {
  10. *mongo.Model
  11. cache cache.Cache
  12. generateCollection func(*mgo.Session) *cachedCollection
  13. }
  14. func MustNewNodeModel(url, collection string, rds *redis.Redis, opts ...cache.Option) *Model {
  15. model, err := NewNodeModel(url, collection, rds, opts...)
  16. if err != nil {
  17. log.Fatal(err)
  18. }
  19. return model
  20. }
  21. func MustNewModel(url, collection string, c cache.CacheConf, opts ...cache.Option) *Model {
  22. model, err := NewModel(url, collection, c, opts...)
  23. if err != nil {
  24. log.Fatal(err)
  25. }
  26. return model
  27. }
  28. func NewNodeModel(url, collection string, rds *redis.Redis, opts ...cache.Option) (*Model, error) {
  29. c := cache.NewNode(rds, sharedCalls, stats, mgo.ErrNotFound, opts...)
  30. return createModel(url, collection, c, func(collection mongo.Collection) *cachedCollection {
  31. return newCollection(collection, c)
  32. })
  33. }
  34. func NewModel(url, collection string, conf cache.CacheConf, opts ...cache.Option) (*Model, error) {
  35. c := cache.New(conf, sharedCalls, stats, mgo.ErrNotFound, opts...)
  36. return createModel(url, collection, c, func(collection mongo.Collection) *cachedCollection {
  37. return newCollection(collection, c)
  38. })
  39. }
  40. func (mm *Model) Count(query interface{}) (int, error) {
  41. return mm.executeInt(func(c *cachedCollection) (int, error) {
  42. return c.Count(query)
  43. })
  44. }
  45. func (mm *Model) DelCache(keys ...string) error {
  46. return mm.cache.Del(keys...)
  47. }
  48. func (mm *Model) GetCache(key string, v interface{}) error {
  49. return mm.cache.Get(key, v)
  50. }
  51. func (mm *Model) GetCollection(session *mgo.Session) *cachedCollection {
  52. return mm.generateCollection(session)
  53. }
  54. func (mm *Model) FindAllNoCache(v interface{}, query interface{}, opts ...QueryOption) error {
  55. return mm.execute(func(c *cachedCollection) error {
  56. return c.FindAllNoCache(v, query, opts...)
  57. })
  58. }
  59. func (mm *Model) FindOne(v interface{}, key string, query interface{}) error {
  60. return mm.execute(func(c *cachedCollection) error {
  61. return c.FindOne(v, key, query)
  62. })
  63. }
  64. func (mm *Model) FindOneNoCache(v interface{}, query interface{}) error {
  65. return mm.execute(func(c *cachedCollection) error {
  66. return c.FindOneNoCache(v, query)
  67. })
  68. }
  69. func (mm *Model) FindOneId(v interface{}, key string, id interface{}) error {
  70. return mm.execute(func(c *cachedCollection) error {
  71. return c.FindOneId(v, key, id)
  72. })
  73. }
  74. func (mm *Model) FindOneIdNoCache(v interface{}, id interface{}) error {
  75. return mm.execute(func(c *cachedCollection) error {
  76. return c.FindOneIdNoCache(v, id)
  77. })
  78. }
  79. func (mm *Model) Insert(docs ...interface{}) error {
  80. return mm.execute(func(c *cachedCollection) error {
  81. return c.Insert(docs...)
  82. })
  83. }
  84. func (mm *Model) Pipe(pipeline interface{}) (mongo.Pipe, error) {
  85. return mm.pipe(func(c *cachedCollection) mongo.Pipe {
  86. return c.Pipe(pipeline)
  87. })
  88. }
  89. func (mm *Model) Remove(selector interface{}, keys ...string) error {
  90. return mm.execute(func(c *cachedCollection) error {
  91. return c.Remove(selector, keys...)
  92. })
  93. }
  94. func (mm *Model) RemoveNoCache(selector interface{}) error {
  95. return mm.execute(func(c *cachedCollection) error {
  96. return c.RemoveNoCache(selector)
  97. })
  98. }
  99. func (mm *Model) RemoveAll(selector interface{}, keys ...string) (*mgo.ChangeInfo, error) {
  100. return mm.change(func(c *cachedCollection) (*mgo.ChangeInfo, error) {
  101. return c.RemoveAll(selector, keys...)
  102. })
  103. }
  104. func (mm *Model) RemoveAllNoCache(selector interface{}) (*mgo.ChangeInfo, error) {
  105. return mm.change(func(c *cachedCollection) (*mgo.ChangeInfo, error) {
  106. return c.RemoveAllNoCache(selector)
  107. })
  108. }
  109. func (mm *Model) RemoveId(id interface{}, keys ...string) error {
  110. return mm.execute(func(c *cachedCollection) error {
  111. return c.RemoveId(id, keys...)
  112. })
  113. }
  114. func (mm *Model) RemoveIdNoCache(id interface{}) error {
  115. return mm.execute(func(c *cachedCollection) error {
  116. return c.RemoveIdNoCache(id)
  117. })
  118. }
  119. func (mm *Model) SetCache(key string, v interface{}) error {
  120. return mm.cache.Set(key, v)
  121. }
  122. func (mm *Model) Update(selector, update interface{}, keys ...string) error {
  123. return mm.execute(func(c *cachedCollection) error {
  124. return c.Update(selector, update, keys...)
  125. })
  126. }
  127. func (mm *Model) UpdateNoCache(selector, update interface{}) error {
  128. return mm.execute(func(c *cachedCollection) error {
  129. return c.UpdateNoCache(selector, update)
  130. })
  131. }
  132. func (mm *Model) UpdateId(id, update interface{}, keys ...string) error {
  133. return mm.execute(func(c *cachedCollection) error {
  134. return c.UpdateId(id, update, keys...)
  135. })
  136. }
  137. func (mm *Model) UpdateIdNoCache(id, update interface{}) error {
  138. return mm.execute(func(c *cachedCollection) error {
  139. return c.UpdateIdNoCache(id, update)
  140. })
  141. }
  142. func (mm *Model) Upsert(selector, update interface{}, keys ...string) (*mgo.ChangeInfo, error) {
  143. return mm.change(func(c *cachedCollection) (*mgo.ChangeInfo, error) {
  144. return c.Upsert(selector, update, keys...)
  145. })
  146. }
  147. func (mm *Model) UpsertNoCache(selector, update interface{}) (*mgo.ChangeInfo, error) {
  148. return mm.change(func(c *cachedCollection) (*mgo.ChangeInfo, error) {
  149. return c.UpsertNoCache(selector, update)
  150. })
  151. }
  152. func (mm *Model) change(fn func(c *cachedCollection) (*mgo.ChangeInfo, error)) (*mgo.ChangeInfo, error) {
  153. session, err := mm.TakeSession()
  154. if err != nil {
  155. return nil, err
  156. }
  157. defer mm.PutSession(session)
  158. return fn(mm.GetCollection(session))
  159. }
  160. func (mm *Model) execute(fn func(c *cachedCollection) error) error {
  161. session, err := mm.TakeSession()
  162. if err != nil {
  163. return err
  164. }
  165. defer mm.PutSession(session)
  166. return fn(mm.GetCollection(session))
  167. }
  168. func (mm *Model) executeInt(fn func(c *cachedCollection) (int, error)) (int, error) {
  169. session, err := mm.TakeSession()
  170. if err != nil {
  171. return 0, err
  172. }
  173. defer mm.PutSession(session)
  174. return fn(mm.GetCollection(session))
  175. }
  176. func (mm *Model) pipe(fn func(c *cachedCollection) mongo.Pipe) (mongo.Pipe, error) {
  177. session, err := mm.TakeSession()
  178. if err != nil {
  179. return nil, err
  180. }
  181. defer mm.PutSession(session)
  182. return fn(mm.GetCollection(session)), nil
  183. }
  184. func createModel(url, collection string, c cache.Cache,
  185. create func(mongo.Collection) *cachedCollection) (*Model, error) {
  186. model, err := mongo.NewModel(url, collection)
  187. if err != nil {
  188. return nil, err
  189. }
  190. return &Model{
  191. Model: model,
  192. cache: c,
  193. generateCollection: func(session *mgo.Session) *cachedCollection {
  194. collection := model.GetCollection(session)
  195. return create(collection)
  196. },
  197. }, nil
  198. }