// NewQueryResultReaderForTablet creates a new QueryResultReader for
// the provided tablet / sql query
func NewQueryResultReaderForTablet(ctx context.Context, ts topo.Server, tabletAlias topo.TabletAlias, sql string) (*QueryResultReader, error) {
tablet, err := ts.GetTablet(ctx, tabletAlias)
if err != nil {
return nil, err
}
endPoint, err := tablet.EndPoint()
if err != nil {
return nil, err
}
conn, err := tabletconn.GetDialer()(ctx, *endPoint, tablet.Keyspace, tablet.Shard, *remoteActionsTimeout)
if err != nil {
return nil, err
}
sr, clientErrFn, err := conn.StreamExecute(ctx, sql, make(map[string]interface{}), 0)
if err != nil {
return nil, err
}
// read the columns, or grab the error
cols, ok := <-sr
if !ok {
return nil, fmt.Errorf("Cannot read Fields for query '%v': %v", sql, clientErrFn())
}
return &QueryResultReader{
Output: sr,
Fields: cols.Fields,
conn: conn,
clientErrFn: clientErrFn,
}, nil
}
func commandVtTabletRollback(ctx context.Context, wr *wrangler.Wrangler, subFlags *flag.FlagSet, args []string) error {
connectTimeout := subFlags.Duration("connect_timeout", 30*time.Second, "Connection timeout for vttablet client")
if err := subFlags.Parse(args); err != nil {
return err
}
if subFlags.NArg() != 2 {
return fmt.Errorf("the <tablet_alias> and <transaction_id> arguments are required for the VtTabletRollback command")
}
transactionID, err := strconv.ParseInt(subFlags.Arg(1), 10, 64)
if err != nil {
return err
}
tabletAlias, err := topoproto.ParseTabletAlias(subFlags.Arg(0))
if err != nil {
return err
}
tabletInfo, err := wr.TopoServer().GetTablet(ctx, tabletAlias)
if err != nil {
return err
}
conn, err := tabletconn.GetDialer()(tabletInfo.Tablet, *connectTimeout)
if err != nil {
return fmt.Errorf("cannot connect to tablet %v: %v", tabletAlias, err)
}
defer conn.Close(ctx)
return conn.Rollback(ctx, &querypb.Target{
Keyspace: tabletInfo.Tablet.Keyspace,
Shard: tabletInfo.Tablet.Shard,
TabletType: tabletInfo.Tablet.Type,
}, transactionID)
}
// getConn reuses an existing connection if possible. Otherwise
// it returns a connection which it will save for future reuse.
// If it returns an error, retry will tell you if getConn can be retried.
// If the context has a deadline and exceeded, it returns error and no-retry immediately.
func (sdc *ShardConn) getConn(ctx context.Context) (conn tabletconn.TabletConn, endPoint topo.EndPoint, err error, retry bool) {
sdc.mu.Lock()
defer sdc.mu.Unlock()
// fail-fast if deadline exceeded
deadline, ok := ctx.Deadline()
if ok {
if time.Now().After(deadline) {
return nil, topo.EndPoint{}, tabletconn.OperationalError("vttablet: deadline exceeded"), false
}
}
if sdc.conn != nil {
return sdc.conn, sdc.conn.EndPoint(), nil, false
}
endPoint, err = sdc.balancer.Get()
if err != nil {
return nil, topo.EndPoint{}, err, false
}
conn, err = tabletconn.GetDialer()(ctx, endPoint, sdc.keyspace, sdc.shard, sdc.timeout)
if err != nil {
sdc.balancer.MarkDown(endPoint.Uid, err.Error())
return nil, endPoint, err, true
}
sdc.conn = conn
return sdc.conn, endPoint, nil, false
}
// getTablet (re)sets the tablet which is used for the streaming query.
// If the method returns an error, the first return value specifies if it is
// okay to retry.
func (r *RestartableResultReader) getTablet() (bool, error) {
if r.tablet != nil {
// If there was a tablet before, return it to the tabletProvider.
r.Close()
r.tablet = nil
r.conn = nil
r.fields = nil
r.output = nil
}
// Get a tablet from the tablet provider.
tablet, err := r.tp.getTablet()
if err != nil {
return true /* retryable */, fmt.Errorf("failed get tablet for streaming query: %v", err)
}
// Get the dialer for it.
conn, err := tabletconn.GetDialer()(tablet, *remoteActionsTimeout)
if err != nil {
return false /* retryable */, fmt.Errorf("failed to get dialer for tablet: %v", err)
}
r.tablet = tablet
r.conn = conn
return false /* retryable */, nil
}
// NewRestartableResultReader creates a new RestartableResultReader for
// the provided tablet and chunk.
// It will automatically create the necessary query to read all rows within
// the chunk.
// NOTE: We assume that the Columns field in "td" was ordered by a preceding
// call to reorderColumnsPrimaryKeyFirst().
func NewRestartableResultReader(ctx context.Context, logger logutil.Logger, ts topo.Server, tabletAlias *topodatapb.TabletAlias, td *tabletmanagerdatapb.TableDefinition, chunk chunk) (*RestartableResultReader, error) {
shortCtx, cancel := context.WithTimeout(ctx, *remoteActionsTimeout)
tablet, err := ts.GetTablet(shortCtx, tabletAlias)
cancel()
if err != nil {
return nil, fmt.Errorf("tablet=%v table=%v chunk=%v: Failed to resolve tablet alias: %v", topoproto.TabletAliasString(tabletAlias), td.Name, chunk, err)
}
conn, err := tabletconn.GetDialer()(tablet.Tablet, *remoteActionsTimeout)
if err != nil {
return nil, fmt.Errorf("tablet=%v table=%v chunk=%v: Failed to get dialer for tablet: %v", topoproto.TabletAliasString(tabletAlias), td.Name, chunk, err)
}
r := &RestartableResultReader{
ctx: ctx,
logger: logger,
tablet: tablet.Tablet,
td: td,
chunk: chunk,
conn: conn,
}
if err := r.startStream(); err != nil {
return nil, err
}
logger.Infof("tablet=%v table=%v chunk=%v: Starting to stream rows using query '%v'.", topoproto.TabletAliasString(tabletAlias), td.Name, chunk, r.query)
return r, nil
}
// TestErrorSuite runs all the tests that expect errors
func TestErrorSuite(t *testing.T, protocol string, endPoint *topodatapb.EndPoint, fake *FakeQueryService) {
// make sure we use the right client
*tabletconn.TabletProtocol = protocol
// create a connection, using sessionId
ctx := context.Background()
conn, err := tabletconn.GetDialer()(ctx, endPoint, testTarget.Keyspace, testTarget.Shard, topodatapb.TabletType_UNKNOWN, 30*time.Second)
if err != nil {
t.Fatalf("dial failed: %v", err)
}
// fake should return an error, make sure errors are handled properly
fake.hasError = true
testBeginError(t, conn)
testCommitError(t, conn)
testRollbackError(t, conn)
testExecuteError(t, conn)
testStreamExecuteError(t, conn, fake)
testExecuteBatchError(t, conn)
testSplitQueryError(t, conn)
testBegin2Error(t, conn)
testCommit2Error(t, conn)
testRollback2Error(t, conn)
testExecute2Error(t, conn)
testStreamExecute2Error(t, conn, fake)
testExecuteBatch2Error(t, conn)
fake.hasError = false
conn.Close()
}
// addL2VTGateConn adds a backend l2vtgate for the provided keyspace / shard.
func (lg *l2VTGateGateway) addL2VTGateConn(addr, keyspace, shard string) error {
lg.mu.Lock()
defer lg.mu.Unlock()
// extract keyrange if it's a range
canonical, kr, err := topo.ValidateShardName(shard)
if err != nil {
return fmt.Errorf("error parsing shard name %v: %v", shard, err)
}
// check for duplicates
for _, c := range lg.connMap[keyspace] {
if c.shard == canonical {
return fmt.Errorf("duplicate %v/%v entry", keyspace, shard)
}
}
// Dial in the background
conn, err := tabletconn.GetDialer()(&topodatapb.Tablet{
Hostname: addr,
}, 0)
if err != nil {
return err
}
lg.connMap[keyspace] = append(lg.connMap[keyspace], &l2VTGateConn{
addr: addr,
keyspace: keyspace,
shard: canonical,
keyRange: kr,
conn: conn,
})
return nil
}
// NewQueryResultReaderForTablet creates a new QueryResultReader for
// the provided tablet / sql query
func NewQueryResultReaderForTablet(ctx context.Context, ts topo.Server, tabletAlias *topodatapb.TabletAlias, sql string) (*QueryResultReader, error) {
shortCtx, cancel := context.WithTimeout(ctx, *remoteActionsTimeout)
tablet, err := ts.GetTablet(shortCtx, tabletAlias)
cancel()
if err != nil {
return nil, err
}
conn, err := tabletconn.GetDialer()(tablet.Tablet, *remoteActionsTimeout)
if err != nil {
return nil, err
}
stream, err := conn.StreamExecute(ctx, &querypb.Target{
Keyspace: tablet.Tablet.Keyspace,
Shard: tablet.Tablet.Shard,
TabletType: tablet.Tablet.Type,
}, sql, make(map[string]interface{}), nil)
if err != nil {
return nil, err
}
// read the columns, or grab the error
cols, err := stream.Recv()
if err != nil {
return nil, fmt.Errorf("Cannot read Fields for query '%v': %v", sql, err)
}
return &QueryResultReader{
output: stream,
fields: cols.Fields,
conn: conn,
}, nil
}
func commandVtTabletBegin(ctx context.Context, wr *wrangler.Wrangler, subFlags *flag.FlagSet, args []string) error {
connectTimeout := subFlags.Duration("connect_timeout", 30*time.Second, "Connection timeout for vttablet client")
if err := subFlags.Parse(args); err != nil {
return err
}
if subFlags.NArg() != 1 {
return fmt.Errorf("the <tablet_alias> argument is required for the VtTabletBegin command")
}
tabletAlias, err := topoproto.ParseTabletAlias(subFlags.Arg(0))
if err != nil {
return err
}
tabletInfo, err := wr.TopoServer().GetTablet(ctx, tabletAlias)
if err != nil {
return err
}
conn, err := tabletconn.GetDialer()(ctx, tabletInfo.Tablet, *connectTimeout)
if err != nil {
return fmt.Errorf("cannot connect to tablet %v: %v", tabletAlias, err)
}
defer conn.Close()
transactionID, err := conn.Begin(ctx)
if err != nil {
return fmt.Errorf("Begin failed: %v", err)
}
result := map[string]int64{
"transaction_id": transactionID,
}
return printJSON(wr.Logger(), result)
}
func (th *tabletHealth) stream(ctx context.Context, ts topo.Server, tabletAlias *topodatapb.TabletAlias) (err error) {
defer func() {
th.mu.Lock()
th.err = err
th.mu.Unlock()
close(th.done)
}()
ti, err := ts.GetTablet(ctx, tabletAlias)
if err != nil {
return err
}
ep, err := topo.TabletEndPoint(ti.Tablet)
if err != nil {
return err
}
// Pass in a tablet type that is not UNKNOWN, so we don't ask
// for sessionId.
conn, err := tabletconn.GetDialer()(ctx, ep, "", "", topodatapb.TabletType_MASTER, 30*time.Second)
if err != nil {
return err
}
defer conn.Close()
stream, err := conn.StreamHealth(ctx)
if err != nil {
return err
}
first := true
for time.Since(th.lastAccessed()) < *tabletHealthKeepAlive {
select {
case <-ctx.Done():
return ctx.Err()
default:
}
result, err := stream.Recv()
if err != nil {
return err
}
th.mu.Lock()
th.result = result
th.mu.Unlock()
if first {
// We got the first result, so we're ready to be accessed.
close(th.ready)
first = false
}
}
return nil
}
// getNewConn creates a new tablet connection with a separate per conn timeout.
// It limits the overall timeout to connTimeoutTotal by checking elapsed time after each blocking call.
func (sdc *ShardConn) getNewConn(ctx context.Context) (conn tabletconn.TabletConn, endPoint *topodatapb.EndPoint, isTimeout bool, err error) {
startTime := time.Now()
endPoints, err := sdc.balancer.Get()
if err != nil {
// Error when getting endpoint
return nil, nil, false, err
}
if len(endPoints) == 0 {
// No valid endpoint
return nil, nil, false, vterrors.FromError(
vtrpcpb.ErrorCode_INTERNAL_ERROR,
fmt.Errorf("no valid endpoint"),
)
}
if time.Now().Sub(startTime) >= sdc.connTimeoutTotal {
return nil, nil, true, vterrors.FromError(
vtrpcpb.ErrorCode_DEADLINE_EXCEEDED,
fmt.Errorf("timeout when getting endpoints"),
)
}
// Iterate through all endpoints to create a connection
perConnTimeout := sdc.getConnTimeoutPerConn(len(endPoints))
allErrors := new(concurrency.AllErrorRecorder)
for _, endPoint := range endPoints {
perConnStartTime := time.Now()
conn, err = tabletconn.GetDialer()(ctx, endPoint, sdc.keyspace, sdc.shard, topodatapb.TabletType_UNKNOWN, perConnTimeout)
if err == nil {
sdc.connectTimings.Record([]string{sdc.keyspace, sdc.shard, strings.ToLower(sdc.tabletType.String())}, perConnStartTime)
sdc.mu.Lock()
defer sdc.mu.Unlock()
sdc.conn = conn
return conn, endPoint, false, nil
}
// Markdown the endpoint if it failed to connect
sdc.balancer.MarkDown(endPoint.Uid, err.Error())
vtErr := vterrors.NewVitessError(
// TODO(aaijazi): what about OperationalErrors here?
vterrors.RecoverVtErrorCode(err), err,
"%v %+v", err, endPoint,
)
allErrors.RecordError(vtErr)
if time.Now().Sub(startTime) >= sdc.connTimeoutTotal {
err = vterrors.FromError(
vtrpcpb.ErrorCode_DEADLINE_EXCEEDED,
fmt.Errorf("timeout when connecting to %+v", endPoint),
)
allErrors.RecordError(err)
return nil, nil, true, allErrors.AggrError(AggregateVtGateErrors)
}
}
return nil, nil, false, allErrors.Error()
}
func (th *tabletHealth) stream(ctx context.Context, ts topo.Server, tabletAlias topo.TabletAlias) (err error) {
defer func() {
th.mu.Lock()
th.err = err
th.mu.Unlock()
close(th.done)
}()
ti, err := ts.GetTablet(ctx, tabletAlias)
if err != nil {
return err
}
ep, err := ti.EndPoint()
if err != nil {
return err
}
// pass in empty keyspace and shard to not ask for sessionId
conn, err := tabletconn.GetDialer()(ctx, *ep, "", "", 30*time.Second)
if err != nil {
return err
}
defer conn.Close()
stream, errFunc, err := conn.StreamHealth(ctx)
if err != nil {
return err
}
first := true
for time.Since(th.lastAccessed()) < *tabletHealthKeepAlive {
select {
case <-ctx.Done():
return ctx.Err()
case result, ok := <-stream:
if !ok {
return errFunc()
}
th.mu.Lock()
th.result = result
th.mu.Unlock()
if first {
// We got the first result, so we're ready to be accessed.
close(th.ready)
first = false
}
}
}
return nil
}
func (th *tabletHealth) stream(ctx context.Context, ts topo.Server, tabletAlias *topodatapb.TabletAlias) (err error) {
defer func() {
th.mu.Lock()
th.err = err
th.mu.Unlock()
close(th.done)
}()
ti, err := ts.GetTablet(ctx, tabletAlias)
if err != nil {
return err
}
conn, err := tabletconn.GetDialer()(ti.Tablet, 30*time.Second)
if err != nil {
return err
}
defer conn.Close(ctx)
stream, err := conn.StreamHealth(ctx)
if err != nil {
return err
}
first := true
for time.Since(th.lastAccessed()) < *tabletHealthKeepAlive {
select {
case <-ctx.Done():
return ctx.Err()
default:
}
result, err := stream.Recv()
if err != nil {
return err
}
th.mu.Lock()
th.result = result
th.mu.Unlock()
if first {
// We got the first result, so we're ready to be accessed.
close(th.ready)
first = false
}
}
return nil
}
// connect creates connection to the endpoint and starts streaming.
func (hcc *healthCheckConn) connect(ctx context.Context, hc *HealthCheckImpl, endPoint *pbt.EndPoint) (<-chan *pbq.StreamHealthResponse, tabletconn.ErrFunc, error) {
conn, err := tabletconn.GetDialer()(ctx, endPoint, "" /*keyspace*/, "" /*shard*/, pbt.TabletType_RDONLY, hc.connTimeout)
if err != nil {
return nil, nil, err
}
stream, errfunc, err := conn.StreamHealth(ctx)
if err != nil {
conn.Close()
return nil, nil, err
}
hcc.mu.Lock()
hcc.conn = conn
hcc.mu.Unlock()
return stream, errfunc, nil
}
func commandVtTabletExecute(ctx context.Context, wr *wrangler.Wrangler, subFlags *flag.FlagSet, args []string) error {
transactionID := subFlags.Int("transaction_id", 0, "transaction id to use, if inside a transaction.")
bindVariables := newBindvars(subFlags)
keyspace := subFlags.String("keyspace", "", "keyspace the tablet belongs to")
shard := subFlags.String("shard", "", "shard the tablet belongs to")
tabletType := subFlags.String("tablet_type", "unknown", "tablet type we expect from the tablet (use unknown to use sessionId)")
connectTimeout := subFlags.Duration("connect_timeout", 30*time.Second, "Connection timeout for vttablet client")
json := subFlags.Bool("json", false, "Output JSON instead of human-readable table")
if err := subFlags.Parse(args); err != nil {
return err
}
if subFlags.NArg() != 2 {
return fmt.Errorf("the <tablet_alias> and <sql> arguments are required for the VtTabletExecute command")
}
tt, err := topoproto.ParseTabletType(*tabletType)
if err != nil {
return err
}
tabletAlias, err := topoproto.ParseTabletAlias(subFlags.Arg(0))
if err != nil {
return err
}
tabletInfo, err := wr.TopoServer().GetTablet(ctx, tabletAlias)
if err != nil {
return err
}
ep, err := topo.TabletEndPoint(tabletInfo.Tablet)
if err != nil {
return fmt.Errorf("cannot get EndPoint from tablet record: %v", err)
}
conn, err := tabletconn.GetDialer()(ctx, ep, *keyspace, *shard, tt, *connectTimeout)
if err != nil {
return fmt.Errorf("cannot connect to tablet %v: %v", tabletAlias, err)
}
defer conn.Close()
qr, err := conn.Execute(ctx, subFlags.Arg(1), *bindVariables, int64(*transactionID))
if err != nil {
return fmt.Errorf("Execute failed: %v", err)
}
if *json {
return printJSON(wr.Logger(), qr)
}
printQueryResult(loggerWriter{wr.Logger()}, qr)
return nil
}
// connect creates connection to the endpoint and starts streaming.
func (hcc *healthCheckConn) connect(hc *HealthCheckImpl, endPoint *topodatapb.EndPoint) (tabletconn.StreamHealthReader, error) {
conn, err := tabletconn.GetDialer()(hcc.ctx, endPoint, "" /*keyspace*/, "" /*shard*/, topodatapb.TabletType_RDONLY, hc.connTimeout)
if err != nil {
return nil, err
}
stream, err := conn.StreamHealth(hcc.ctx)
if err != nil {
conn.Close()
return nil, err
}
hcc.mu.Lock()
hcc.conn = conn
hcc.lastError = nil
hcc.mu.Unlock()
return stream, nil
}
func commandVtTabletUpdateStream(ctx context.Context, wr *wrangler.Wrangler, subFlags *flag.FlagSet, args []string) error {
count := subFlags.Int("count", 1, "number of responses to wait for")
timestamp := subFlags.Int("timestamp", 0, "timestamp to start the stream from")
position := subFlags.String("position", "", "position to start the stream from")
connectTimeout := subFlags.Duration("connect_timeout", 30*time.Second, "Connection timeout for vttablet client")
if err := subFlags.Parse(args); err != nil {
return err
}
if subFlags.NArg() != 1 {
return fmt.Errorf("The <tablet alias> argument is required for the VtTabletUpdateStream command.")
}
tabletAlias, err := topoproto.ParseTabletAlias(subFlags.Arg(0))
if err != nil {
return err
}
tabletInfo, err := wr.TopoServer().GetTablet(ctx, tabletAlias)
if err != nil {
return err
}
conn, err := tabletconn.GetDialer()(tabletInfo.Tablet, *connectTimeout)
if err != nil {
return fmt.Errorf("cannot connect to tablet %v: %v", tabletAlias, err)
}
stream, err := conn.UpdateStream(ctx, &querypb.Target{
Keyspace: tabletInfo.Tablet.Keyspace,
Shard: tabletInfo.Tablet.Shard,
TabletType: tabletInfo.Tablet.Type,
}, *position, int64(*timestamp))
if err != nil {
return err
}
for i := 0; i < *count; i++ {
se, err := stream.Recv()
if err != nil {
return fmt.Errorf("stream ended early: %v", err)
}
data, err := json.Marshal(se)
if err != nil {
wr.Logger().Errorf("cannot json-marshal structure: %v", err)
} else {
wr.Logger().Printf("%v\n", string(data))
}
}
return nil
}
func commandVtTabletStreamHealth(ctx context.Context, wr *wrangler.Wrangler, subFlags *flag.FlagSet, args []string) error {
count := subFlags.Int("count", 1, "number of responses to wait for")
connectTimeout := subFlags.Duration("connect_timeout", 30*time.Second, "Connection timeout for vttablet client")
if err := subFlags.Parse(args); err != nil {
return err
}
if subFlags.NArg() != 1 {
return fmt.Errorf("The <tablet alias> argument is required for the VtTabletStreamHealth command.")
}
tabletAlias, err := topoproto.ParseTabletAlias(subFlags.Arg(0))
if err != nil {
return err
}
tabletInfo, err := wr.TopoServer().GetTablet(ctx, tabletAlias)
if err != nil {
return err
}
ep, err := topo.TabletEndPoint(tabletInfo.Tablet)
if err != nil {
return fmt.Errorf("cannot get EndPoint from tablet record: %v", err)
}
// pass in a non-UNKNOWN tablet type to not use sessionId
conn, err := tabletconn.GetDialer()(ctx, ep, "", "", pb.TabletType_MASTER, *connectTimeout)
if err != nil {
return fmt.Errorf("cannot connect to tablet %v: %v", tabletAlias, err)
}
stream, errFunc, err := conn.StreamHealth(ctx)
if err != nil {
return err
}
for i := 0; i < *count; i++ {
shr, ok := <-stream
if !ok {
return fmt.Errorf("stream ended early: %v", errFunc())
}
data, err := json.Marshal(shr)
if err != nil {
wr.Logger().Errorf("cannot json-marshal structure: %v", err)
} else {
wr.Logger().Printf("%v\n", string(data))
}
}
return nil
}
func commandVtTabletExecute(ctx context.Context, wr *wrangler.Wrangler, subFlags *flag.FlagSet, args []string) error {
transactionID := subFlags.Int("transaction_id", 0, "transaction id to use, if inside a transaction.")
bindVariables := newBindvars(subFlags)
connectTimeout := subFlags.Duration("connect_timeout", 30*time.Second, "Connection timeout for vttablet client")
options := subFlags.String("options", "", "execute options values as a text encoded proto of the ExecuteOptions structure")
json := subFlags.Bool("json", false, "Output JSON instead of human-readable table")
if err := subFlags.Parse(args); err != nil {
return err
}
if subFlags.NArg() != 2 {
return fmt.Errorf("the <tablet_alias> and <sql> arguments are required for the VtTabletExecute command")
}
tabletAlias, err := topoproto.ParseTabletAlias(subFlags.Arg(0))
if err != nil {
return err
}
tabletInfo, err := wr.TopoServer().GetTablet(ctx, tabletAlias)
if err != nil {
return err
}
executeOptions, err := parseExecuteOptions(*options)
if err != nil {
return err
}
conn, err := tabletconn.GetDialer()(tabletInfo.Tablet, *connectTimeout)
if err != nil {
return fmt.Errorf("cannot connect to tablet %v: %v", tabletAlias, err)
}
defer conn.Close(ctx)
qr, err := conn.Execute(ctx, &querypb.Target{
Keyspace: tabletInfo.Tablet.Keyspace,
Shard: tabletInfo.Tablet.Shard,
TabletType: tabletInfo.Tablet.Type,
}, subFlags.Arg(1), *bindVariables, int64(*transactionID), executeOptions)
if err != nil {
return fmt.Errorf("Execute failed: %v", err)
}
if *json {
return printJSON(wr.Logger(), qr)
}
printQueryResult(loggerWriter{wr.Logger()}, qr)
return nil
}
// connect creates connection to the tablet and starts streaming.
func (hcc *healthCheckConn) connect(hc *HealthCheckImpl) (tabletconn.StreamHealthReader, error) {
// Keyspace, shard and tabletType are the ones from the tablet
// record, but they won't be used just yet.
conn, err := tabletconn.GetDialer()(hcc.tabletStats.Tablet, hc.connTimeout)
if err != nil {
return nil, err
}
stream, err := conn.StreamHealth(hcc.ctx)
if err != nil {
conn.Close()
return nil, err
}
hcc.mu.Lock()
hcc.conn = conn
hcc.tabletStats.LastError = nil
hcc.mu.Unlock()
return stream, nil
}
// getConn reuses an existing connection if possible. Otherwise
// it returns a connection which it will save for future reuse.
// If it returns an error, retry will tell you if getConn can be retried.
func (sdc *ShardConn) getConn(context interface{}) (conn tabletconn.TabletConn, err error, retry bool) {
sdc.mu.Lock()
defer sdc.mu.Unlock()
if sdc.conn != nil {
return sdc.conn, nil, false
}
endPoint, err := sdc.balancer.Get()
if err != nil {
return nil, err, false
}
conn, err = tabletconn.GetDialer()(context, endPoint, sdc.keyspace, sdc.shard)
if err != nil {
sdc.balancer.MarkDown(endPoint.Uid)
return nil, err, true
}
sdc.conn = conn
return sdc.conn, nil, false
}
func commandVtTabletBegin(ctx context.Context, wr *wrangler.Wrangler, subFlags *flag.FlagSet, args []string) error {
keyspace := subFlags.String("keyspace", "", "keyspace the tablet belongs to")
shard := subFlags.String("shard", "", "shard the tablet belongs to")
tabletType := subFlags.String("tablet_type", "unknown", "tablet type we expect from the tablet (use unknown to use sessionId)")
connectTimeout := subFlags.Duration("connect_timeout", 30*time.Second, "Connection timeout for vttablet client")
if err := subFlags.Parse(args); err != nil {
return err
}
if subFlags.NArg() != 1 {
return fmt.Errorf("the <tablet_alias> argument is required for the VtTabletBegin command")
}
tt, err := topoproto.ParseTabletType(*tabletType)
if err != nil {
return err
}
tabletAlias, err := topoproto.ParseTabletAlias(subFlags.Arg(0))
if err != nil {
return err
}
tabletInfo, err := wr.TopoServer().GetTablet(ctx, tabletAlias)
if err != nil {
return err
}
ep, err := topo.TabletEndPoint(tabletInfo.Tablet)
if err != nil {
return fmt.Errorf("cannot get EndPoint from tablet record: %v", err)
}
conn, err := tabletconn.GetDialer()(ctx, ep, *keyspace, *shard, tt, *connectTimeout)
if err != nil {
return fmt.Errorf("cannot connect to tablet %v: %v", tabletAlias, err)
}
defer conn.Close()
transactionID, err := conn.Begin(ctx)
if err != nil {
return fmt.Errorf("Begin failed: %v", err)
}
result := map[string]int64{
"transaction_id": transactionID,
}
return printJSON(wr, result)
}
func (conn *Conn) dial() (err error) {
// build the endpoint in the right format
host, port, err := netutil.SplitHostPort(conn.dbi.Host)
if err != nil {
return err
}
endPoint := topo.EndPoint{
Host: host,
NamedPortMap: map[string]int{
"_vtocc": port,
},
}
// and dial
tabletConn, err := tabletconn.GetDialer()(nil, endPoint, conn.keyspace(), conn.shard())
if err != nil {
return err
}
conn.tabletConn = tabletConn
return
}
// getNewConn creates a new tablet connection with a separate per conn timeout.
// It limits the overall timeout to connTimeoutTotal by checking elapsed time after each blocking call.
func (sdc *ShardConn) getNewConn(ctx context.Context) (conn tabletconn.TabletConn, endPoint *pb.EndPoint, isTimeout bool, err error) {
startTime := time.Now()
endPoints, err := sdc.balancer.Get()
if err != nil {
// Error when getting endpoint
return nil, nil, false, err
}
if len(endPoints) == 0 {
// No valid endpoint
return nil, nil, false, fmt.Errorf("no valid endpoint")
}
if time.Now().Sub(startTime) >= sdc.connTimeoutTotal {
return nil, nil, true, fmt.Errorf("timeout when getting endpoints")
}
// Iterate through all endpoints to create a connection
perConnTimeout := sdc.getConnTimeoutPerConn(len(endPoints))
allErrors := new(concurrency.AllErrorRecorder)
for _, endPoint := range endPoints {
perConnStartTime := time.Now()
conn, err = tabletconn.GetDialer()(ctx, endPoint, sdc.keyspace, sdc.shard, perConnTimeout)
if err == nil {
sdc.connectTimings.Record([]string{sdc.keyspace, sdc.shard, string(sdc.tabletType)}, perConnStartTime)
sdc.mu.Lock()
defer sdc.mu.Unlock()
sdc.conn = conn
return conn, endPoint, false, nil
}
// Markdown the endpoint if it failed to connect
sdc.balancer.MarkDown(endPoint.Uid, err.Error())
allErrors.RecordError(fmt.Errorf("%v %+v", err, endPoint))
if time.Now().Sub(startTime) >= sdc.connTimeoutTotal {
err = fmt.Errorf("timeout when connecting to %+v", endPoint)
allErrors.RecordError(err)
return nil, nil, true, allErrors.Error()
}
}
return nil, nil, false, allErrors.Error()
}
func commandVtTabletRollback(ctx context.Context, wr *wrangler.Wrangler, subFlags *flag.FlagSet, args []string) error {
keyspace := subFlags.String("keyspace", "", "keyspace the tablet belongs to")
shard := subFlags.String("shard", "", "shard the tablet belongs to")
tabletType := subFlags.String("tablet_type", "unknown", "tablet type we expect from the tablet (use unknown to use sessionId)")
connectTimeout := subFlags.Duration("connect_timeout", 30*time.Second, "Connection timeout for vttablet client")
if err := subFlags.Parse(args); err != nil {
return err
}
if subFlags.NArg() != 2 {
return fmt.Errorf("the <tablet_alias> and <transaction_id> arguments are required for the VtTabletRollback command")
}
transactionID, err := strconv.ParseInt(subFlags.Arg(1), 10, 64)
if err != nil {
return err
}
tt, err := topo.ParseTabletType(*tabletType)
if err != nil {
return err
}
tabletAlias, err := topo.ParseTabletAliasString(subFlags.Arg(0))
if err != nil {
return err
}
tabletInfo, err := wr.TopoServer().GetTablet(ctx, tabletAlias)
if err != nil {
return err
}
ep, err := topo.TabletEndPoint(tabletInfo.Tablet)
if err != nil {
return fmt.Errorf("cannot get EndPoint from tablet record: %v", err)
}
conn, err := tabletconn.GetDialer()(ctx, ep, *keyspace, *shard, tt, *connectTimeout)
if err != nil {
return fmt.Errorf("cannot connect to tablet %v: %v", tabletAlias, err)
}
defer conn.Close()
return conn.Rollback(ctx, transactionID)
}
func (th *TabletHealth) update(thc *tabletHealthCache, tabletAlias topo.TabletAlias) {
defer thc.delete(tabletAlias)
ctx := context.Background()
ti, err := thc.ts.GetTablet(ctx, tabletAlias)
if err != nil {
return
}
ep, err := ti.EndPoint()
if err != nil {
return
}
// pass in empty keyspace and shard to not ask for sessionId
conn, err := tabletconn.GetDialer()(ctx, *ep, "", "", 30*time.Second)
if err != nil {
return
}
stream, errFunc, err := conn.StreamHealth(ctx)
if err != nil {
return
}
for shr := range stream {
th.mu.Lock()
if !reflect.DeepEqual(shr, th.StreamHealthResponse) {
th.StreamHealthResponse = shr
th.Version++
th.result, th.lastError = json.MarshalIndent(th, "", " ")
}
th.mu.Unlock()
}
// we call errFunc as some implementations may use this to
// free resources.
errFunc()
}
func commandVtTabletExecute(ctx context.Context, wr *wrangler.Wrangler, subFlags *flag.FlagSet, args []string) error {
transactionID := subFlags.Int("transaction_id", 0, "transaction id to use, if inside a transaction.")
bindVariables := newBindvars(subFlags)
keyspace := subFlags.String("keyspace", "", "keyspace the tablet belongs to")
shard := subFlags.String("shard", "", "shard the tablet belongs to")
connectTimeout := subFlags.Duration("connect_timeout", 30*time.Second, "Connection timeout for vttablet client")
if err := subFlags.Parse(args); err != nil {
return err
}
if subFlags.NArg() != 2 {
return fmt.Errorf("the <tablet_alis> and <sql> arguments are required for the VtTabletExecute command")
}
tabletAlias, err := topo.ParseTabletAliasString(subFlags.Arg(0))
if err != nil {
return err
}
tabletInfo, err := wr.TopoServer().GetTablet(ctx, tabletAlias)
if err != nil {
return err
}
ep, err := tabletInfo.EndPoint()
if err != nil {
return fmt.Errorf("cannot get EndPoint from tablet record: %v", err)
}
// pass in empty keyspace and shard to not ask for sessionId
conn, err := tabletconn.GetDialer()(ctx, ep, *keyspace, *shard, *connectTimeout)
if err != nil {
return fmt.Errorf("cannot connect to tablet %v: %v", tabletAlias, err)
}
defer conn.Close()
qr, err := conn.Execute(ctx, subFlags.Arg(1), *bindVariables, int64(*transactionID))
if err != nil {
return fmt.Errorf("Execute failed: %v", err)
}
wr.Logger().Printf("%v\n", jscfg.ToJSON(qr))
return nil
}
// TestSuite runs all the tests
func TestSuite(t *testing.T, protocol string, endPoint *pbt.EndPoint, fake *FakeQueryService) {
// make sure we use the right client
*tabletconn.TabletProtocol = protocol
// create a connection, using sessionId
ctx := context.Background()
conn, err := tabletconn.GetDialer()(ctx, endPoint, testTarget.Keyspace, testTarget.Shard, pbt.TabletType_UNKNOWN, 30*time.Second)
if err != nil {
t.Fatalf("dial failed: %v", err)
}
// run the normal tests
testBegin(t, conn)
testCommit(t, conn)
testRollback(t, conn)
testExecute(t, conn)
testStreamExecute(t, conn)
testExecuteBatch(t, conn)
testSplitQuery(t, conn)
testStreamHealth(t, conn)
// fake should return an error, make sure errors are handled properly
fake.hasError = true
testBeginError(t, conn)
testCommitError(t, conn)
testRollbackError(t, conn)
testExecuteError(t, conn)
testStreamExecuteError(t, conn, fake)
testExecuteBatchError(t, conn)
testSplitQueryError(t, conn)
testBegin2Error(t, conn)
testCommit2Error(t, conn)
testRollback2Error(t, conn)
testExecute2Error(t, conn)
testStreamExecute2Error(t, conn, fake)
testExecuteBatch2Error(t, conn)
fake.hasError = false
// create a new connection that expects the extra fields
conn.Close()
conn, err = tabletconn.GetDialer()(ctx, endPoint, testTarget.Keyspace, testTarget.Shard, pbt.TabletType_REPLICA, 30*time.Second)
if err != nil {
t.Fatalf("dial failed: %v", err)
}
// run the tests that expect extra fields
fake.checkExtraFields = true
testBegin2(t, conn)
testCommit2(t, conn)
testRollback2(t, conn)
testExecute2(t, conn)
testStreamExecute2(t, conn)
testExecuteBatch2(t, conn)
testSplitQuery(t, conn)
// force panics, make sure they're caught (with extra fields)
fake.panics = true
testBegin2Panics(t, conn)
testCommit2Panics(t, conn)
testRollback2Panics(t, conn)
testExecute2Panics(t, conn)
testStreamExecute2Panics(t, conn, fake)
testExecuteBatch2Panics(t, conn)
testSplitQueryPanics(t, conn)
testStreamHealthPanics(t, conn)
// force panic without extra fields
conn.Close()
conn, err = tabletconn.GetDialer()(ctx, endPoint, testTarget.Keyspace, testTarget.Shard, pbt.TabletType_UNKNOWN, 30*time.Second)
if err != nil {
t.Fatalf("dial failed: %v", err)
}
fake.checkExtraFields = false
testBeginPanics(t, conn)
testCommitPanics(t, conn)
testRollbackPanics(t, conn)
testExecutePanics(t, conn)
testExecuteBatchPanics(t, conn)
testStreamExecutePanics(t, conn, fake)
fake.panics = false
conn.Close()
}
// TestSuite runs all the tests.
// If fake.TestingGateway is set, we only test the calls that can go through
// a gateway.
func TestSuite(t *testing.T, protocol string, tablet *topodatapb.Tablet, fake *FakeQueryService) {
tests := []func(*testing.T, tabletconn.TabletConn, *FakeQueryService){
// positive test cases
testBegin,
testCommit,
testRollback,
testExecute,
testBeginExecute,
testStreamExecute,
testExecuteBatch,
testBeginExecuteBatch,
testSplitQuery,
// error test cases
testBeginError,
testCommitError,
testRollbackError,
testExecuteError,
testBeginExecuteErrorInBegin,
testBeginExecuteErrorInExecute,
testStreamExecuteError,
testExecuteBatchError,
testBeginExecuteBatchErrorInBegin,
testBeginExecuteBatchErrorInExecuteBatch,
testSplitQueryError,
// panic test cases
testBeginPanics,
testCommitPanics,
testRollbackPanics,
testExecutePanics,
testBeginExecutePanics,
testStreamExecutePanics,
testExecuteBatchPanics,
testBeginExecuteBatchPanics,
testSplitQueryPanics,
}
if !fake.TestingGateway {
tests = append(tests, []func(*testing.T, tabletconn.TabletConn, *FakeQueryService){
// positive test cases
testStreamHealth,
// error test cases
testStreamHealthError,
// panic test cases
testStreamHealthPanics,
}...)
}
// make sure we use the right client
*tabletconn.TabletProtocol = protocol
// create a connection
conn, err := tabletconn.GetDialer()(tablet, 30*time.Second)
if err != nil {
t.Fatalf("dial failed: %v", err)
}
// run the tests
for _, c := range tests {
c(t, conn, fake)
}
// and we're done
conn.Close()
}
