-
Notifications
You must be signed in to change notification settings - Fork 3.8k
/
command_result.go
424 lines (380 loc) · 12.7 KB
/
command_result.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
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
// Copyright 2018 The Cockroach Authors.
//
// Use of this software is governed by the Business Source License
// included in the file licenses/BSL.txt.
//
// As of the Change Date specified in that file, in accordance with
// the Business Source License, use of this software will be governed
// by the Apache License, Version 2.0, included in the file
// licenses/APL.txt.
package pgwire
import (
"context"
"fmt"
"github.com/cockroachdb/cockroach/pkg/sql"
"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgwirebase"
"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
"github.com/cockroachdb/errors"
"github.com/lib/pq/oid"
)
type completionMsgType int
const (
_ completionMsgType = iota
commandComplete
bindComplete
closeComplete
parseComplete
emptyQueryResponse
readyForQuery
flush
// Some commands, like Describe, don't need a completion message.
noCompletionMsg
)
// commandResult is an implementation of sql.CommandResult that streams a
// commands results over a pgwire network connection.
type commandResult struct {
// conn is the parent connection of this commandResult.
conn *conn
// conv indicates the conversion settings for SQL values.
conv sessiondata.DataConversionConfig
// pos identifies the position of the command within the connection.
pos sql.CmdPos
err error
// errExpected, if set, enforces that an error had been set when Close is
// called.
errExpected bool
typ completionMsgType
// If typ == commandComplete, this is the tag to be written in the
// CommandComplete message.
cmdCompleteTag string
stmtType tree.StatementType
descOpt sql.RowDescOpt
rowsAffected int
// formatCodes describes the encoding of each column of result rows. It is nil
// for statements not returning rows (or for results for commands other than
// executing statements). It can also be nil for queries returning rows,
// meaning that all columns will be encoded in the text format (this is the
// case for queries executed through the simple protocol). Otherwise, it needs
// to have an entry for every column.
formatCodes []pgwirebase.FormatCode
// oids is a map from result column index to its Oid, similar to formatCodes
// (except oids must always be set).
oids []oid.Oid
// bufferingDisabled is conditionally set during planning of certain
// statements.
bufferingDisabled bool
// released is set when the command result has been released so that its
// memory can be reused. It is also used to assert against use-after-free
// errors.
released bool
}
// Close is part of the CommandResult interface.
func (r *commandResult) Close(t sql.TransactionStatusIndicator) {
r.assertNotReleased()
defer r.release()
if r.errExpected && r.err == nil {
panic("expected err to be set on result by Close, but wasn't")
}
r.conn.writerState.fi.registerCmd(r.pos)
if r.err != nil {
r.conn.bufferErr(r.err)
return
}
// Send a completion message, specific to the type of result.
switch r.typ {
case commandComplete:
tag := cookTag(
r.cmdCompleteTag, r.conn.writerState.tagBuf[:0], r.stmtType, r.rowsAffected,
)
r.conn.bufferCommandComplete(tag)
case parseComplete:
r.conn.bufferParseComplete()
case bindComplete:
r.conn.bufferBindComplete()
case closeComplete:
r.conn.bufferCloseComplete()
case readyForQuery:
r.conn.bufferReadyForQuery(byte(t))
// The error is saved on conn.err.
_ /* err */ = r.conn.Flush(r.pos)
case emptyQueryResponse:
r.conn.bufferEmptyQueryResponse()
case flush:
// The error is saved on conn.err.
_ /* err */ = r.conn.Flush(r.pos)
case noCompletionMsg:
// nothing to do
default:
panic(fmt.Sprintf("unknown type: %v", r.typ))
}
}
// CloseWithErr is part of the CommandResult interface.
func (r *commandResult) CloseWithErr(err error) {
r.assertNotReleased()
defer r.release()
if r.err != nil {
panic(fmt.Sprintf("can't overwrite err: %s with err: %s", r.err, err))
}
r.conn.writerState.fi.registerCmd(r.pos)
r.conn.bufferErr(err)
}
// Discard is part of the CommandResult interface.
func (r *commandResult) Discard() {
r.assertNotReleased()
defer r.release()
}
// Err is part of the CommandResult interface.
func (r *commandResult) Err() error {
r.assertNotReleased()
return r.err
}
// SetError is part of the CommandResult interface.
//
// We're not going to write any bytes to the buffer in order to support future
// SetError() calls. The error will only be serialized at Close() time.
func (r *commandResult) SetError(err error) {
r.assertNotReleased()
r.err = err
}
// AddRow is part of the CommandResult interface.
func (r *commandResult) AddRow(ctx context.Context, row tree.Datums) error {
r.assertNotReleased()
if r.err != nil {
panic(fmt.Sprintf("can't call AddRow after having set error: %s",
r.err))
}
r.conn.writerState.fi.registerCmd(r.pos)
if err := r.conn.GetErr(); err != nil {
return err
}
if r.err != nil {
panic("can't send row after error")
}
r.rowsAffected++
r.conn.bufferRow(ctx, row, r.formatCodes, r.conv, r.oids)
var err error
if r.bufferingDisabled {
err = r.conn.Flush(r.pos)
} else {
_ /* flushed */, err = r.conn.maybeFlush(r.pos)
}
return err
}
// DisableBuffering is part of the CommandResult interface.
func (r *commandResult) DisableBuffering() {
r.assertNotReleased()
r.bufferingDisabled = true
}
// SetColumns is part of the CommandResult interface.
func (r *commandResult) SetColumns(ctx context.Context, cols sqlbase.ResultColumns) {
r.assertNotReleased()
r.conn.writerState.fi.registerCmd(r.pos)
if r.descOpt == sql.NeedRowDesc {
_ /* err */ = r.conn.writeRowDescription(ctx, cols, r.formatCodes, &r.conn.writerState.buf)
}
r.oids = make([]oid.Oid, len(cols))
for i, col := range cols {
r.oids[i] = mapResultOid(col.Typ.Oid())
}
}
// SetInferredTypes is part of the DescribeResult interface.
func (r *commandResult) SetInferredTypes(types []oid.Oid) {
r.assertNotReleased()
r.conn.writerState.fi.registerCmd(r.pos)
r.conn.bufferParamDesc(types)
}
// SetNoDataRowDescription is part of the DescribeResult interface.
func (r *commandResult) SetNoDataRowDescription() {
r.assertNotReleased()
r.conn.writerState.fi.registerCmd(r.pos)
r.conn.bufferNoDataMsg()
}
// SetPrepStmtOutput is part of the DescribeResult interface.
func (r *commandResult) SetPrepStmtOutput(ctx context.Context, cols sqlbase.ResultColumns) {
r.assertNotReleased()
r.conn.writerState.fi.registerCmd(r.pos)
_ /* err */ = r.conn.writeRowDescription(ctx, cols, nil /* formatCodes */, &r.conn.writerState.buf)
}
// SetPortalOutput is part of the DescribeResult interface.
func (r *commandResult) SetPortalOutput(
ctx context.Context, cols sqlbase.ResultColumns, formatCodes []pgwirebase.FormatCode,
) {
r.assertNotReleased()
r.conn.writerState.fi.registerCmd(r.pos)
_ /* err */ = r.conn.writeRowDescription(ctx, cols, formatCodes, &r.conn.writerState.buf)
}
// IncrementRowsAffected is part of the CommandResult interface.
func (r *commandResult) IncrementRowsAffected(n int) {
r.assertNotReleased()
r.rowsAffected += n
}
// RowsAffected is part of the CommandResult interface.
func (r *commandResult) RowsAffected() int {
r.assertNotReleased()
return r.rowsAffected
}
// ResetStmtType is part of the CommandResult interface.
func (r *commandResult) ResetStmtType(stmt tree.Statement) {
r.assertNotReleased()
r.stmtType = stmt.StatementType()
r.cmdCompleteTag = stmt.StatementTag()
}
// release frees the commandResult and allows its memory to be reused.
func (r *commandResult) release() {
*r = commandResult{released: true}
}
// assertNotReleased asserts that the commandResult is not being used after
// being freed by one of the methods in the CommandResultClose interface. The
// assertion can have false negatives, where it fails to detect a use-after-free
// condition, but will never result in a false positive.
func (r *commandResult) assertNotReleased() {
if r.released {
panic("commandResult used after being released")
}
}
func (c *conn) allocCommandResult() *commandResult {
r := &c.res
if r.released {
r.released = false
} else {
// In practice, each conn only ever uses a single commandResult at a
// time, so we could make this panic. However, doing so would entail
// complicating the ClientComm interface, which doesn't seem worth it.
r = new(commandResult)
}
return r
}
func (c *conn) newCommandResult(
descOpt sql.RowDescOpt,
pos sql.CmdPos,
stmt tree.Statement,
formatCodes []pgwirebase.FormatCode,
conv sessiondata.DataConversionConfig,
limit int,
portalName string,
implicitTxn bool,
) sql.CommandResult {
r := c.allocCommandResult()
*r = commandResult{
conn: c,
conv: conv,
pos: pos,
typ: commandComplete,
cmdCompleteTag: stmt.StatementTag(),
stmtType: stmt.StatementType(),
descOpt: descOpt,
formatCodes: formatCodes,
}
if limit == 0 {
return r
}
return &limitedCommandResult{
limit: limit,
portalName: portalName,
implicitTxn: implicitTxn,
commandResult: r,
}
}
func (c *conn) newMiscResult(pos sql.CmdPos, typ completionMsgType) *commandResult {
r := c.allocCommandResult()
*r = commandResult{
conn: c,
pos: pos,
typ: typ,
}
return r
}
// limitedCommandResult is a commandResult that has a limit, after which calls
// to AddRow will block until the associated client connection asks for more
// rows. It essentially implements the "execute portal with limit" part of the
// Postgres protocol.
//
// This design is known to be flawed. It only supports a specific subset of the
// protocol. We only allow a portal suspension in an explicit transaction where
// the suspended portal is completely exhausted before any other pgwire command
// is executed, otherwise an error is produced. You cannot, for example,
// interleave portal executions (a portal must be executed to completion before
// another can be executed). It also breaks the software layering by adding an
// additional state machine here, instead of teaching the state machine in the
// sql package about portals.
//
// This has been done because refactoring the executor to be able to correctly
// suspend a portal will require a lot of work, and we wanted to move
// forward. The work included is things like auditing all of the defers and
// post-execution stuff (like stats collection) to have it only execute once
// per statement instead of once per portal.
type limitedCommandResult struct {
*commandResult
portalName string
implicitTxn bool
seenTuples int
// If set, an error will be sent to the client if more rows are produced than
// this limit.
limit int
}
// AddRow is part of the CommandResult interface.
func (r *limitedCommandResult) AddRow(ctx context.Context, row tree.Datums) error {
if err := r.commandResult.AddRow(ctx, row); err != nil {
return err
}
r.seenTuples++
if r.seenTuples == r.limit {
// If we've seen up to the limit of rows, send a "portal suspended" message
// and wait for another exec portal message.
r.conn.bufferPortalSuspended()
if err := r.conn.Flush(r.pos); err != nil {
return err
}
r.seenTuples = 0
return r.moreResultsNeeded()
}
if _ /* flushed */, err := r.conn.maybeFlush(r.pos); err != nil {
return err
}
return nil
}
// moreResultsNeeded is a restricted connection handler that waits for more
// requests for rows from the active portal, during the "execute portal" flow
// when a limit has been specified.
func (r *limitedCommandResult) moreResultsNeeded() error {
// In an implicit transaction, we're supposed to suspend then close the
// portal after the first exec message. We don't support that yet, so
// just error.
if r.implicitTxn {
return errors.WithDetail(sql.LimitedResultErr, "implicit transactions not supported")
}
r.conn.stmtBuf.AdvanceOne()
for {
cmd, _, err := r.conn.stmtBuf.CurCmd()
if err != nil {
return err
}
switch c := cmd.(type) {
case sql.ExecPortal:
// The happy case: the client wants more rows from the portal.
if c.Name != r.portalName {
return errors.WithDetail(sql.LimitedResultErr, "portals must be executed to completion")
}
r.limit = c.Limit
// In order to get the correct command tag, we need to reset the seen rows.
r.rowsAffected = 0
return nil
case sql.Sync:
// The client wants to see a ready for query message
// back. Send it then run the for loop again.
r.conn.stmtBuf.AdvanceOne()
// We can hard code InTxnBlock here because we don't
// support implicit transactions, so we know we're in
// a transaction.
r.conn.bufferReadyForQuery(byte(sql.InTxnBlock))
if err := r.conn.Flush(r.pos); err != nil {
return err
}
default:
// We got some other message, but we only support executing to completion.
return errors.WithDetail(sql.LimitedResultErr, "portals must be executed to completion")
}
}
}