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 }