-
Notifications
You must be signed in to change notification settings - Fork 1
/
query.go
162 lines (139 loc) · 3.38 KB
/
query.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
package astra
import (
"fmt"
pb "github.com/stargate/stargate-grpc-go-client/stargate/pkg/proto"
"google.golang.org/protobuf/types/known/wrapperspb"
)
type BatchType uint8
// Batch types for BatchQuery.
// See https://docs.datastax.com/en/cql-oss/3.x/cql/cql_reference/cqlBatch.html
const (
BatchLogged BatchType = iota
BatchUnlogged
BatchCounter
)
type params struct {
keyspace string
}
func (p *params) toQueryParamsProto() *pb.QueryParameters {
if p == nil {
return nil
}
res := &pb.QueryParameters{}
if p.keyspace != "" {
res.Keyspace = &wrapperspb.StringValue{Value: p.keyspace}
}
return res
}
func (p *params) toBatchParamsProto() *pb.BatchParameters {
if p == nil {
return nil
}
res := &pb.BatchParameters{}
if p.keyspace != "" {
res.Keyspace = &wrapperspb.StringValue{Value: p.keyspace}
}
return res
}
type queryParams struct {
params *params
}
func (p *queryParams) createIfEmpty() {
if p.params == nil {
p.params = ¶ms{}
}
}
func (p *queryParams) keyspace(value string) {
p.createIfEmpty()
p.params.keyspace = value
}
// Query is a configurable and executable Stargate query. Use Client.Query to
// create a Query.
type Query struct {
client *Client
cql string
values []any
queryParams
}
// Keyspace sets the keyspace to use for the query.
func (q *Query) Keyspace(value string) *Query {
q.queryParams.keyspace(value)
return q
}
// Exec executes the Query using the client that created it and returns the
// resultant rows.
func (q *Query) Exec() (Rows, error) {
return q.client.execQuery(q)
}
func (q *Query) toQueryProto() (*pb.Query, error) {
vs, err := valuesToProto(q.values)
if err != nil {
return nil, fmt.Errorf("failed to convert values to proto: %v", err)
}
return &pb.Query{
Cql: q.cql,
Values: &pb.Values{Values: vs},
Parameters: q.queryParams.params.toQueryParamsProto(),
}, nil
}
func (q *Query) toBatchQueryProto() (*pb.BatchQuery, error) {
vs, err := valuesToProto(q.values)
if err != nil {
return nil, fmt.Errorf("failed to convert values to proto: %v", err)
}
return &pb.BatchQuery{
Cql: q.cql,
Values: &pb.Values{Values: vs},
}, nil
}
// BatchQuery is a configurable and executable Stargate batch query. Use
// Client.Batch to create a BatchQuery.
type BatchQuery struct {
client *Client
batchType BatchType
queries []*Query
queryParams
}
// BatchType sets the batch type to use for the batch query.
func (b *BatchQuery) BatchType(batchType BatchType) *BatchQuery {
b.batchType = batchType
return b
}
// Keyspace sets the keyspace to use for the batch query.
func (b *BatchQuery) Keyspace(value string) *BatchQuery {
b.queryParams.keyspace(value)
return b
}
func (b *BatchQuery) toProto() (*pb.Batch, error) {
qs := make([]*pb.BatchQuery, len(b.queries))
for i, q := range b.queries {
bq, err := q.toBatchQueryProto()
if err != nil {
return nil, err
}
qs[i] = bq
}
var t pb.Batch_Type
switch b.batchType {
case BatchLogged:
t = pb.Batch_LOGGED
case BatchUnlogged:
t = pb.Batch_UNLOGGED
case BatchCounter:
t = pb.Batch_COUNTER
default:
return nil, fmt.Errorf("unknown batch type: %v", b.batchType)
}
res := &pb.Batch{
Queries: qs,
Parameters: b.params.toBatchParamsProto(),
}
if t != pb.Batch_LOGGED {
res.Type = t
}
return res, nil
}
// Exec executes the BatchQuery using the client that created it.
func (b *BatchQuery) Exec() error {
return b.client.execBatch(b)
}