-
Notifications
You must be signed in to change notification settings - Fork 0
/
stmt.go
125 lines (103 loc) · 2.45 KB
/
stmt.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
package bigquery
import (
"context"
"database/sql/driver"
"cloud.google.com/go/bigquery"
)
var (
_ driver.Stmt = (*stmt)(nil)
_ driver.StmtExecContext = (*stmt)(nil)
_ driver.StmtQueryContext = (*stmt)(nil)
)
type stmt struct {
conn *conn
query string
}
func (s *stmt) Close() error {
return nil
}
func (s *stmt) NumInput() int {
return -1
}
func (s *stmt) Exec(args []driver.Value) (driver.Result, error) {
return nil, driver.ErrSkip
}
func (s *stmt) Query(args []driver.Value) (driver.Rows, error) {
return nil, driver.ErrSkip
}
func (s *stmt) ExecContext(ctx context.Context, args []driver.NamedValue) (driver.Result, error) {
iterator, err := s.iterator(ctx, args)
if err != nil {
return nil, err
}
return &result{
iterator: iterator,
}, nil
}
func (s *stmt) QueryContext(ctx context.Context, args []driver.NamedValue) (driver.Rows, error) {
iterator, err := s.iterator(ctx, args)
if err != nil {
return nil, err
}
return &rows{
iterator: iterator,
}, nil
}
func (s *stmt) iterator(ctx context.Context, args []driver.NamedValue) (*bigquery.RowIterator, error) {
query := s.buildQuery(args)
s.conn.getQueryOpt(query)
job, err := query.Run(ctx)
if err != nil {
return nil, err
}
s.conn.getJobOpt(job)
if query.DryRun {
return nil, nil
}
if sessionID := getSessionID(job); sessionID != "" {
s.conn.sessionID = sessionID
}
return job.Read(ctx)
}
func getSessionID(job *bigquery.Job) string {
status := job.LastStatus()
if status == nil {
return ""
}
if status.Statistics == nil {
return ""
}
if status.Statistics.SessionInfo == nil {
return ""
}
return status.Statistics.SessionInfo.SessionID
}
func (s *stmt) buildQuery(args []driver.NamedValue) *bigquery.Query {
query := s.conn.client.Query(s.query)
query.DefaultDatasetID = s.conn.config.Dataset
query.Parameters = s.buildParameters(args)
query.ConnectionProperties = s.buildConnectionProperties()
query.CreateSession = s.conn.sessionID == ""
return query
}
func (s *stmt) buildParameters(args []driver.NamedValue) []bigquery.QueryParameter {
params := make([]bigquery.QueryParameter, len(args))
for i, arg := range args {
params[i] = bigquery.QueryParameter{
Name: arg.Name,
Value: arg.Value,
}
}
return params
}
func (s *stmt) buildConnectionProperties() []*bigquery.ConnectionProperty {
if s.conn.sessionID == "" {
return nil
}
return []*bigquery.ConnectionProperty{
{
Key: "session_id",
Value: s.conn.sessionID,
},
}
}