Golang proto.Query類代碼示例

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