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, ) }