forked from vitessio/vitess
-
Notifications
You must be signed in to change notification settings - Fork 13
/
streaming_rows.go
110 lines (98 loc) · 2.25 KB
/
streaming_rows.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
// Copyright 2015, Google Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package client
import (
"database/sql/driver"
"errors"
"io"
"golang.org/x/net/context"
mproto "github.com/youtube/vitess/go/mysql/proto"
"github.com/youtube/vitess/go/vt/vtgate/vtgateconn"
)
// streamingRows creates a database/sql/driver compliant Row iterator
// for a streaming query.
type streamingRows struct {
qrc <-chan *mproto.QueryResult
errFunc vtgateconn.ErrFunc
failed error
fields []mproto.Field
qr *mproto.QueryResult
index int
cancel context.CancelFunc
}
// newStreamingRows creates a new streamingRows from qrc and errFunc.
func newStreamingRows(qrc <-chan *mproto.QueryResult, errFunc vtgateconn.ErrFunc, cancel context.CancelFunc) driver.Rows {
return &streamingRows{
qrc: qrc,
errFunc: errFunc,
cancel: cancel,
}
}
func (ri *streamingRows) Columns() []string {
ri.checkFields()
if ri.failed != nil {
return nil
}
cols := make([]string, 0, len(ri.fields))
for _, field := range ri.fields {
cols = append(cols, field.Name)
}
return cols
}
func (ri *streamingRows) Close() error {
return nil
}
func (ri *streamingRows) Next(dest []driver.Value) error {
ri.checkFields()
if ri.failed != nil {
return ri.failed
}
for ri.qr == nil || ri.index == len(ri.qr.Rows) {
ri.fetchNext()
if ri.failed != nil {
return ri.failed
}
ri.index = 0
}
err := populateRow(dest, ri.fields, ri.qr.Rows[ri.index])
ri.index++
return err
}
// checkFields fetches the first packet from the channel, which
// should contain the field info. It sets ri.failed if it fails.
func (ri *streamingRows) checkFields() {
if ri.failed != nil {
return
}
if ri.fields != nil {
return
}
qr, ok := <-ri.qrc
if !ok {
ri.setErr()
return
}
ri.fields = qr.Fields
if ri.fields == nil {
ri.failed = errors.New("first packet did not return fields")
}
}
// fetchNext fetches the next packet from the channel.
func (ri *streamingRows) fetchNext() {
qr, ok := <-ri.qrc
if !ok {
ri.setErr()
return
}
ri.qr = qr
}
func (ri *streamingRows) setErr() {
ri.failed = io.EOF
if err := ri.errFunc(); err != nil {
ri.failed = err
}
if ri.cancel != nil {
ri.cancel()
}
}