-
Notifications
You must be signed in to change notification settings - Fork 0
/
cursor.go
362 lines (285 loc) · 8.46 KB
/
cursor.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
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
package mongolang
/*
Methods to support MongoDB Cursor.
In keeping with golang requirements, all accessible methods start with uppercase.
Note that a number of calls return a cursor and the actual DB read
does not actually start until a document is returned from a call.
This makes the following call possible:
cursor := db.Coll("podcosts").Find().Sort(bson.M{}).Skip(10).Limit(100)
In the above call note that:
- .Find() must be the first cursor method - it actually creates the cursor struct
- .Sort(), .Skip(), .Limit() can be in any order after the .find()
- no documents have yet been returned and therefore no DB reads have occurred
The following commands will then actually retrieve a document and begin db reads.
Once these cursor methods are called none of the previously mentioned methods can be called
on the same cursor.
- HasNext() - Next() - ForEach()
- ToArray() - Count() - Pretty()
- Close() - IsClosed
HasNext() and Next() will access the next document in the cursor (if one exists).
If no next document exists, the cursor will be closed.
IsClosed() returns true if the cursor is closed.
All others will close the cursor after reading all of the documents for the cursor.
Once the cursor is closed any attempt to use it will cause a panic.
TODO: methods which read the remaining documents should check to see if one was buffered via
a HasNext()
*/
import (
"bytes"
"context"
"encoding/json"
"fmt"
"errors"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/mongo/options"
)
// cursorOkay if cursor is properly related to a Collection
// and the Collection is okay.
// NOTE that this does NOT specifically check that there
// haven't been any errors.
// To do that, check that Err() == nil.
func (c *Cursor) cursorOkay() bool {
if c.Collection == nil {
return false
}
return c.Collection.collOkay()
}
// If there is an error, return it.
// Else returns nil.
func (c *Cursor) Err() error {
if !c.cursorOkay() {
if c.Collection == nil {
return ErrInvalidCursor
}
}
return c.Collection.Err()
}
func (c *Cursor) requireOpenCursor() bool {
var err = c.Err()
if err != nil {
return false
}
if c.IsClosed {
c.setErr(ErrClosedCursor)
return false
}
return true
}
func (c *Cursor) requireOpenFindCursor() bool {
if c.requireOpenCursor() && c.IsFindCursor {
return true
}
c.setErr(ErrNotFindCursor)
return false
}
// setErr if this is a properly established cursor
// and there isn't already an error.
func (c *Cursor) setErr(err error) {
if c.cursorOkay() {
c.Collection.setErr(err)
}
}
// Close closes a cursor
// Note that it's possible to reuse the cursor, though not recommended?
func (c *Cursor) Close() error {
var err error
if !c.IsClosed {
if c.MongoCursor != nil {
err = c.MongoCursor.Close(context.Background())
c.MongoCursor = nil
}
c.IsClosed = true
c.NextDoc = nil
c.Filter = nil
c.FindOptions = options.FindOptions{}
c.AggrPipeline = nil
c.AggrOptions = options.AggregateOptions{}
} else {
err = ErrClosedCursor
}
return err
}
// getMongoCursor ensures that we have an open MongoDB Cursor.
// If the cursor is currently nil, creates a new Find or Aggregate cursor.
func (c *Cursor) getMongoCursor() error {
if c.MongoCursor == nil {
var err error
if c.IsFindCursor {
c.MongoCursor, err = c.Collection.MongoColl.Find(context.Background(), c.Filter, &c.FindOptions)
} else {
c.MongoCursor, err = c.Collection.MongoColl.Aggregate(context.Background(), c.AggrPipeline, &c.AggrOptions)
}
// mark as not closed here so that if error, c.Close() reinitializes cursor
c.IsClosed = false
if err != nil {
c.Close()
c.setErr(err)
return err
}
}
return nil
}
// bufferNext reads the next document, if it exists, into the Cursor buffer.
// Returns true if there is a next document.
// If no next document, automatically closes the cursor.
// Since this is an internal method it will panic if called with a closed cursor.
func (c *Cursor) bufferNext() bool {
if c.IsClosed {
panic("internal error: bufferNext() called after cursor closed")
}
err := c.getMongoCursor()
if err != nil {
return false
}
hasNext := c.MongoCursor.Next(context.Background())
if !hasNext {
c.Close()
return false
}
c.NextDoc = &bson.D{}
err = c.MongoCursor.Decode(c.NextDoc)
if err != nil {
c.Close()
// Close will set c.Err
// Don't lose error from Decode call
c.NextDoc = nil
return false
}
return true
}
// sort, skip, limit - pre-cursor open methods
//
// NOT allowed on aggregate cursor
// Sort specifies the bson.D to be used to sort the cursor results
func (c *Cursor) Sort(sortSequence interface{}) *Cursor {
var err error
if c.requireOpenFindCursor() {
c.FindOptions.Sort, err = verifyParm(sortSequence, (bsonDAllowed | bsonMAllowed))
if err != nil {
c.setErr(err)
}
}
return c
}
// Skip specifies the number of documents to skip before returning the first document
func (c *Cursor) Skip(skipCount int64) *Cursor {
if c.requireOpenFindCursor() {
c.FindOptions.Skip = &skipCount
}
return c
}
// Limit specifies the max number of documents to return
func (c *Cursor) Limit(limitCount int64) *Cursor {
if c.requireOpenFindCursor() {
c.FindOptions.Limit = &limitCount
}
return c
}
/*
HasNext() and Next() - used to read through a cursor.
Closes the cursor if no next document.
*/
// HasNext returns true if the cursor has a next document available.
func (c *Cursor) HasNext() bool {
if !c.requireOpenCursor() {
c.NextDoc = nil
return false
}
if c.NextDoc != nil {
return true
}
return c.bufferNext()
}
// Next returns the next document for the cursor as a bson.D struct.
// TODO: allow a struct to be passed similar to cursor.ToArray(...)
func (c *Cursor) Next() *bson.D {
if !c.requireOpenCursor() {
return &bson.D{}
}
if c.NextDoc == nil {
hasNext := c.bufferNext()
if !hasNext {
c.setErr(errors.New("Next() called when there isn't a next document"))
return &bson.D{}
}
}
// "unbuffer" the next document
doc := c.NextDoc
c.NextDoc = nil
return doc
}
/*
Read all of the documents for a cursor then close the cursor.
- ForEach() - ToArray() - Count() - Pretty()
- String() (fulfills the Stringer interface for printing, etc.)
*/
// ForEach calls the specified function once for each remaining cursor document
// passing the function a bson.D document.
func (c *Cursor) ForEach(f func(*bson.D)) {
fmt.Println("ForEach(...) not yet implemented")
}
// ToArray returns all of the remaining documents for a cursor
// in a bson.D slice. NOTE: currently seems to return all docs
// even those already read via a cursor.Next() call.
// Optional parm is a pointer to a slice which typically would contain
// a custom struct or bson.A struct. In this case, ToArray returns an
// empty []bson.D slice.
// This may change at some future date but currently it is difficult
// to deal with a slice of any type. For consistency with Mongo Shell
// ToArray should return a slice. However we also need the ability
// to return all of the documents to a slice which is a custom struct
// or a bson.A struct.
func (c *Cursor) ToArray(parm ...interface{}) []bson.D {
result := []bson.D{}
if !c.requireOpenCursor() {
return result
}
err := c.getMongoCursor()
if err != nil {
return result
}
if len(parm) > 0 {
err = c.MongoCursor.All(context.Background(), parm[0])
} else {
err = c.MongoCursor.All(context.Background(), &result)
}
c.MongoCursor = nil
c.Close()
c.setErr(err)
return result
}
// Count returns a count of the (remaining) documents for the cursor.
func (c *Cursor) Count() int {
if !c.requireOpenCursor() {
return 0
}
return len(c.ToArray())
}
// Pretty returns a pretty string version of the remaining documents for a cursor.
// Unlike String() it shows the bson.D as key:value instead of {Key:key Value:value}
func (c *Cursor) Pretty() string {
if !c.requireOpenCursor() {
return ""
}
var buf bytes.Buffer
docs := c.ToArray()
if c.Err() != nil {
return ""
}
for _, doc := range docs {
buf.WriteString("{ \n")
for _, v := range doc {
fmt.Fprintf(&buf, " %s : %v \n", v.Key, v.Value)
}
buf.WriteString("} \n")
}
return buf.String()
}
// String fulfills the Stringer interface.
// Calling this will return a string containing the "pretty" print
// contents of the ToArray() function. ToArray() returns an array
// with all of the documents remaining for the cursor and closes the cursor.
func (c *Cursor) String() string {
json, _ := json.MarshalIndent(c.ToArray(), "", " ")
return string(json)
}