func TestStreamSelectKeyrange(t *testing.T) {
router, _, _, _ := createRouterEnv()
q := proto.Query{
Sql: "select * from user where keyrange('', '\x20')",
TabletType: topo.TYPE_MASTER,
}
result, err := routerStream(router, &q)
if err != nil {
t.Error(err)
}
wantResult := singleRowResult
if !reflect.DeepEqual(result, wantResult) {
t.Errorf("result: %+v, want %+v", result, wantResult)
}
q.Sql = "select * from user where keyrange('\x40', '\x60')"
result, err = routerStream(router, &q)
if err != nil {
t.Error(err)
}
wantResult = singleRowResult
if !reflect.DeepEqual(result, wantResult) {
t.Errorf("result: %+v, want %+v", result, wantResult)
}
}
func TestStreamSelectIN(t *testing.T) {
router, _, _, sbclookup := createRouterEnv()
q := proto.Query{
Sql: "select * from user where id in (1)",
TabletType: topo.TYPE_MASTER,
}
result, err := routerStream(router, &q)
if err != nil {
t.Error(err)
}
wantResult := singleRowResult
if !reflect.DeepEqual(result, wantResult) {
t.Errorf("result: %+v, want %+v", result, wantResult)
}
q.Sql = "select * from user where id in (1, 3)"
result, err = routerStream(router, &q)
if err != nil {
t.Error(err)
}
wantResult = &mproto.QueryResult{
Fields: singleRowResult.Fields,
Rows: [][]sqltypes.Value{
singleRowResult.Rows[0],
singleRowResult.Rows[0],
},
RowsAffected: 2,
}
if !reflect.DeepEqual(result, wantResult) {
t.Errorf("result: %+v, want %+v", result, wantResult)
}
q.Sql = "select * from user where name = 'foo'"
result, err = routerStream(router, &q)
if err != nil {
t.Error(err)
}
wantResult = singleRowResult
if !reflect.DeepEqual(result, wantResult) {
t.Errorf("result: %+v, want %+v", result, wantResult)
}
wantQueries := []tproto.BoundQuery{{
Sql: "select user_id from name_user_map where name = :name",
BindVariables: map[string]interface{}{
"name": "foo",
},
}}
if !reflect.DeepEqual(sbclookup.Queries, wantQueries) {
t.Errorf("sbclookup.Queries: %+v, want %+v\n", sbclookup.Queries, wantQueries)
}
}
func TestVTGateExecute(t *testing.T) {
sandbox := createSandbox(KsTestUnsharded)
sbc := &sandboxConn{}
sandbox.MapTestConn("0", sbc)
q := proto.Query{
Sql: "select * from t1",
TabletType: topo.TYPE_MASTER,
}
qr := new(proto.QueryResult)
err := rpcVTGate.Execute(context.Background(), &q, qr)
if err != nil {
t.Errorf("want nil, got %v", err)
}
wantqr := new(proto.QueryResult)
wantqr.Result = singleRowResult
if !reflect.DeepEqual(wantqr, qr) {
t.Errorf("want \n%+v, got \n%+v", singleRowResult, qr)
}
if qr.Session != nil {
t.Errorf("want nil, got %+v\n", qr.Session)
}
q.Session = new(proto.Session)
rpcVTGate.Begin(context.Background(), q.Session)
if !q.Session.InTransaction {
t.Errorf("want true, got false")
}
rpcVTGate.Execute(context.Background(), &q, qr)
wantSession := &proto.Session{
InTransaction: true,
ShardSessions: []*proto.ShardSession{{
Keyspace: KsTestUnsharded,
Shard: "0",
TabletType: topo.TYPE_MASTER,
TransactionId: 1,
}},
}
if !reflect.DeepEqual(wantSession, q.Session) {
t.Errorf("want \n%+v, got \n%+v", wantSession, q.Session)
}
rpcVTGate.Commit(context.Background(), q.Session)
if sbc.CommitCount != 1 {
t.Errorf("want 1, got %d", sbc.CommitCount)
}
q.Session = new(proto.Session)
rpcVTGate.Begin(context.Background(), q.Session)
rpcVTGate.Execute(context.Background(), &q, qr)
rpcVTGate.Rollback(context.Background(), q.Session)
}
func TestSelectSingleShardKey(t *testing.T) {
schema, err := planbuilder.LoadSchemaJSON(locateFile("router_test.json"))
if err != nil {
t.Fatal(err)
}
s := createSandbox("TestRouter")
sbc1 := &sandboxConn{}
sbc2 := &sandboxConn{}
s.MapTestConn("-20", sbc1)
s.MapTestConn("40-60", sbc2)
serv := new(sandboxTopo)
scatterConn := NewScatterConn(serv, "", "aa", 1*time.Second, 10, 1*time.Millisecond)
router := NewRouter(serv, "aa", schema, "", scatterConn)
q := proto.Query{
Sql: "select * from user where id = 1",
TabletType: topo.TYPE_MASTER,
}
_, err = router.Execute(&context.DummyContext{}, &q)
if err != nil {
t.Errorf("want nil, got %v", err)
}
if sbc1.ExecCount != 1 {
t.Errorf("want 1, got %v\n", sbc1.ExecCount)
}
if sbc2.ExecCount != 0 {
t.Errorf("want 0, got %v\n", sbc2.ExecCount)
}
q.Sql = "select * from user where id = 3"
_, err = router.Execute(&context.DummyContext{}, &q)
if err != nil {
t.Errorf("want nil, got %v", err)
}
if sbc1.ExecCount != 1 {
t.Errorf("want 1, got %v\n", sbc1.ExecCount)
}
if sbc2.ExecCount != 1 {
t.Errorf("want 1, got %v\n", sbc2.ExecCount)
}
}
// Execute routes a non-streaming query.
func (rtr *Router) Execute(ctx context.Context, query *proto.Query) (*mproto.QueryResult, error) {
if query.BindVariables == nil {
query.BindVariables = make(map[string]interface{})
}
vcursor := newRequestContext(ctx, query, rtr)
plan := rtr.planner.GetPlan(string(query.Sql))
switch plan.ID {
case planbuilder.UpdateEqual:
return rtr.execUpdateEqual(vcursor, plan)
case planbuilder.DeleteEqual:
return rtr.execDeleteEqual(vcursor, plan)
case planbuilder.InsertSharded:
return rtr.execInsertSharded(vcursor, plan)
}
var err error
var params *scatterParams
switch plan.ID {
case planbuilder.SelectUnsharded, planbuilder.UpdateUnsharded,
planbuilder.DeleteUnsharded, planbuilder.InsertUnsharded:
params, err = rtr.paramsUnsharded(vcursor, plan)
case planbuilder.SelectEqual:
params, err = rtr.paramsSelectEqual(vcursor, plan)
case planbuilder.SelectIN:
params, err = rtr.paramsSelectIN(vcursor, plan)
case planbuilder.SelectKeyrange:
params, err = rtr.paramsSelectKeyrange(vcursor, plan)
case planbuilder.SelectScatter:
params, err = rtr.paramsSelectScatter(vcursor, plan)
default:
return nil, fmt.Errorf("cannot route query: %s: %s", query.Sql, plan.Reason)
}
if err != nil {
return nil, err
}
return rtr.scatterConn.ExecuteMulti(
ctx,
params.query,
params.ks,
params.shardVars,
query.TabletType,
NewSafeSession(vcursor.query.Session),
query.NotInTransaction,
)
}
// StreamExecute executes a streaming query.
func (rtr *Router) StreamExecute(ctx context.Context, query *proto.Query, sendReply func(*mproto.QueryResult) error) error {
if query.BindVariables == nil {
query.BindVariables = make(map[string]interface{})
}
vcursor := newRequestContext(ctx, query, rtr)
plan := rtr.planner.GetPlan(string(query.Sql))
var err error
var params *scatterParams
switch plan.ID {
case planbuilder.SelectUnsharded:
params, err = rtr.paramsUnsharded(vcursor, plan)
case planbuilder.SelectEqual:
params, err = rtr.paramsSelectEqual(vcursor, plan)
case planbuilder.SelectIN:
params, err = rtr.paramsSelectIN(vcursor, plan)
case planbuilder.SelectKeyrange:
params, err = rtr.paramsSelectKeyrange(vcursor, plan)
case planbuilder.SelectScatter:
params, err = rtr.paramsSelectScatter(vcursor, plan)
default:
return fmt.Errorf("query %q cannot be used for streaming", query.Sql)
}
if err != nil {
return err
}
return rtr.scatterConn.StreamExecuteMulti(
ctx,
params.query,
params.ks,
params.shardVars,
query.TabletType,
NewSafeSession(vcursor.query.Session),
sendReply,
query.NotInTransaction,
)
}