forked from scylladb/gocqlx
/
iterx.go
271 lines (230 loc) · 6.79 KB
/
iterx.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
// Copyright (C) 2017 ScyllaDB
// Use of this source code is governed by a ALv2-style
// license that can be found in the LICENSE file.
package gocqlx
import (
"errors"
"fmt"
"reflect"
"github.com/gocql/gocql"
"github.com/scylladb/gocqlx/reflectx"
)
// Get is a convenience function for creating iterator and calling Get.
//
// DEPRECATED use Queryx.Get or Queryx.GetRelease.
func Get(dest interface{}, q *gocql.Query) error {
return Iter(q).Get(dest)
}
// Select is a convenience function for creating iterator and calling Select.
//
// DEPRECATED use Queryx.Select or Queryx.SelectRelease.
func Select(dest interface{}, q *gocql.Query) error {
return Iter(q).Select(dest)
}
// Iterx is a wrapper around gocql.Iter which adds struct scanning capabilities.
type Iterx struct {
*gocql.Iter
err error
unsafe bool
Mapper *reflectx.Mapper
// these fields cache memory use for a rows during iteration w/ structScan
started bool
fields [][]int
values []interface{}
}
// Iter creates a new Iterx from gocql.Query using a default mapper.
func Iter(q *gocql.Query) *Iterx {
return &Iterx{
Iter: q.Iter(),
Mapper: DefaultMapper,
}
}
// Unsafe forces the iterator to ignore missing fields. By default when scanning
// a struct if result row has a column that cannot be mapped to any destination
// field an error is reported. With unsafe such columns are ignored.
func (iter *Iterx) Unsafe() *Iterx {
iter.unsafe = true
return iter
}
// Get scans first row into a destination and closes the iterator. If the
// destination type is a struct pointer, then StructScan will be used.
// If the destination is some other type, then the row must only have one column
// which can scan into that type.
//
// If no rows were selected, ErrNotFound is returned.
func (iter *Iterx) Get(dest interface{}) error {
iter.scanAny(dest, false)
iter.Close()
return iter.checkErrAndNotFound()
}
func (iter *Iterx) scanAny(dest interface{}, structOnly bool) bool {
value := reflect.ValueOf(dest)
if value.Kind() != reflect.Ptr {
iter.err = errors.New("must pass a pointer, not a value, to StructScan destination")
return false
}
if value.IsNil() {
iter.err = errors.New("nil pointer passed to StructScan destination")
return false
}
// no results or query error
if iter.Iter.NumRows() == 0 {
return false
}
base := reflectx.Deref(value.Type())
scannable := isScannable(base)
if structOnly && scannable {
iter.err = structOnlyError(base)
return false
}
if scannable && len(iter.Columns()) > 1 {
iter.err = fmt.Errorf("scannable dest type %s with >1 columns (%d) in result", base.Kind(), len(iter.Columns()))
return false
}
if scannable {
return iter.Scan(dest)
}
return iter.StructScan(dest)
}
// Select scans all rows into a destination, which must be a pointer to slice
// of any type and closes the iterator. If the destination slice type is
// a struct, then StructScan will be used on each row. If the destination is
// some other type, then each row must only have one column which can scan into
// that type.
//
// If no rows were selected, ErrNotFound is NOT returned.
func (iter *Iterx) Select(dest interface{}) error {
iter.scanAll(dest, false)
iter.Close()
return iter.err
}
func (iter *Iterx) scanAll(dest interface{}, structOnly bool) bool {
value := reflect.ValueOf(dest)
// json.Unmarshal returns errors for these
if value.Kind() != reflect.Ptr {
iter.err = errors.New("must pass a pointer, not a value, to StructScan destination")
return false
}
if value.IsNil() {
iter.err = errors.New("nil pointer passed to StructScan destination")
return false
}
// no results or query error
if iter.Iter.NumRows() == 0 {
return false
}
slice, err := baseType(value.Type(), reflect.Slice)
if err != nil {
iter.err = err
return false
}
isPtr := slice.Elem().Kind() == reflect.Ptr
base := reflectx.Deref(slice.Elem())
scannable := isScannable(base)
if structOnly && scannable {
iter.err = structOnlyError(base)
return false
}
// if it's a base type make sure it only has 1 column; if not return an error
if scannable && len(iter.Columns()) > 1 {
iter.err = fmt.Errorf("non-struct dest type %s with >1 columns (%d)", base.Kind(), len(iter.Columns()))
return false
}
var (
alloc bool
v reflect.Value
vp reflect.Value
ok bool
)
for {
// create a new struct type (which returns PtrTo) and indirect it
vp = reflect.New(base)
// scan into the struct field pointers
if !scannable {
ok = iter.StructScan(vp.Interface())
} else {
ok = iter.Scan(vp.Interface())
}
if !ok {
break
}
// allocate memory for the page data
if !alloc {
v = reflect.MakeSlice(slice, 0, iter.Iter.NumRows())
alloc = true
}
if isPtr {
v = reflect.Append(v, vp)
} else {
v = reflect.Append(v, reflect.Indirect(vp))
}
}
// update dest if allocated slice
if alloc {
reflect.Indirect(value).Set(v)
}
return true
}
// StructScan is like gocql.Iter.Scan, but scans a single row into a single
// struct. Use this and iterate manually when the memory load of Select() might
// be prohibitive. StructScan caches the reflect work of matching up column
// positions to fields to avoid that overhead per scan, which means it is not
// safe to run StructScan on the same Iterx instance with different struct
// types.
func (iter *Iterx) StructScan(dest interface{}) bool {
v := reflect.ValueOf(dest)
if v.Kind() != reflect.Ptr {
iter.err = errors.New("must pass a pointer, not a value, to StructScan destination")
return false
}
// no results or query error
if iter.Iter.NumRows() == 0 {
return false
}
if !iter.started {
columns := columnNames(iter.Iter.Columns())
m := iter.Mapper
iter.fields = m.TraversalsByName(v.Type(), columns)
// if we are not unsafe and are missing fields, return an error
if !iter.unsafe {
if f, err := missingFields(iter.fields); err != nil {
iter.err = fmt.Errorf("missing destination name %q in %T", columns[f], dest)
return false
}
}
iter.values = make([]interface{}, len(columns))
iter.started = true
}
err := fieldsByTraversal(v, iter.fields, iter.values, true)
if err != nil {
iter.err = err
return false
}
// scan into the struct field pointers and append to our results
return iter.Iter.Scan(iter.values...)
}
func columnNames(ci []gocql.ColumnInfo) []string {
r := make([]string, len(ci))
for i, column := range ci {
r[i] = column.Name
}
return r
}
// Close closes the iterator and returns any errors that happened during
// the query or the iteration.
func (iter *Iterx) Close() error {
err := iter.Iter.Close()
if iter.err == nil {
iter.err = err
}
return iter.err
}
// checkErrAndNotFound handle error and NotFound in one method.
func (iter *Iterx) checkErrAndNotFound() error {
if iter.err != nil {
return iter.err
} else if iter.Iter.NumRows() == 0 {
return gocql.ErrNotFound
}
return nil
}