vitess.io/vitess@v0.16.2/go/vt/vttablet/sandboxconn/sandboxconn.go (about)

     1  /*
     2  Copyright 2019 The Vitess Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  // Package sandboxconn provides a fake QueryService implementation for tests.
    18  // It can return real results, and simulate error cases.
    19  package sandboxconn
    20  
    21  import (
    22  	"context"
    23  	"fmt"
    24  	"sync"
    25  	"time"
    26  
    27  	"vitess.io/vitess/go/vt/log"
    28  	"vitess.io/vitess/go/vt/sqlparser"
    29  
    30  	"vitess.io/vitess/go/sqltypes"
    31  	"vitess.io/vitess/go/sync2"
    32  	"vitess.io/vitess/go/vt/vterrors"
    33  	"vitess.io/vitess/go/vt/vttablet/queryservice"
    34  
    35  	binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata"
    36  	querypb "vitess.io/vitess/go/vt/proto/query"
    37  	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
    38  	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
    39  )
    40  
    41  // SandboxConn satisfies the QueryService interface
    42  type SandboxConn struct {
    43  	tablet *topodatapb.Tablet
    44  
    45  	// These errors work for all functions.
    46  	MustFailCodes map[vtrpcpb.Code]int
    47  
    48  	// These errors are triggered only for specific functions.
    49  	// For now these are just for the 2PC functions.
    50  	MustFailPrepare             int
    51  	MustFailCommitPrepared      int
    52  	MustFailRollbackPrepared    int
    53  	MustFailCreateTransaction   int
    54  	MustFailStartCommit         int
    55  	MustFailSetRollback         int
    56  	MustFailConcludeTransaction int
    57  	// MustFailExecute is keyed by the statement type and stores the number
    58  	// of times to fail when it sees that statement type.
    59  	// Once, exhausted it will start returning non-error response.
    60  	MustFailExecute map[sqlparser.StatementType]int
    61  
    62  	// These Count vars report how often the corresponding
    63  	// functions were called.
    64  	ExecCount                sync2.AtomicInt64
    65  	BeginCount               sync2.AtomicInt64
    66  	CommitCount              sync2.AtomicInt64
    67  	RollbackCount            sync2.AtomicInt64
    68  	AsTransactionCount       sync2.AtomicInt64
    69  	PrepareCount             sync2.AtomicInt64
    70  	CommitPreparedCount      sync2.AtomicInt64
    71  	RollbackPreparedCount    sync2.AtomicInt64
    72  	CreateTransactionCount   sync2.AtomicInt64
    73  	StartCommitCount         sync2.AtomicInt64
    74  	SetRollbackCount         sync2.AtomicInt64
    75  	ConcludeTransactionCount sync2.AtomicInt64
    76  	ReadTransactionCount     sync2.AtomicInt64
    77  	ReserveCount             sync2.AtomicInt64
    78  	ReleaseCount             sync2.AtomicInt64
    79  	GetSchemaCount           sync2.AtomicInt64
    80  
    81  	// Queries stores the non-batch requests received.
    82  	Queries []*querypb.BoundQuery
    83  
    84  	// BatchQueries stores the batch requests received
    85  	// Each batch request is inlined as a slice of Queries.
    86  	BatchQueries [][]*querypb.BoundQuery
    87  
    88  	// Options stores the options received by all calls.
    89  	Options []*querypb.ExecuteOptions
    90  
    91  	// results specifies the results to be returned.
    92  	// They're consumed as results are returned. If there are
    93  	// no results left, SingleRowResult is returned.
    94  	results []*sqltypes.Result
    95  
    96  	// ReadTransactionResults is used for returning results for ReadTransaction.
    97  	ReadTransactionResults []*querypb.TransactionMetadata
    98  
    99  	MessageIDs []*querypb.Value
   100  
   101  	// vstream expectations.
   102  	StartPos      string
   103  	VStreamEvents [][]*binlogdatapb.VEvent
   104  	VStreamErrors []error
   105  	VStreamCh     chan *binlogdatapb.VEvent
   106  
   107  	// transaction id generator
   108  	TransactionID sync2.AtomicInt64
   109  
   110  	// reserve id generator
   111  	ReserveID sync2.AtomicInt64
   112  
   113  	mapMu     sync.Mutex //protects the map txIDToRID
   114  	txIDToRID map[int64]int64
   115  
   116  	sExecMu sync.Mutex
   117  	execMu  sync.Mutex
   118  
   119  	// this error will only happen once
   120  	EphemeralShardErr error
   121  
   122  	NotServing bool
   123  
   124  	getSchemaResult []map[string]string
   125  }
   126  
   127  var _ queryservice.QueryService = (*SandboxConn)(nil) // compile-time interface check
   128  
   129  // NewSandboxConn returns a new SandboxConn targeted to the provided tablet.
   130  func NewSandboxConn(t *topodatapb.Tablet) *SandboxConn {
   131  	return &SandboxConn{
   132  		tablet:          t,
   133  		MustFailCodes:   make(map[vtrpcpb.Code]int),
   134  		MustFailExecute: make(map[sqlparser.StatementType]int),
   135  		txIDToRID:       make(map[int64]int64),
   136  	}
   137  }
   138  
   139  func (sbc *SandboxConn) getError() error {
   140  	for code, count := range sbc.MustFailCodes {
   141  		if count == 0 {
   142  			continue
   143  		}
   144  		sbc.MustFailCodes[code] = count - 1
   145  		return vterrors.New(code, fmt.Sprintf("%v error", code))
   146  	}
   147  	if sbc.EphemeralShardErr != nil {
   148  		err := sbc.EphemeralShardErr
   149  		sbc.EphemeralShardErr = nil
   150  		return err
   151  	}
   152  	return nil
   153  }
   154  
   155  // SetResults sets what this con should return next time.
   156  func (sbc *SandboxConn) SetResults(r []*sqltypes.Result) {
   157  	sbc.results = r
   158  }
   159  
   160  // SetSchemaResult sets what GetSchema should return on each call.
   161  func (sbc *SandboxConn) SetSchemaResult(r []map[string]string) {
   162  	sbc.getSchemaResult = r
   163  }
   164  
   165  // Execute is part of the QueryService interface.
   166  func (sbc *SandboxConn) Execute(ctx context.Context, target *querypb.Target, query string, bindVars map[string]*querypb.BindVariable, transactionID, reservedID int64, options *querypb.ExecuteOptions) (*sqltypes.Result, error) {
   167  	sbc.execMu.Lock()
   168  	defer sbc.execMu.Unlock()
   169  	sbc.ExecCount.Add(1)
   170  	if sbc.NotServing {
   171  		return nil, vterrors.New(vtrpcpb.Code_CLUSTER_EVENT, vterrors.NotServing)
   172  	}
   173  	if sbc.tablet.Type != target.TabletType {
   174  		return nil, vterrors.Errorf(vtrpcpb.Code_FAILED_PRECONDITION, "%s: %v, want: %v", vterrors.WrongTablet, target.TabletType, sbc.tablet.Type)
   175  	}
   176  	bv := make(map[string]*querypb.BindVariable)
   177  	for k, v := range bindVars {
   178  		bv[k] = v
   179  	}
   180  	sbc.Queries = append(sbc.Queries, &querypb.BoundQuery{
   181  		Sql:           query,
   182  		BindVariables: bv,
   183  	})
   184  	sbc.Options = append(sbc.Options, options)
   185  	if err := sbc.getError(); err != nil {
   186  		return nil, err
   187  	}
   188  
   189  	stmt, _ := sqlparser.Parse(query) // knowingly ignoring the error
   190  	if sbc.MustFailExecute[sqlparser.ASTToStatementType(stmt)] > 0 {
   191  		sbc.MustFailExecute[sqlparser.ASTToStatementType(stmt)] = sbc.MustFailExecute[sqlparser.ASTToStatementType(stmt)] - 1
   192  		return nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "failed query: %v", query)
   193  	}
   194  	return sbc.getNextResult(stmt), nil
   195  }
   196  
   197  // StreamExecute is part of the QueryService interface.
   198  func (sbc *SandboxConn) StreamExecute(ctx context.Context, target *querypb.Target, query string, bindVars map[string]*querypb.BindVariable, transactionID int64, reservedID int64, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) error {
   199  	sbc.sExecMu.Lock()
   200  	sbc.ExecCount.Add(1)
   201  	bv := make(map[string]*querypb.BindVariable)
   202  	for k, v := range bindVars {
   203  		bv[k] = v
   204  	}
   205  	sbc.Queries = append(sbc.Queries, &querypb.BoundQuery{
   206  		Sql:           query,
   207  		BindVariables: bv,
   208  	})
   209  	sbc.Options = append(sbc.Options, options)
   210  	err := sbc.getError()
   211  	if err != nil {
   212  		sbc.sExecMu.Unlock()
   213  		return err
   214  	}
   215  	parse, _ := sqlparser.Parse(query)
   216  
   217  	if sbc.results == nil {
   218  		nextRs := sbc.getNextResult(parse)
   219  		sbc.sExecMu.Unlock()
   220  		return callback(nextRs)
   221  	}
   222  
   223  	for len(sbc.results) > 0 {
   224  		nextRs := sbc.getNextResult(parse)
   225  		sbc.sExecMu.Unlock()
   226  		err := callback(nextRs)
   227  		if err != nil {
   228  			return err
   229  		}
   230  		sbc.sExecMu.Lock()
   231  	}
   232  
   233  	sbc.sExecMu.Unlock()
   234  	return nil
   235  }
   236  
   237  // Begin is part of the QueryService interface.
   238  func (sbc *SandboxConn) Begin(ctx context.Context, target *querypb.Target, options *querypb.ExecuteOptions) (queryservice.TransactionState, error) {
   239  	return sbc.begin(ctx, target, nil, 0, options)
   240  }
   241  
   242  func (sbc *SandboxConn) begin(ctx context.Context, target *querypb.Target, preQueries []string, reservedID int64, options *querypb.ExecuteOptions) (queryservice.TransactionState, error) {
   243  	sbc.BeginCount.Add(1)
   244  	err := sbc.getError()
   245  	if err != nil {
   246  		return queryservice.TransactionState{}, err
   247  	}
   248  
   249  	transactionID := reservedID
   250  	if transactionID == 0 {
   251  		transactionID = sbc.TransactionID.Add(1)
   252  	}
   253  	for _, preQuery := range preQueries {
   254  		_, err := sbc.Execute(ctx, target, preQuery, nil, transactionID, reservedID, options)
   255  		if err != nil {
   256  			return queryservice.TransactionState{}, err
   257  		}
   258  	}
   259  	return queryservice.TransactionState{TransactionID: transactionID, TabletAlias: sbc.tablet.Alias}, nil
   260  }
   261  
   262  // Commit is part of the QueryService interface.
   263  func (sbc *SandboxConn) Commit(ctx context.Context, target *querypb.Target, transactionID int64) (int64, error) {
   264  	sbc.CommitCount.Add(1)
   265  	reservedID := sbc.getTxReservedID(transactionID)
   266  	if reservedID != 0 {
   267  		reservedID = sbc.ReserveID.Add(1)
   268  	}
   269  	return reservedID, sbc.getError()
   270  }
   271  
   272  // Rollback is part of the QueryService interface.
   273  func (sbc *SandboxConn) Rollback(ctx context.Context, target *querypb.Target, transactionID int64) (int64, error) {
   274  	sbc.RollbackCount.Add(1)
   275  	reservedID := sbc.getTxReservedID(transactionID)
   276  	if reservedID != 0 {
   277  		reservedID = sbc.ReserveID.Add(1)
   278  	}
   279  	return reservedID, sbc.getError()
   280  }
   281  
   282  // Prepare prepares the specified transaction.
   283  func (sbc *SandboxConn) Prepare(ctx context.Context, target *querypb.Target, transactionID int64, dtid string) (err error) {
   284  	sbc.PrepareCount.Add(1)
   285  	if sbc.MustFailPrepare > 0 {
   286  		sbc.MustFailPrepare--
   287  		return vterrors.New(vtrpcpb.Code_FAILED_PRECONDITION, "error: err")
   288  	}
   289  	return sbc.getError()
   290  }
   291  
   292  // CommitPrepared commits the prepared transaction.
   293  func (sbc *SandboxConn) CommitPrepared(ctx context.Context, target *querypb.Target, dtid string) (err error) {
   294  	sbc.CommitPreparedCount.Add(1)
   295  	if sbc.MustFailCommitPrepared > 0 {
   296  		sbc.MustFailCommitPrepared--
   297  		return vterrors.New(vtrpcpb.Code_FAILED_PRECONDITION, "error: err")
   298  	}
   299  	return sbc.getError()
   300  }
   301  
   302  // RollbackPrepared rolls back the prepared transaction.
   303  func (sbc *SandboxConn) RollbackPrepared(ctx context.Context, target *querypb.Target, dtid string, originalID int64) (err error) {
   304  	sbc.RollbackPreparedCount.Add(1)
   305  	if sbc.MustFailRollbackPrepared > 0 {
   306  		sbc.MustFailRollbackPrepared--
   307  		return vterrors.New(vtrpcpb.Code_FAILED_PRECONDITION, "error: err")
   308  	}
   309  	return sbc.getError()
   310  }
   311  
   312  // CreateTransaction creates the metadata for a 2PC transaction.
   313  func (sbc *SandboxConn) CreateTransaction(ctx context.Context, target *querypb.Target, dtid string, participants []*querypb.Target) (err error) {
   314  	sbc.CreateTransactionCount.Add(1)
   315  	if sbc.MustFailCreateTransaction > 0 {
   316  		sbc.MustFailCreateTransaction--
   317  		return vterrors.New(vtrpcpb.Code_FAILED_PRECONDITION, "error: err")
   318  	}
   319  	return sbc.getError()
   320  }
   321  
   322  // StartCommit atomically commits the transaction along with the
   323  // decision to commit the associated 2pc transaction.
   324  func (sbc *SandboxConn) StartCommit(ctx context.Context, target *querypb.Target, transactionID int64, dtid string) (err error) {
   325  	sbc.StartCommitCount.Add(1)
   326  	if sbc.MustFailStartCommit > 0 {
   327  		sbc.MustFailStartCommit--
   328  		return vterrors.New(vtrpcpb.Code_FAILED_PRECONDITION, "error: err")
   329  	}
   330  	return sbc.getError()
   331  }
   332  
   333  // SetRollback transitions the 2pc transaction to the Rollback state.
   334  // If a transaction id is provided, that transaction is also rolled back.
   335  func (sbc *SandboxConn) SetRollback(ctx context.Context, target *querypb.Target, dtid string, transactionID int64) (err error) {
   336  	sbc.SetRollbackCount.Add(1)
   337  	if sbc.MustFailSetRollback > 0 {
   338  		sbc.MustFailSetRollback--
   339  		return vterrors.New(vtrpcpb.Code_FAILED_PRECONDITION, "error: err")
   340  	}
   341  	return sbc.getError()
   342  }
   343  
   344  // ConcludeTransaction deletes the 2pc transaction metadata
   345  // essentially resolving it.
   346  func (sbc *SandboxConn) ConcludeTransaction(ctx context.Context, target *querypb.Target, dtid string) (err error) {
   347  	sbc.ConcludeTransactionCount.Add(1)
   348  	if sbc.MustFailConcludeTransaction > 0 {
   349  		sbc.MustFailConcludeTransaction--
   350  		return vterrors.New(vtrpcpb.Code_FAILED_PRECONDITION, "error: err")
   351  	}
   352  	return sbc.getError()
   353  }
   354  
   355  // ReadTransaction returns the metadata for the sepcified dtid.
   356  func (sbc *SandboxConn) ReadTransaction(ctx context.Context, target *querypb.Target, dtid string) (metadata *querypb.TransactionMetadata, err error) {
   357  	sbc.ReadTransactionCount.Add(1)
   358  	if err := sbc.getError(); err != nil {
   359  		return nil, err
   360  	}
   361  	if len(sbc.ReadTransactionResults) >= 1 {
   362  		res := sbc.ReadTransactionResults[0]
   363  		sbc.ReadTransactionResults = sbc.ReadTransactionResults[1:]
   364  		return res, nil
   365  	}
   366  	return nil, nil
   367  }
   368  
   369  // BeginExecute is part of the QueryService interface.
   370  func (sbc *SandboxConn) BeginExecute(ctx context.Context, target *querypb.Target, preQueries []string, query string, bindVars map[string]*querypb.BindVariable, reservedID int64, options *querypb.ExecuteOptions) (queryservice.TransactionState, *sqltypes.Result, error) {
   371  	state, err := sbc.begin(ctx, target, preQueries, reservedID, options)
   372  	if state.TransactionID != 0 {
   373  		sbc.setTxReservedID(state.TransactionID, reservedID)
   374  	}
   375  	if err != nil {
   376  		return queryservice.TransactionState{}, nil, err
   377  	}
   378  	result, err := sbc.Execute(ctx, target, query, bindVars, state.TransactionID, reservedID, options)
   379  	return state, result, err
   380  }
   381  
   382  // BeginStreamExecute is part of the QueryService interface.
   383  func (sbc *SandboxConn) BeginStreamExecute(ctx context.Context, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, reservedID int64, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) (queryservice.TransactionState, error) {
   384  	state, err := sbc.begin(ctx, target, preQueries, reservedID, options)
   385  	if state.TransactionID != 0 {
   386  		sbc.setTxReservedID(state.TransactionID, reservedID)
   387  	}
   388  	if err != nil {
   389  		return queryservice.TransactionState{}, err
   390  	}
   391  	err = sbc.StreamExecute(ctx, target, sql, bindVariables, state.TransactionID, reservedID, options, callback)
   392  	return state, err
   393  }
   394  
   395  // MessageStream is part of the QueryService interface.
   396  func (sbc *SandboxConn) MessageStream(ctx context.Context, target *querypb.Target, name string, callback func(*sqltypes.Result) error) (err error) {
   397  	if err := sbc.getError(); err != nil {
   398  		return err
   399  	}
   400  	r := sbc.getNextResult(nil)
   401  	if r == nil {
   402  		return nil
   403  	}
   404  	callback(r)
   405  	return nil
   406  }
   407  
   408  // MessageAck is part of the QueryService interface.
   409  func (sbc *SandboxConn) MessageAck(ctx context.Context, target *querypb.Target, name string, ids []*querypb.Value) (count int64, err error) {
   410  	sbc.MessageIDs = ids
   411  	return int64(len(ids)), nil
   412  }
   413  
   414  // SandboxSQRowCount is the default number of fake splits returned.
   415  var SandboxSQRowCount = int64(10)
   416  
   417  // StreamHealth is not implemented.
   418  func (sbc *SandboxConn) StreamHealth(ctx context.Context, callback func(*querypb.StreamHealthResponse) error) error {
   419  	return fmt.Errorf("not implemented in test")
   420  }
   421  
   422  // ExpectVStreamStartPos makes the conn verify that that the next vstream request has the right startPos.
   423  func (sbc *SandboxConn) ExpectVStreamStartPos(startPos string) {
   424  	sbc.StartPos = startPos
   425  }
   426  
   427  // AddVStreamEvents adds a set of VStream events to be returned.
   428  func (sbc *SandboxConn) AddVStreamEvents(events []*binlogdatapb.VEvent, err error) {
   429  	sbc.VStreamEvents = append(sbc.VStreamEvents, events)
   430  	sbc.VStreamErrors = append(sbc.VStreamErrors, err)
   431  }
   432  
   433  // VStream is part of the QueryService interface.
   434  func (sbc *SandboxConn) VStream(ctx context.Context, request *binlogdatapb.VStreamRequest, send func([]*binlogdatapb.VEvent) error) error {
   435  	if sbc.StartPos != "" && sbc.StartPos != request.Position {
   436  		log.Errorf("startPos(%v): %v, want %v", request.Target, request.Position, sbc.StartPos)
   437  		return fmt.Errorf("startPos(%v): %v, want %v", request.Target, request.Position, sbc.StartPos)
   438  	}
   439  	done := false
   440  	// for testing the minimize stream skew feature (TestStreamSkew) we need the ability to send events in specific sequences from
   441  	// multiple streams. We introduce a channel in the sandbox that we listen on and vstream those events
   442  	// as we receive them. We also need to simulate vstreamer heartbeats since the skew detection logic depends on it
   443  	// in case of shards where there are no real events within a second
   444  	if sbc.VStreamCh != nil {
   445  		lastTimestamp := int64(0)
   446  		for !done {
   447  			timer := time.NewTimer(1 * time.Second)
   448  			select {
   449  			case <-timer.C:
   450  				events := []*binlogdatapb.VEvent{{
   451  					Type:        binlogdatapb.VEventType_HEARTBEAT,
   452  					Timestamp:   lastTimestamp,
   453  					CurrentTime: lastTimestamp,
   454  				}, {
   455  					Type:        binlogdatapb.VEventType_COMMIT,
   456  					Timestamp:   lastTimestamp,
   457  					CurrentTime: lastTimestamp,
   458  				}}
   459  
   460  				if err := send(events); err != nil {
   461  					log.Infof("error sending event in test sandbox %s", err.Error())
   462  					return err
   463  				}
   464  				lastTimestamp++
   465  
   466  			case ev := <-sbc.VStreamCh:
   467  				if ev == nil {
   468  					done = true
   469  				}
   470  				if err := send([]*binlogdatapb.VEvent{ev}); err != nil {
   471  					log.Infof("error sending event in test sandbox %s", err.Error())
   472  					return err
   473  				}
   474  				lastTimestamp = ev.Timestamp
   475  			}
   476  		}
   477  	} else {
   478  		// this path is followed for all vstream tests other than the skew tests
   479  		for len(sbc.VStreamEvents) != 0 {
   480  			ev := sbc.VStreamEvents[0]
   481  			err := sbc.VStreamErrors[0]
   482  			sbc.VStreamEvents = sbc.VStreamEvents[1:]
   483  			sbc.VStreamErrors = sbc.VStreamErrors[1:]
   484  			if ev == nil {
   485  				return err
   486  			}
   487  			if err := send(ev); err != nil {
   488  				return err
   489  			}
   490  		}
   491  	}
   492  	// Don't return till context is canceled.
   493  	<-ctx.Done()
   494  	return ctx.Err()
   495  }
   496  
   497  // VStreamRows is part of the QueryService interface.
   498  func (sbc *SandboxConn) VStreamRows(ctx context.Context, request *binlogdatapb.VStreamRowsRequest, send func(*binlogdatapb.VStreamRowsResponse) error) error {
   499  	return fmt.Errorf("not implemented in test")
   500  }
   501  
   502  // VStreamResults is part of the QueryService interface.
   503  func (sbc *SandboxConn) VStreamResults(ctx context.Context, target *querypb.Target, query string, send func(*binlogdatapb.VStreamResultsResponse) error) error {
   504  	return fmt.Errorf("not implemented in test")
   505  }
   506  
   507  // QueryServiceByAlias is part of the Gateway interface.
   508  func (sbc *SandboxConn) QueryServiceByAlias(_ *topodatapb.TabletAlias, _ *querypb.Target) (queryservice.QueryService, error) {
   509  	return sbc, nil
   510  }
   511  
   512  // HandlePanic is part of the QueryService interface.
   513  func (sbc *SandboxConn) HandlePanic(err *error) {
   514  }
   515  
   516  // ReserveBeginExecute implements the QueryService interface
   517  func (sbc *SandboxConn) ReserveBeginExecute(ctx context.Context, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, options *querypb.ExecuteOptions) (queryservice.ReservedTransactionState, *sqltypes.Result, error) {
   518  	reservedID := sbc.reserve(ctx, target, preQueries, bindVariables, 0, options)
   519  	state, result, err := sbc.BeginExecute(ctx, target, postBeginQueries, sql, bindVariables, reservedID, options)
   520  	if state.TransactionID != 0 {
   521  		sbc.setTxReservedID(state.TransactionID, reservedID)
   522  	}
   523  	return queryservice.ReservedTransactionState{
   524  		ReservedID:    reservedID,
   525  		TransactionID: state.TransactionID,
   526  		TabletAlias:   state.TabletAlias,
   527  	}, result, err
   528  }
   529  
   530  // ReserveBeginStreamExecute is part of the QueryService interface.
   531  func (sbc *SandboxConn) ReserveBeginStreamExecute(ctx context.Context, target *querypb.Target, preQueries []string, postBeginQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) (queryservice.ReservedTransactionState, error) {
   532  	reservedID := sbc.reserve(ctx, target, preQueries, bindVariables, 0, options)
   533  	state, err := sbc.BeginStreamExecute(ctx, target, postBeginQueries, sql, bindVariables, reservedID, options, callback)
   534  	if state.TransactionID != 0 {
   535  		sbc.setTxReservedID(state.TransactionID, reservedID)
   536  	}
   537  	return queryservice.ReservedTransactionState{
   538  		ReservedID:    reservedID,
   539  		TransactionID: state.TransactionID,
   540  		TabletAlias:   state.TabletAlias,
   541  	}, err
   542  }
   543  
   544  // ReserveExecute implements the QueryService interface
   545  func (sbc *SandboxConn) ReserveExecute(ctx context.Context, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, options *querypb.ExecuteOptions) (queryservice.ReservedState, *sqltypes.Result, error) {
   546  	reservedID := sbc.reserve(ctx, target, preQueries, bindVariables, transactionID, options)
   547  	result, err := sbc.Execute(ctx, target, sql, bindVariables, transactionID, reservedID, options)
   548  	if transactionID != 0 {
   549  		sbc.setTxReservedID(transactionID, reservedID)
   550  	}
   551  	return queryservice.ReservedState{
   552  		ReservedID:  reservedID,
   553  		TabletAlias: sbc.tablet.Alias,
   554  	}, result, err
   555  }
   556  
   557  // ReserveStreamExecute is part of the QueryService interface.
   558  func (sbc *SandboxConn) ReserveStreamExecute(ctx context.Context, target *querypb.Target, preQueries []string, sql string, bindVariables map[string]*querypb.BindVariable, transactionID int64, options *querypb.ExecuteOptions, callback func(*sqltypes.Result) error) (queryservice.ReservedState, error) {
   559  	reservedID := sbc.reserve(ctx, target, preQueries, bindVariables, transactionID, options)
   560  	err := sbc.StreamExecute(ctx, target, sql, bindVariables, transactionID, reservedID, options, callback)
   561  	if transactionID != 0 {
   562  		sbc.setTxReservedID(transactionID, reservedID)
   563  	}
   564  	return queryservice.ReservedState{
   565  		ReservedID:  reservedID,
   566  		TabletAlias: sbc.tablet.Alias,
   567  	}, err
   568  }
   569  
   570  func (sbc *SandboxConn) reserve(ctx context.Context, target *querypb.Target, preQueries []string, bindVariables map[string]*querypb.BindVariable, transactionID int64, options *querypb.ExecuteOptions) int64 {
   571  	sbc.ReserveCount.Add(1)
   572  	for _, query := range preQueries {
   573  		sbc.Execute(ctx, target, query, bindVariables, transactionID, 0, options)
   574  	}
   575  	if transactionID != 0 {
   576  		return transactionID
   577  	}
   578  	return sbc.ReserveID.Add(1)
   579  }
   580  
   581  // Release implements the QueryService interface
   582  func (sbc *SandboxConn) Release(ctx context.Context, target *querypb.Target, transactionID, reservedID int64) error {
   583  	sbc.ReleaseCount.Add(1)
   584  	return sbc.getError()
   585  }
   586  
   587  // GetSchema implements the QueryService interface
   588  func (sbc *SandboxConn) GetSchema(ctx context.Context, target *querypb.Target, tableType querypb.SchemaTableType, tableNames []string, callback func(schemaRes *querypb.GetSchemaResponse) error) error {
   589  	sbc.GetSchemaCount.Add(1)
   590  	if len(sbc.getSchemaResult) == 0 {
   591  		return nil
   592  	}
   593  	resp := sbc.getSchemaResult[0]
   594  	sbc.getSchemaResult = sbc.getSchemaResult[1:]
   595  	return callback(&querypb.GetSchemaResponse{TableDefinition: resp})
   596  }
   597  
   598  // Close does not change ExecCount
   599  func (sbc *SandboxConn) Close(ctx context.Context) error {
   600  	return nil
   601  }
   602  
   603  // Tablet is part of the QueryService interface.
   604  func (sbc *SandboxConn) Tablet() *topodatapb.Tablet {
   605  	return sbc.tablet
   606  }
   607  
   608  // ChangeTabletType changes the tablet type.
   609  func (sbc *SandboxConn) ChangeTabletType(typ topodatapb.TabletType) {
   610  	sbc.tablet.Type = typ
   611  }
   612  
   613  func (sbc *SandboxConn) getNextResult(stmt sqlparser.Statement) *sqltypes.Result {
   614  	switch stmt.(type) {
   615  	case *sqlparser.Savepoint,
   616  		*sqlparser.SRollback,
   617  		*sqlparser.Release:
   618  		return &sqltypes.Result{}
   619  	}
   620  	if len(sbc.results) != 0 {
   621  		r := sbc.results[0]
   622  		sbc.results = sbc.results[1:]
   623  		return r
   624  	}
   625  	if stmt == nil {
   626  		// if we didn't get a valid query, we'll assume we need a SELECT
   627  		return getSingleRowResult()
   628  	}
   629  	switch stmt.(type) {
   630  	case *sqlparser.Select,
   631  		*sqlparser.Union,
   632  		*sqlparser.Show,
   633  		sqlparser.Explain,
   634  		*sqlparser.OtherRead:
   635  		return getSingleRowResult()
   636  	case *sqlparser.Set,
   637  		sqlparser.DDLStatement,
   638  		*sqlparser.AlterVschema,
   639  		*sqlparser.Use,
   640  		*sqlparser.OtherAdmin:
   641  		return &sqltypes.Result{}
   642  	}
   643  
   644  	// for everything else we fake a single row being affected
   645  	return &sqltypes.Result{RowsAffected: 1}
   646  }
   647  
   648  func (sbc *SandboxConn) setTxReservedID(transactionID int64, reservedID int64) {
   649  	sbc.mapMu.Lock()
   650  	defer sbc.mapMu.Unlock()
   651  	sbc.txIDToRID[transactionID] = reservedID
   652  }
   653  
   654  func (sbc *SandboxConn) getTxReservedID(txID int64) int64 {
   655  	sbc.mapMu.Lock()
   656  	defer sbc.mapMu.Unlock()
   657  	return sbc.txIDToRID[txID]
   658  }
   659  
   660  // StringQueries returns the queries executed as a slice of strings
   661  func (sbc *SandboxConn) StringQueries() []string {
   662  	result := make([]string, len(sbc.Queries))
   663  	for i, query := range sbc.Queries {
   664  		result[i] = query.Sql
   665  	}
   666  	return result
   667  }
   668  
   669  // getSingleRowResult is used to get a SingleRowResult but it creates separate fields because some tests change the fields
   670  // If these fields are not created separately then the constants value also changes which leads to some other tests failing later
   671  func getSingleRowResult() *sqltypes.Result {
   672  	singleRowResult := &sqltypes.Result{
   673  		InsertID:    SingleRowResult.InsertID,
   674  		StatusFlags: SingleRowResult.StatusFlags,
   675  		Rows:        SingleRowResult.Rows,
   676  	}
   677  
   678  	fields := SingleRowResult.Fields
   679  	for _, field := range fields {
   680  		singleRowResult.Fields = append(singleRowResult.Fields, &querypb.Field{
   681  			Name: field.Name,
   682  			Type: field.Type,
   683  		})
   684  	}
   685  
   686  	return singleRowResult
   687  }
   688  
   689  // SingleRowResult is returned when there is no pre-stored result.
   690  var SingleRowResult = &sqltypes.Result{
   691  	Fields: []*querypb.Field{
   692  		{Name: "id", Type: sqltypes.Int32},
   693  		{Name: "value", Type: sqltypes.VarChar},
   694  	},
   695  	InsertID: 0,
   696  	Rows: [][]sqltypes.Value{{
   697  		sqltypes.NewInt32(1),
   698  		sqltypes.NewVarChar("foo"),
   699  	}},
   700  	StatusFlags: sqltypes.ServerStatusAutocommit,
   701  }
   702  
   703  // StreamRowResult is SingleRowResult with RowsAffected set to 0.
   704  var StreamRowResult = &sqltypes.Result{
   705  	Fields: []*querypb.Field{
   706  		{Name: "id", Type: sqltypes.Int32},
   707  		{Name: "value", Type: sqltypes.VarChar},
   708  	},
   709  	Rows: [][]sqltypes.Value{{
   710  		sqltypes.NewInt32(1),
   711  		sqltypes.NewVarChar("foo"),
   712  	}},
   713  }