github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/sql/conn_io.go (about)

     1  // Copyright 2017 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package sql
    12  
    13  import (
    14  	"context"
    15  	"fmt"
    16  	"io"
    17  	"sync"
    18  	"time"
    19  
    20  	"github.com/cockroachdb/cockroach/pkg/sql/parser"
    21  	"github.com/cockroachdb/cockroach/pkg/sql/pgwire/pgwirebase"
    22  	"github.com/cockroachdb/cockroach/pkg/sql/sem/tree"
    23  	"github.com/cockroachdb/cockroach/pkg/sql/sessiondata"
    24  	"github.com/cockroachdb/cockroach/pkg/sql/sqlbase"
    25  	"github.com/cockroachdb/cockroach/pkg/util/log"
    26  	"github.com/cockroachdb/cockroach/pkg/util/ring"
    27  	"github.com/cockroachdb/cockroach/pkg/util/syncutil"
    28  	"github.com/cockroachdb/errors"
    29  	"github.com/lib/pq/oid"
    30  )
    31  
    32  // This file contains utils and interfaces used by a connExecutor to communicate
    33  // with a SQL client. There's StmtBuf used for input and ClientComm used for
    34  // output.
    35  
    36  // CmdPos represents the index of a command relative to the start of a
    37  // connection. The first command received on a connection has position 0.
    38  type CmdPos int64
    39  
    40  // TransactionStatusIndicator represents a pg identifier for the transaction state.
    41  type TransactionStatusIndicator byte
    42  
    43  const (
    44  	// IdleTxnBlock means the session is outside of a transaction.
    45  	IdleTxnBlock TransactionStatusIndicator = 'I'
    46  	// InTxnBlock means the session is inside a transaction.
    47  	InTxnBlock TransactionStatusIndicator = 'T'
    48  	// InFailedTxnBlock means the session is inside a transaction, but the
    49  	// transaction is in the Aborted state.
    50  	InFailedTxnBlock TransactionStatusIndicator = 'E'
    51  )
    52  
    53  // StmtBuf maintains a list of commands that a SQL client has sent for execution
    54  // over a network connection. The commands are SQL queries to be executed,
    55  // statements to be prepared, etc. At any point in time the buffer contains
    56  // outstanding commands that have yet to be executed, and it can also contain
    57  // some history of commands that we might want to retry - in the case of a
    58  // retriable error, we'd like to retry all the commands pertaining to the
    59  // current SQL transaction.
    60  //
    61  // The buffer is supposed to be used by one reader and one writer. The writer
    62  // adds commands to the buffer using Push(). The reader reads one command at a
    63  // time using CurCmd(). The consumer is then supposed to create command results
    64  // (the buffer is not involved in this).
    65  // The buffer internally maintains a cursor representing the reader's position.
    66  // The reader has to manually move the cursor using AdvanceOne(),
    67  // seekToNextBatch() and rewind().
    68  // In practice, the writer is a module responsible for communicating with a SQL
    69  // client (i.e. pgwire.conn) and the reader is a connExecutor.
    70  //
    71  // The StmtBuf supports grouping commands into "batches" delimited by sync
    72  // commands. A reader can then at any time chose to skip over commands from the
    73  // current batch. This is used to implement Postgres error semantics: when an
    74  // error happens during processing of a command, some future commands might need
    75  // to be skipped. Batches correspond either to multiple queries received in a
    76  // single query string (when the SQL client sends a semicolon-separated list of
    77  // queries as part of the "simple" protocol), or to different commands pipelined
    78  // by the cliend, separated from "sync" messages.
    79  //
    80  // push() can be called concurrently with CurCmd().
    81  //
    82  // The connExecutor will use the buffer to maintain a window around the
    83  // command it is currently executing. It will maintain enough history for
    84  // executing commands again in case of an automatic retry. The connExecutor is
    85  // in charge of trimming completed commands from the buffer when it's done with
    86  // them.
    87  type StmtBuf struct {
    88  	mu struct {
    89  		syncutil.Mutex
    90  
    91  		// closed, if set, means that the writer has closed the buffer. See Close().
    92  		closed bool
    93  
    94  		// cond is signaled when new commands are pushed.
    95  		cond *sync.Cond
    96  
    97  		// data contains the elements of the buffer.
    98  		data ring.Buffer // []Command
    99  
   100  		// startPos indicates the index of the first command currently in data
   101  		// relative to the start of the connection.
   102  		startPos CmdPos
   103  		// curPos is the current position of the cursor going through the commands.
   104  		// At any time, curPos indicates the position of the command to be returned
   105  		// by CurCmd().
   106  		curPos CmdPos
   107  		// lastPos indicates the position of the last command that was pushed into
   108  		// the buffer.
   109  		lastPos CmdPos
   110  	}
   111  }
   112  
   113  // Command is an interface implemented by all commands pushed by pgwire into the
   114  // buffer.
   115  type Command interface {
   116  	fmt.Stringer
   117  	// command returns a string representation of the command type (e.g.
   118  	// "prepare stmt", "exec stmt").
   119  	command() string
   120  }
   121  
   122  // ExecStmt is the command for running a query sent through the "simple" pgwire
   123  // protocol.
   124  type ExecStmt struct {
   125  	// Information returned from parsing: AST, SQL, NumPlaceholders.
   126  	// Note that AST can be nil, in which case executing it should produce an
   127  	// "empty query response" message.
   128  	parser.Statement
   129  
   130  	// TimeReceived is the time at which the exec message was received
   131  	// from the client. Used to compute the service latency.
   132  	TimeReceived time.Time
   133  	// ParseStart/ParseEnd are the timing info for parsing of the query. Used for
   134  	// stats reporting.
   135  	ParseStart time.Time
   136  	ParseEnd   time.Time
   137  }
   138  
   139  // command implements the Command interface.
   140  func (ExecStmt) command() string { return "exec stmt" }
   141  
   142  func (e ExecStmt) String() string {
   143  	// We have the original SQL, but we still use String() because it obfuscates
   144  	// passwords.
   145  	s := "(empty)"
   146  	// e.AST could be nil in the case of a completely empty query.
   147  	if e.AST != nil {
   148  		s = e.AST.String()
   149  	}
   150  	return fmt.Sprintf("ExecStmt: %s", s)
   151  }
   152  
   153  var _ Command = ExecStmt{}
   154  
   155  // ExecPortal is the Command for executing a portal.
   156  type ExecPortal struct {
   157  	Name string
   158  	// limit is a feature of pgwire that we don't really support. We accept it and
   159  	// don't complain as long as the statement produces fewer results than this.
   160  	Limit int
   161  	// TimeReceived is the time at which the exec message was received
   162  	// from the client. Used to compute the service latency.
   163  	TimeReceived time.Time
   164  }
   165  
   166  // command implements the Command interface.
   167  func (ExecPortal) command() string { return "exec portal" }
   168  
   169  func (e ExecPortal) String() string {
   170  	return fmt.Sprintf("ExecPortal name: %q", e.Name)
   171  }
   172  
   173  var _ Command = ExecPortal{}
   174  
   175  // PrepareStmt is the command for creating a prepared statement.
   176  type PrepareStmt struct {
   177  	// Name of the prepared statement (optional).
   178  	Name string
   179  
   180  	// Information returned from parsing: AST, SQL, NumPlaceholders.
   181  	// Note that AST can be nil, in which case executing it should produce an
   182  	// "empty query response" message.
   183  	parser.Statement
   184  
   185  	TypeHints tree.PlaceholderTypes
   186  	// RawTypeHints is the representation of type hints exactly as specified by
   187  	// the client.
   188  	RawTypeHints []oid.Oid
   189  	ParseStart   time.Time
   190  	ParseEnd     time.Time
   191  }
   192  
   193  // command implements the Command interface.
   194  func (PrepareStmt) command() string { return "prepare stmt" }
   195  
   196  func (p PrepareStmt) String() string {
   197  	// We have the original SQL, but we still use String() because it obfuscates
   198  	// passwords.
   199  	s := "(empty)"
   200  	// p.AST could be nil in the case of a completely empty query.
   201  	if p.AST != nil {
   202  		s = p.AST.String()
   203  	}
   204  	return fmt.Sprintf("PrepareStmt: %s", s)
   205  }
   206  
   207  var _ Command = PrepareStmt{}
   208  
   209  // DescribeStmt is the Command for producing info about a prepared statement or
   210  // portal.
   211  type DescribeStmt struct {
   212  	Name string
   213  	Type pgwirebase.PrepareType
   214  }
   215  
   216  // command implements the Command interface.
   217  func (DescribeStmt) command() string { return "describe stmt" }
   218  
   219  func (d DescribeStmt) String() string {
   220  	return fmt.Sprintf("Describe: %q", d.Name)
   221  }
   222  
   223  var _ Command = DescribeStmt{}
   224  
   225  // BindStmt is the Command for creating a portal from a prepared statement.
   226  type BindStmt struct {
   227  	PreparedStatementName string
   228  	PortalName            string
   229  	// OutFormats contains the requested formats for the output columns.
   230  	// It either contains a bunch of format codes, in which case the number will
   231  	// need to match the number of output columns of the portal, or contains a single
   232  	// code, in which case that code will be applied to all columns.
   233  	OutFormats []pgwirebase.FormatCode
   234  	// Args are the arguments for the prepared statement.
   235  	// They are passed in without decoding because decoding requires type
   236  	// inference to have been performed.
   237  	//
   238  	// A nil element means a tree.DNull argument.
   239  	Args [][]byte
   240  	// ArgFormatCodes are the codes to be used to deserialize the Args.
   241  	// It either contains a bunch of format codes, in which case the number will
   242  	// need to match the number of arguments for the portal, or contains a single
   243  	// code, in which case that code will be applied to all arguments.
   244  	ArgFormatCodes []pgwirebase.FormatCode
   245  
   246  	// internalArgs, if not nil, represents the arguments for the prepared
   247  	// statements as produced by the internal clients. These don't need to go
   248  	// through encoding/decoding of the args. However, the types of the datums
   249  	// must correspond exactly to the inferred types (but note that the types of
   250  	// the datums are passes as type hints to the PrepareStmt command, so the
   251  	// inferred types should reflect that).
   252  	// If internalArgs is specified, Args and ArgFormatCodes are ignored.
   253  	internalArgs []tree.Datum
   254  }
   255  
   256  // command implements the Command interface.
   257  func (BindStmt) command() string { return "bind stmt" }
   258  
   259  func (b BindStmt) String() string {
   260  	return fmt.Sprintf("BindStmt: %q->%q", b.PreparedStatementName, b.PortalName)
   261  }
   262  
   263  var _ Command = BindStmt{}
   264  
   265  // DeletePreparedStmt is the Command for freeing a prepared statement.
   266  type DeletePreparedStmt struct {
   267  	Name string
   268  	Type pgwirebase.PrepareType
   269  }
   270  
   271  // command implements the Command interface.
   272  func (DeletePreparedStmt) command() string { return "delete stmt" }
   273  
   274  func (d DeletePreparedStmt) String() string {
   275  	return fmt.Sprintf("DeletePreparedStmt: %q", d.Name)
   276  }
   277  
   278  var _ Command = DeletePreparedStmt{}
   279  
   280  // Sync is a command that serves two purposes:
   281  // 1) It marks the end of one batch of commands and the beginning of the next.
   282  // stmtBuf.seekToNextBatch will seek to this marker.
   283  // 2) It generates a ReadyForQuery protocol message.
   284  //
   285  // A Sync command is generated for both the simple and the extended pgwire
   286  // protocol variants. So, it doesn't strictly correspond to a pgwire sync
   287  // message - those are not sent in the simple protocol. We synthesize Sync
   288  // commands though because their handling matches the simple protocol too.
   289  type Sync struct{}
   290  
   291  // command implements the Command interface.
   292  func (Sync) command() string { return "sync" }
   293  
   294  func (Sync) String() string {
   295  	return "Sync"
   296  }
   297  
   298  var _ Command = Sync{}
   299  
   300  // Flush is a Command asking for the results of all previous commands to be
   301  // delivered to the client.
   302  type Flush struct{}
   303  
   304  // command implements the Command interface.
   305  func (Flush) command() string { return "flush" }
   306  
   307  func (Flush) String() string {
   308  	return "Flush"
   309  }
   310  
   311  var _ Command = Flush{}
   312  
   313  // CopyIn is the command for execution of the Copy-in pgwire subprotocol.
   314  type CopyIn struct {
   315  	Stmt *tree.CopyFrom
   316  	// Conn is the network connection. Execution of the CopyFrom statement takes
   317  	// control of the connection.
   318  	Conn pgwirebase.Conn
   319  	// CopyDone is decremented once execution finishes, signaling that control of
   320  	// the connection is being handed back to the network routine.
   321  	CopyDone *sync.WaitGroup
   322  }
   323  
   324  // command implements the Command interface.
   325  func (CopyIn) command() string { return "copy" }
   326  
   327  func (CopyIn) String() string {
   328  	return "CopyIn"
   329  }
   330  
   331  var _ Command = CopyIn{}
   332  
   333  // DrainRequest represents a notice that the server is draining and command
   334  // processing should stop soon.
   335  //
   336  // DrainRequest commands don't produce results.
   337  type DrainRequest struct{}
   338  
   339  // command implements the Command interface.
   340  func (DrainRequest) command() string { return "drain" }
   341  
   342  func (DrainRequest) String() string {
   343  	return "Drain"
   344  }
   345  
   346  var _ Command = DrainRequest{}
   347  
   348  // SendError is a command that, upon execution, send a specific error to the
   349  // client. This is used by pgwire to schedule errors to be sent at an
   350  // appropriate time.
   351  type SendError struct {
   352  	// Err is a *pgerror.Error.
   353  	Err error
   354  }
   355  
   356  // command implements the Command interface.
   357  func (SendError) command() string { return "send error" }
   358  
   359  func (s SendError) String() string {
   360  	return fmt.Sprintf("SendError: %s", s.Err)
   361  }
   362  
   363  var _ Command = SendError{}
   364  
   365  // NewStmtBuf creates a StmtBuf.
   366  func NewStmtBuf() *StmtBuf {
   367  	var buf StmtBuf
   368  	buf.Init()
   369  	return &buf
   370  }
   371  
   372  // Init initializes a StmtBuf. It exists to avoid the allocation imposed by
   373  // NewStmtBuf.
   374  func (buf *StmtBuf) Init() {
   375  	buf.mu.lastPos = -1
   376  	buf.mu.cond = sync.NewCond(&buf.mu.Mutex)
   377  }
   378  
   379  // Close marks the buffer as closed. Once Close() is called, no further push()es
   380  // are allowed. If a reader is blocked on a CurCmd() call, it is unblocked with
   381  // io.EOF. Any further CurCmd() call also returns io.EOF (even if some
   382  // commands were already available in the buffer before the Close()).
   383  //
   384  // Close() is idempotent.
   385  func (buf *StmtBuf) Close() {
   386  	buf.mu.Lock()
   387  	buf.mu.closed = true
   388  	buf.mu.cond.Signal()
   389  	buf.mu.Unlock()
   390  }
   391  
   392  // Push adds a Command to the end of the buffer. If a CurCmd() call was blocked
   393  // waiting for this command to arrive, it will be woken up.
   394  //
   395  // An error is returned if the buffer has been closed.
   396  func (buf *StmtBuf) Push(ctx context.Context, cmd Command) error {
   397  	buf.mu.Lock()
   398  	defer buf.mu.Unlock()
   399  	if buf.mu.closed {
   400  		return errors.AssertionFailedf("buffer is closed")
   401  	}
   402  	buf.mu.data.AddLast(cmd)
   403  	buf.mu.lastPos++
   404  
   405  	buf.mu.cond.Signal()
   406  	return nil
   407  }
   408  
   409  // CurCmd returns the Command currently indicated by the cursor. Besides the
   410  // Command itself, the command's position is also returned; the position can be
   411  // used to later rewind() to this Command.
   412  //
   413  // If the cursor is positioned over an empty slot, the call blocks until the
   414  // next Command is pushed into the buffer.
   415  //
   416  // If the buffer has previously been Close()d, or is closed while this is
   417  // blocked, io.EOF is returned.
   418  func (buf *StmtBuf) CurCmd() (Command, CmdPos, error) {
   419  	buf.mu.Lock()
   420  	defer buf.mu.Unlock()
   421  	for {
   422  		if buf.mu.closed {
   423  			return nil, 0, io.EOF
   424  		}
   425  		curPos := buf.mu.curPos
   426  		cmdIdx, err := buf.translatePosLocked(curPos)
   427  		if err != nil {
   428  			return nil, 0, err
   429  		}
   430  		len := buf.mu.data.Len()
   431  		if cmdIdx < len {
   432  			return buf.mu.data.Get(cmdIdx).(Command), curPos, nil
   433  		}
   434  		if cmdIdx != len {
   435  			return nil, 0, errors.AssertionFailedf(
   436  				"can only wait for next command; corrupt cursor: %d", errors.Safe(curPos))
   437  		}
   438  		// Wait for the next Command to arrive to the buffer.
   439  		buf.mu.cond.Wait()
   440  	}
   441  }
   442  
   443  // translatePosLocked translates an absolute position of a command (counting
   444  // from the connection start) to the index of the respective command in the
   445  // buffer (so, it returns an index relative to the start of the buffer).
   446  //
   447  // Attempting to translate a position that's below buf.startPos returns an
   448  // error.
   449  func (buf *StmtBuf) translatePosLocked(pos CmdPos) (int, error) {
   450  	if pos < buf.mu.startPos {
   451  		return 0, errors.AssertionFailedf(
   452  			"position %d no longer in buffer (buffer starting at %d)",
   453  			errors.Safe(pos), errors.Safe(buf.mu.startPos))
   454  	}
   455  	return int(pos - buf.mu.startPos), nil
   456  }
   457  
   458  // ltrim iterates over the buffer forward and removes all commands up to
   459  // (not including) the command at pos.
   460  //
   461  // It's illegal to ltrim to a position higher than the current cursor.
   462  func (buf *StmtBuf) ltrim(ctx context.Context, pos CmdPos) {
   463  	buf.mu.Lock()
   464  	defer buf.mu.Unlock()
   465  	if pos < buf.mu.startPos {
   466  		log.Fatalf(ctx, "invalid ltrim position: %d. buf starting at: %d",
   467  			pos, buf.mu.startPos)
   468  	}
   469  	if buf.mu.curPos < pos {
   470  		log.Fatalf(ctx, "invalid ltrim position: %d when cursor is: %d",
   471  			pos, buf.mu.curPos)
   472  	}
   473  	// Remove commands one by one.
   474  	for {
   475  		if buf.mu.startPos == pos {
   476  			break
   477  		}
   478  		buf.mu.data.RemoveFirst()
   479  		buf.mu.startPos++
   480  	}
   481  }
   482  
   483  // AdvanceOne advances the cursor one Command over. The command over which
   484  // the cursor will be positioned when this returns may not be in the buffer
   485  // yet. The previous CmdPos is returned.
   486  func (buf *StmtBuf) AdvanceOne() CmdPos {
   487  	buf.mu.Lock()
   488  	prev := buf.mu.curPos
   489  	buf.mu.curPos++
   490  	buf.mu.Unlock()
   491  	return prev
   492  }
   493  
   494  // seekToNextBatch moves the cursor position to the start of the next batch of
   495  // commands, skipping past remaining commands from the current batch (if any).
   496  // Batches are delimited by Sync commands. Sync is considered to be the first
   497  // command in a batch, so once this returns, the cursor will be positioned over
   498  // a Sync command. If the cursor is positioned on a Sync when this is called,
   499  // that Sync will be skipped.
   500  //
   501  // This method blocks until a Sync command is pushed to the buffer.
   502  //
   503  // It is an error to start seeking when the cursor is positioned on an empty
   504  // slot.
   505  func (buf *StmtBuf) seekToNextBatch() error {
   506  	buf.mu.Lock()
   507  	curPos := buf.mu.curPos
   508  	cmdIdx, err := buf.translatePosLocked(curPos)
   509  	if err != nil {
   510  		buf.mu.Unlock()
   511  		return err
   512  	}
   513  	if cmdIdx == buf.mu.data.Len() {
   514  		buf.mu.Unlock()
   515  		return errors.AssertionFailedf("invalid seek start point")
   516  	}
   517  	buf.mu.Unlock()
   518  
   519  	var foundSync bool
   520  	for !foundSync {
   521  		buf.AdvanceOne()
   522  		_, pos, err := buf.CurCmd()
   523  		if err != nil {
   524  			return err
   525  		}
   526  		buf.mu.Lock()
   527  		cmdIdx, err := buf.translatePosLocked(pos)
   528  		if err != nil {
   529  			buf.mu.Unlock()
   530  			return err
   531  		}
   532  
   533  		if _, ok := buf.mu.data.Get(cmdIdx).(Sync); ok {
   534  			foundSync = true
   535  		}
   536  
   537  		buf.mu.Unlock()
   538  	}
   539  	return nil
   540  }
   541  
   542  // Rewind resets the buffer's position to pos.
   543  func (buf *StmtBuf) Rewind(ctx context.Context, pos CmdPos) {
   544  	buf.mu.Lock()
   545  	defer buf.mu.Unlock()
   546  	if pos < buf.mu.startPos {
   547  		log.Fatalf(ctx, "attempting to rewind below buffer start")
   548  	}
   549  	buf.mu.curPos = pos
   550  }
   551  
   552  // Len returns the buffer's length.
   553  func (buf *StmtBuf) Len() int {
   554  	buf.mu.Lock()
   555  	defer buf.mu.Unlock()
   556  	return buf.mu.data.Len()
   557  }
   558  
   559  // RowDescOpt specifies whether a result needs a row description message.
   560  type RowDescOpt bool
   561  
   562  const (
   563  	// NeedRowDesc specifies that a row description message is needed.
   564  	NeedRowDesc RowDescOpt = false
   565  	// DontNeedRowDesc specifies that a row description message is not needed.
   566  	DontNeedRowDesc RowDescOpt = true
   567  )
   568  
   569  // ClientComm is the interface used by the connExecutor for creating results to
   570  // be communicated to client and for exerting some control over this
   571  // communication.
   572  //
   573  // ClientComm is implemented by the pgwire connection.
   574  type ClientComm interface {
   575  	// createStatementResult creates a StatementResult for stmt.
   576  	//
   577  	// descOpt specifies if result needs to inform the client about row schema. If
   578  	// it doesn't, a SetColumns call becomes a no-op.
   579  	//
   580  	// pos is the stmt's position within the connection and is used to enforce
   581  	// that results are created in order and also to discard results through
   582  	// ClientLock.rtrim(pos).
   583  	//
   584  	// formatCodes describe how each column in the result rows is to be encoded.
   585  	// It should be nil if statement type != Rows. Otherwise, it can be nil, in
   586  	// which case every column will be encoded using the text encoding, otherwise
   587  	// it needs to contain a value for every column.
   588  	CreateStatementResult(
   589  		stmt tree.Statement,
   590  		descOpt RowDescOpt,
   591  		pos CmdPos,
   592  		formatCodes []pgwirebase.FormatCode,
   593  		conv sessiondata.DataConversionConfig,
   594  		limit int,
   595  		portalName string,
   596  		implicitTxn bool,
   597  	) CommandResult
   598  	// CreatePrepareResult creates a result for a PrepareStmt command.
   599  	CreatePrepareResult(pos CmdPos) ParseResult
   600  	// CreateDescribeResult creates a result for a DescribeStmt command.
   601  	CreateDescribeResult(pos CmdPos) DescribeResult
   602  	// CreateBindResult creates a result for a BindStmt command.
   603  	CreateBindResult(pos CmdPos) BindResult
   604  	// CreateDeleteResult creates a result for a DeletePreparedStmt command.
   605  	CreateDeleteResult(pos CmdPos) DeleteResult
   606  	// CreateSyncResult creates a result for a Sync command.
   607  	CreateSyncResult(pos CmdPos) SyncResult
   608  	// CreateFlushResult creates a result for a Flush command.
   609  	CreateFlushResult(pos CmdPos) FlushResult
   610  	// CreateErrorResult creates a result on which only errors can be communicated
   611  	// to the client.
   612  	CreateErrorResult(pos CmdPos) ErrorResult
   613  	// CreateEmptyQueryResult creates a result for an empty-string query.
   614  	CreateEmptyQueryResult(pos CmdPos) EmptyQueryResult
   615  	// CreateCopyInResult creates a result for a Copy-in command.
   616  	CreateCopyInResult(pos CmdPos) CopyInResult
   617  	// CreateDrainResult creates a result for a Drain command.
   618  	CreateDrainResult(pos CmdPos) DrainResult
   619  
   620  	// lockCommunication ensures that no further results are delivered to the
   621  	// client. The returned ClientLock can be queried to see what results have
   622  	// been already delivered to the client and to discard results that haven't
   623  	// been delivered.
   624  	//
   625  	// ClientLock.Close() needs to be called on the returned lock once
   626  	// communication can be unlocked (i.e. results can be delivered to the client
   627  	// again).
   628  	LockCommunication() ClientLock
   629  
   630  	// Flush delivers all the previous results to the client. The results might
   631  	// have been buffered, in which case this flushes the buffer.
   632  	Flush(pos CmdPos) error
   633  }
   634  
   635  // CommandResult represents the result of a statement. It which needs to be
   636  // ultimately delivered to the client. pgwire.conn implements this.
   637  type CommandResult interface {
   638  	RestrictedCommandResult
   639  	CommandResultClose
   640  }
   641  
   642  // CommandResultErrBase is the subset of CommandResult dealing with setting a
   643  // query execution error.
   644  type CommandResultErrBase interface {
   645  	// SetError accumulates an execution error that needs to be reported to the
   646  	// client. No further calls other than SetError(), Close() and Discard() are
   647  	// allowed.
   648  	//
   649  	// Calling SetError() a second time overwrites the previously set error.
   650  	SetError(error)
   651  
   652  	// Err returns the error previously set with SetError(), if any.
   653  	Err() error
   654  }
   655  
   656  // ResultBase is the common interface implemented by all the different command
   657  // results.
   658  type ResultBase interface {
   659  	CommandResultErrBase
   660  	CommandResultClose
   661  }
   662  
   663  // CommandResultClose is a subset of CommandResult dealing with the closing of
   664  // the result.
   665  type CommandResultClose interface {
   666  	// Close marks a result as complete. No further uses of the CommandResult are
   667  	// allowed after this call. All results must be eventually closed through
   668  	// Close()/Discard(), except in case query processing has encountered an
   669  	// irrecoverable error and the client connection will be closed; in such
   670  	// cases it is not mandated that these functions are called on the result
   671  	// that may have been open at the time the error occurred.
   672  	// NOTE(andrei): We might want to tighten the contract if the results get any
   673  	// state that needs to be closed even when the whole connection is about to be
   674  	// terminated.
   675  	Close(context.Context, TransactionStatusIndicator)
   676  
   677  	// Discard is called to mark the fact that the result is being disposed off.
   678  	// No completion message will be sent to the client. The expectation is that
   679  	// either the no other methods on the result had previously been used (and so
   680  	// no data has been buffered for the client), or there is a communication lock
   681  	// in effect and the buffer will be rewound - in either case, the client will
   682  	// never see any bytes pertaining to this result.
   683  	Discard()
   684  }
   685  
   686  // RestrictedCommandResult is a subset of CommandResult meant to make it clear
   687  // that its clients don't close the CommandResult.
   688  type RestrictedCommandResult interface {
   689  	CommandResultErrBase
   690  
   691  	// AppendParamStatusUpdate appends a parameter status update to the result.
   692  	// This gets flushed only when the CommandResult is closed.
   693  	AppendParamStatusUpdate(string, string)
   694  
   695  	// AppendNotice appends a notice to the result.
   696  	// This gets flushed only when the CommandResult is closed.
   697  	AppendNotice(noticeErr error)
   698  
   699  	// SetColumns informs the client about the schema of the result. The columns
   700  	// can be nil.
   701  	//
   702  	// This needs to be called (once) before AddRow.
   703  	SetColumns(context.Context, sqlbase.ResultColumns)
   704  
   705  	// ResetStmtType allows a client to change the statement type of the current
   706  	// result, from the original one set when the result was created trough
   707  	// ClientComm.createStatementResult.
   708  	ResetStmtType(stmt tree.Statement)
   709  
   710  	// AddRow accumulates a result row.
   711  	//
   712  	// The implementation cannot hold on to the row slice; it needs to make a
   713  	// shallow copy if it needs to.
   714  	AddRow(ctx context.Context, row tree.Datums) error
   715  
   716  	// IncrementRowsAffected increments a counter by n. This is used for all
   717  	// result types other than tree.Rows.
   718  	IncrementRowsAffected(n int)
   719  
   720  	// RowsAffected returns either the number of times AddRow was called, or the
   721  	// sum of all n passed into IncrementRowsAffected.
   722  	RowsAffected() int
   723  
   724  	// DisableBuffering can be called during execution to ensure that
   725  	// the results accumulated so far, and all subsequent rows added
   726  	// to this CommandResult, will be flushed immediately to the client.
   727  	// This is currently used for sinkless changefeeds.
   728  	DisableBuffering()
   729  }
   730  
   731  // DescribeResult represents the result of a Describe command (for either
   732  // describing a prepared statement or a portal).
   733  type DescribeResult interface {
   734  	ResultBase
   735  
   736  	// SetInferredTypes tells the client about the inferred placeholder types.
   737  	SetInferredTypes([]oid.Oid)
   738  	// SetNoDataDescription is used to tell the client that the prepared statement
   739  	// or portal produces no rows.
   740  	SetNoDataRowDescription()
   741  	// SetPrepStmtOutput tells the client about the results schema of a prepared
   742  	// statement.
   743  	SetPrepStmtOutput(context.Context, sqlbase.ResultColumns)
   744  	// SetPortalOutput tells the client about the results schema and formatting of
   745  	// a portal.
   746  	SetPortalOutput(context.Context, sqlbase.ResultColumns, []pgwirebase.FormatCode)
   747  }
   748  
   749  // ParseResult represents the result of a Parse command.
   750  type ParseResult interface {
   751  	ResultBase
   752  }
   753  
   754  // BindResult represents the result of a Bind command.
   755  type BindResult interface {
   756  	ResultBase
   757  }
   758  
   759  // ErrorResult represents the result of a SendError command.
   760  type ErrorResult interface {
   761  	ResultBase
   762  }
   763  
   764  // DeleteResult represents the result of a DeletePreparedStatement command.
   765  type DeleteResult interface {
   766  	ResultBase
   767  }
   768  
   769  // SyncResult represents the result of a Sync command. When closed, a
   770  // readyForQuery message will be generated and all buffered data will be
   771  // flushed.
   772  type SyncResult interface {
   773  	ResultBase
   774  }
   775  
   776  // FlushResult represents the result of a Flush command. When this result is
   777  // closed, all previously accumulated results are flushed to the client.
   778  type FlushResult interface {
   779  	ResultBase
   780  }
   781  
   782  // DrainResult represents the result of a Drain command. Closing this result
   783  // produces no output for the client.
   784  type DrainResult interface {
   785  	ResultBase
   786  }
   787  
   788  // EmptyQueryResult represents the result of an empty query (a query
   789  // representing a blank string).
   790  type EmptyQueryResult interface {
   791  	ResultBase
   792  }
   793  
   794  // CopyInResult represents the result of a CopyIn command. Closing this result
   795  // produces no output for the client.
   796  type CopyInResult interface {
   797  	ResultBase
   798  }
   799  
   800  // ClientLock is an interface returned by ClientComm.lockCommunication(). It
   801  // represents a lock on the delivery of results to a SQL client. While such a
   802  // lock is used, no more results are delivered. The lock itself can be used to
   803  // query what results have already been delivered and to discard results that
   804  // haven't been delivered.
   805  type ClientLock interface {
   806  	// Close unlocks the ClientComm from whence this ClientLock came from. After
   807  	// Close is called, buffered results may again be sent to the client,
   808  	// according to the result streaming policy.
   809  	//
   810  	// Once Close() is called, the ClientLock cannot be used anymore.
   811  	Close()
   812  
   813  	// ClientPos returns the position of the latest command for which results
   814  	// have been sent to the client. The position is relative to the start of the
   815  	// connection.
   816  	ClientPos() CmdPos
   817  
   818  	// RTrim iterates backwards through the results and drops all results with
   819  	// position >= pos.
   820  	// It is illegal to call rtrim with a position <= clientPos(). In other words,
   821  	// results can
   822  	RTrim(ctx context.Context, pos CmdPos)
   823  }
   824  
   825  // rewindCapability is used pass rewinding instructions in between different
   826  // layers of the connExecutor state machine. It ties together a position to
   827  // which we want to rewind within the stream of commands with:
   828  // a) a ClientLock that guarantees that the rewind to the respective position is
   829  // (and remains) possible.
   830  // b) the StmtBuf that needs to be rewound at the same time as the results.
   831  //
   832  // rewindAndUnlock() needs to be called eventually in order to actually perform
   833  // the rewinding and unlock the respective ClientComm.
   834  type rewindCapability struct {
   835  	cl  ClientLock
   836  	buf *StmtBuf
   837  
   838  	rewindPos CmdPos
   839  }
   840  
   841  // rewindAndUnlock performs the rewinding described by the rewindCapability and
   842  // unlocks the respective ClientComm.
   843  func (rc *rewindCapability) rewindAndUnlock(ctx context.Context) {
   844  	rc.cl.RTrim(ctx, rc.rewindPos)
   845  	rc.buf.Rewind(ctx, rc.rewindPos)
   846  	rc.cl.Close()
   847  }
   848  
   849  // close closes the underlying ClientLock.
   850  func (rc *rewindCapability) close() {
   851  	rc.cl.Close()
   852  }
   853  
   854  type resCloseType bool
   855  
   856  const closed resCloseType = true
   857  const discarded resCloseType = false
   858  
   859  // bufferedCommandResult is a CommandResult that buffers rows and can call a
   860  // provided callback when closed.
   861  type bufferedCommandResult struct {
   862  	err          error
   863  	rows         []tree.Datums
   864  	rowsAffected int
   865  	cols         sqlbase.ResultColumns
   866  
   867  	// errOnly, if set, makes AddRow() panic. This can be used when the execution
   868  	// of the query is not expected to produce any results.
   869  	errOnly bool
   870  
   871  	// closeCallback, if set, is called when Close()/Discard() is called.
   872  	closeCallback func(*bufferedCommandResult, resCloseType, error)
   873  }
   874  
   875  var _ RestrictedCommandResult = &bufferedCommandResult{}
   876  var _ CommandResultClose = &bufferedCommandResult{}
   877  
   878  // SetColumns is part of the RestrictedCommandResult interface.
   879  func (r *bufferedCommandResult) SetColumns(_ context.Context, cols sqlbase.ResultColumns) {
   880  	if r.errOnly {
   881  		panic("SetColumns() called when errOnly is set")
   882  	}
   883  	r.cols = cols
   884  }
   885  
   886  // AppendParamStatusUpdate is part of the RestrictedCommandResult interface.
   887  func (r *bufferedCommandResult) AppendParamStatusUpdate(key string, val string) {
   888  	panic("unimplemented")
   889  }
   890  
   891  // AppendNotice is part of the RestrictedCommandResult interface.
   892  func (r *bufferedCommandResult) AppendNotice(noticeErr error) {
   893  	panic("unimplemented")
   894  }
   895  
   896  // ResetStmtType is part of the RestrictedCommandResult interface.
   897  func (r *bufferedCommandResult) ResetStmtType(stmt tree.Statement) {
   898  	panic("unimplemented")
   899  }
   900  
   901  // AddRow is part of the RestrictedCommandResult interface.
   902  func (r *bufferedCommandResult) AddRow(ctx context.Context, row tree.Datums) error {
   903  	if r.errOnly {
   904  		panic("AddRow() called when errOnly is set")
   905  	}
   906  	rowCopy := make(tree.Datums, len(row))
   907  	copy(rowCopy, row)
   908  	r.rows = append(r.rows, rowCopy)
   909  	return nil
   910  }
   911  
   912  func (r *bufferedCommandResult) DisableBuffering() {
   913  	panic("cannot disable buffering here")
   914  }
   915  
   916  // SetError is part of the RestrictedCommandResult interface.
   917  func (r *bufferedCommandResult) SetError(err error) {
   918  	r.err = err
   919  }
   920  
   921  // Err is part of the RestrictedCommandResult interface.
   922  func (r *bufferedCommandResult) Err() error {
   923  	return r.err
   924  }
   925  
   926  // IncrementRowsAffected is part of the RestrictedCommandResult interface.
   927  func (r *bufferedCommandResult) IncrementRowsAffected(n int) {
   928  	r.rowsAffected += n
   929  }
   930  
   931  // RowsAffected is part of the RestrictedCommandResult interface.
   932  func (r *bufferedCommandResult) RowsAffected() int {
   933  	return r.rowsAffected
   934  }
   935  
   936  // Close is part of the CommandResultClose interface.
   937  func (r *bufferedCommandResult) Close(context.Context, TransactionStatusIndicator) {
   938  	if r.closeCallback != nil {
   939  		r.closeCallback(r, closed, nil /* err */)
   940  	}
   941  }
   942  
   943  // Discard is part of the CommandResult interface.
   944  func (r *bufferedCommandResult) Discard() {
   945  	if r.closeCallback != nil {
   946  		r.closeCallback(r, discarded, nil /* err */)
   947  	}
   948  }
   949  
   950  // SetInferredTypes is part of the DescribeResult interface.
   951  func (r *bufferedCommandResult) SetInferredTypes([]oid.Oid) {}
   952  
   953  // SetNoDataRowDescription is part of the DescribeResult interface.
   954  func (r *bufferedCommandResult) SetNoDataRowDescription() {}
   955  
   956  // SetPrepStmtOutput is part of the DescribeResult interface.
   957  func (r *bufferedCommandResult) SetPrepStmtOutput(context.Context, sqlbase.ResultColumns) {}
   958  
   959  // SetPortalOutput is part of the DescribeResult interface.
   960  func (r *bufferedCommandResult) SetPortalOutput(
   961  	context.Context, sqlbase.ResultColumns, []pgwirebase.FormatCode,
   962  ) {
   963  }