-
-
Notifications
You must be signed in to change notification settings - Fork 3.9k
/
model.go
177 lines (148 loc) · 4.29 KB
/
model.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
package mongo
import (
"log"
"time"
"github.com/globalsign/mgo"
"github.com/zeromicro/go-zero/core/breaker"
)
// A Model is a mongo model.
type Model struct {
session *concurrentSession
db *mgo.Database
collection string
brk breaker.Breaker
opts []Option
}
// MustNewModel returns a Model, exits on errors.
func MustNewModel(url, collection string, opts ...Option) *Model {
model, err := NewModel(url, collection, opts...)
if err != nil {
log.Fatal(err)
}
return model
}
// NewModel returns a Model.
func NewModel(url, collection string, opts ...Option) (*Model, error) {
session, err := getConcurrentSession(url)
if err != nil {
return nil, err
}
return &Model{
session: session,
// If name is empty, the database name provided in the dialed URL is used instead
db: session.DB(""),
collection: collection,
brk: breaker.GetBreaker(url),
opts: opts,
}, nil
}
// Find finds a record with given query.
func (mm *Model) Find(query interface{}) (Query, error) {
return mm.query(func(c Collection) Query {
return c.Find(query)
})
}
// FindId finds a record with given id.
func (mm *Model) FindId(id interface{}) (Query, error) {
return mm.query(func(c Collection) Query {
return c.FindId(id)
})
}
// GetCollection returns a Collection with given session.
func (mm *Model) GetCollection(session *mgo.Session) Collection {
return newCollection(mm.db.C(mm.collection).With(session), mm.brk)
}
// Insert inserts docs into mm.
func (mm *Model) Insert(docs ...interface{}) error {
return mm.execute(func(c Collection) error {
return c.Insert(docs...)
})
}
// Pipe returns a Pipe with given pipeline.
func (mm *Model) Pipe(pipeline interface{}) (Pipe, error) {
return mm.pipe(func(c Collection) Pipe {
return c.Pipe(pipeline)
})
}
// PutSession returns the given session.
func (mm *Model) PutSession(session *mgo.Session) {
mm.session.putSession(session)
}
// Remove removes the records with given selector.
func (mm *Model) Remove(selector interface{}) error {
return mm.execute(func(c Collection) error {
return c.Remove(selector)
})
}
// RemoveAll removes all with given selector and returns a mgo.ChangeInfo.
func (mm *Model) RemoveAll(selector interface{}) (*mgo.ChangeInfo, error) {
return mm.change(func(c Collection) (*mgo.ChangeInfo, error) {
return c.RemoveAll(selector)
})
}
// RemoveId removes a record with given id.
func (mm *Model) RemoveId(id interface{}) error {
return mm.execute(func(c Collection) error {
return c.RemoveId(id)
})
}
// TakeSession gets a session.
func (mm *Model) TakeSession() (*mgo.Session, error) {
return mm.session.takeSession(mm.opts...)
}
// Update updates a record with given selector.
func (mm *Model) Update(selector, update interface{}) error {
return mm.execute(func(c Collection) error {
return c.Update(selector, update)
})
}
// UpdateId updates a record with given id.
func (mm *Model) UpdateId(id, update interface{}) error {
return mm.execute(func(c Collection) error {
return c.UpdateId(id, update)
})
}
// Upsert upserts a record with given selector, and returns a mgo.ChangeInfo.
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
}
// WithTimeout customizes an operation with given timeout.
func WithTimeout(timeout time.Duration) Option {
return func(opts *options) {
opts.timeout = timeout
}
}