vitess.io/vitess@v0.16.2/go/vtbench/client.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 vtbench
    18  
    19  import (
    20  	"context"
    21  	"fmt"
    22  	"sync"
    23  
    24  	"google.golang.org/grpc"
    25  
    26  	"vitess.io/vitess/go/mysql"
    27  	"vitess.io/vitess/go/sqltypes"
    28  	"vitess.io/vitess/go/vt/grpcclient"
    29  	"vitess.io/vitess/go/vt/key"
    30  	"vitess.io/vitess/go/vt/topo/topoproto"
    31  	"vitess.io/vitess/go/vt/vtgate/vtgateconn"
    32  	"vitess.io/vitess/go/vt/vttablet/queryservice"
    33  	"vitess.io/vitess/go/vt/vttablet/tabletconn"
    34  
    35  	querypb "vitess.io/vitess/go/vt/proto/query"
    36  	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
    37  )
    38  
    39  type clientConn interface {
    40  	connect(ctx context.Context, cp ConnParams) error
    41  	execute(ctx context.Context, query string, bindVars map[string]*querypb.BindVariable) (*sqltypes.Result, error)
    42  }
    43  
    44  type mysqlClientConn struct {
    45  	conn *mysql.Conn
    46  }
    47  
    48  func (c *mysqlClientConn) connect(ctx context.Context, cp ConnParams) error {
    49  	conn, err := mysql.Connect(ctx, &mysql.ConnParams{
    50  		Host:       cp.Hosts[0],
    51  		Port:       cp.Port,
    52  		DbName:     cp.DB,
    53  		Uname:      cp.Username,
    54  		Pass:       cp.Password,
    55  		UnixSocket: cp.UnixSocket,
    56  	})
    57  
    58  	if err != nil {
    59  		return err
    60  	}
    61  
    62  	c.conn = conn
    63  
    64  	return nil
    65  }
    66  
    67  func (c *mysqlClientConn) execute(ctx context.Context, query string, bindVars map[string]*querypb.BindVariable) (*sqltypes.Result, error) {
    68  	if len(bindVars) != 0 {
    69  		panic("cannot have bind vars for mysql protocol")
    70  	}
    71  
    72  	return c.conn.ExecuteFetch(query, 10000 /* maxrows */, true /* wantfields*/)
    73  }
    74  
    75  // used to ensure grpc.WithBlock() is added once to the options
    76  var withBlockOnce sync.Once
    77  
    78  type grpcVtgateConn struct {
    79  	session *vtgateconn.VTGateSession
    80  }
    81  
    82  var vtgateConns = map[string]*vtgateconn.VTGateConn{}
    83  
    84  func (c *grpcVtgateConn) connect(ctx context.Context, cp ConnParams) error {
    85  	withBlockOnce.Do(func() {
    86  		grpcclient.RegisterGRPCDialOptions(func(opts []grpc.DialOption) ([]grpc.DialOption, error) {
    87  			return append(opts, grpc.WithBlock()), nil
    88  		})
    89  	})
    90  
    91  	address := fmt.Sprintf("%v:%v", cp.Hosts[0], cp.Port)
    92  
    93  	conn, ok := vtgateConns[address]
    94  	if !ok {
    95  		var err error
    96  		conn, err = vtgateconn.DialProtocol(ctx, "grpc", address)
    97  		if err != nil {
    98  			return err
    99  		}
   100  		vtgateConns[address] = conn
   101  	}
   102  
   103  	c.session = conn.Session(cp.DB, nil)
   104  
   105  	return nil
   106  }
   107  
   108  func (c *grpcVtgateConn) execute(ctx context.Context, query string, bindVars map[string]*querypb.BindVariable) (*sqltypes.Result, error) {
   109  	return c.session.Execute(ctx, query, bindVars)
   110  }
   111  
   112  type grpcVttabletConn struct {
   113  	qs     queryservice.QueryService
   114  	target querypb.Target
   115  }
   116  
   117  var vttabletConns = map[string]queryservice.QueryService{}
   118  
   119  func (c *grpcVttabletConn) connect(ctx context.Context, cp ConnParams) error {
   120  	withBlockOnce.Do(func() {
   121  		grpcclient.RegisterGRPCDialOptions(func(opts []grpc.DialOption) ([]grpc.DialOption, error) {
   122  			return append(opts, grpc.WithBlock()), nil
   123  		})
   124  	})
   125  
   126  	// parse the "db" into the keyspace/shard target
   127  	keyspace, tabletType, dest, err := topoproto.ParseDestination(cp.DB, topodatapb.TabletType_PRIMARY)
   128  	if err != nil {
   129  		return err
   130  	}
   131  
   132  	qs, ok := vttabletConns[cp.Hosts[0]]
   133  	if !ok {
   134  		tablet := topodatapb.Tablet{
   135  			Hostname: cp.Hosts[0],
   136  			PortMap:  map[string]int32{"grpc": int32(cp.Port)},
   137  			Keyspace: keyspace,
   138  		}
   139  		var err error
   140  		qs, err = tabletconn.GetDialer()(&tablet, true)
   141  		if err != nil {
   142  			return err
   143  		}
   144  		vttabletConns[cp.Hosts[0]] = qs
   145  	}
   146  
   147  	c.qs = qs
   148  
   149  	shard, ok := dest.(key.DestinationShard)
   150  	if !ok {
   151  		return fmt.Errorf("invalid destination shard")
   152  	}
   153  
   154  	c.target = querypb.Target{
   155  		Keyspace:   keyspace,
   156  		Shard:      string(shard),
   157  		TabletType: tabletType,
   158  	}
   159  
   160  	return nil
   161  }
   162  
   163  func (c *grpcVttabletConn) execute(ctx context.Context, query string, bindVars map[string]*querypb.BindVariable) (*sqltypes.Result, error) {
   164  	return c.qs.Execute(ctx, &c.target, query, bindVars, 0, 0, nil)
   165  }