Ejemplo n.º 1
0
// ExecuteKeyRanges executes a non-streaming query based on the specified keyranges.
func (vtg *VTGate) ExecuteKeyRanges(ctx context.Context, sql string, bindVariables map[string]interface{}, keyspace string, keyRanges []*pb.KeyRange, tabletType pb.TabletType, session *proto.Session, notInTransaction bool, reply *proto.QueryResult) error {
	startTime := time.Now()
	statsKey := []string{"ExecuteKeyRanges", keyspace, strings.ToLower(tabletType.String())}
	defer vtg.timings.Record(statsKey, startTime)

	x := vtg.inFlight.Add(1)
	defer vtg.inFlight.Add(-1)
	if 0 < vtg.maxInFlight && vtg.maxInFlight < x {
		return errTooManyInFlight
	}

	sql = sqlannotation.AddFilteredReplicationUnfriendlyIfDML(sql)

	qr, err := vtg.resolver.ExecuteKeyRanges(ctx, sql, bindVariables, keyspace, keyRanges, tabletType, session, notInTransaction)
	if err == nil {
		reply.Result = qr
		vtg.rowsReturned.Add(statsKey, int64(len(qr.Rows)))
	} else {
		query := map[string]interface{}{
			"Sql":              sql,
			"BindVariables":    bindVariables,
			"Keyspace":         keyspace,
			"KeyRanges":        keyRanges,
			"TabletType":       strings.ToLower(tabletType.String()),
			"Session":          session,
			"NotInTransaction": notInTransaction,
		}
		reply.Error = handleExecuteError(err, statsKey, query, vtg.logExecuteKeyRanges).Error()
		reply.Err = vterrors.RPCErrFromVtError(err)
	}
	reply.Session = session
	return nil
}
Ejemplo n.º 2
0
// StreamExecuteKeyRanges executes a streaming query on the specified KeyRanges.
// The KeyRanges are resolved to shards using the serving graph.
// This function currently temporarily enforces the restriction of executing on
// one shard since it cannot merge-sort the results to guarantee ordering of
// response which is needed for checkpointing.
// The api supports supplying multiple keyranges to make it future proof.
func (vtg *VTGate) StreamExecuteKeyRanges(context context.Context, query *proto.KeyRangeQuery, sendReply func(*proto.QueryResult) error) (err error) {
	defer handlePanic(&err)

	startTime := time.Now()
	statsKey := []string{"StreamExecuteKeyRanges", query.Keyspace, string(query.TabletType)}
	defer vtg.timings.Record(statsKey, startTime)

	x := vtg.inFlight.Add(1)
	defer vtg.inFlight.Add(-1)
	if 0 < vtg.maxInFlight && vtg.maxInFlight < x {
		return ErrTooManyInFlight
	}

	err = vtg.resolver.StreamExecuteKeyRanges(
		context,
		query,
		func(mreply *mproto.QueryResult) error {
			reply := new(proto.QueryResult)
			reply.Result = mreply
			// Note we don't populate reply.Session here,
			// as it may change incrementaly as responses are sent.
			return sendReply(reply)
		})

	if err != nil {
		normalErrors.Add(statsKey, 1)
		vtg.logStreamExecuteKeyRanges.Errorf("%v, query: %+v", err, query)
	}
	// Now we can send the final Sessoin info.
	if query.Session != nil {
		sendReply(&proto.QueryResult{Session: query.Session})
	}
	return err
}
Ejemplo n.º 3
0
func TestVTGateStreamExecuteShards(t *testing.T) {
	s := createSandbox("TestVTGateStreamExecuteShards")
	sbc := &sandboxConn{}
	s.MapTestConn("0", sbc)
	// Test for successful execution
	var qrs []*proto.QueryResult
	err := rpcVTGate.StreamExecuteShards(context.Background(),
		"query",
		nil,
		"TestVTGateStreamExecuteShards",
		[]string{"0"},
		topodatapb.TabletType_MASTER,
		func(r *proto.QueryResult) error {
			qrs = append(qrs, r)
			return nil
		})
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}
	row := new(proto.QueryResult)
	row.Result = singleRowResult
	want := []*proto.QueryResult{row}
	if !reflect.DeepEqual(want, qrs) {
		t.Errorf("want \n%+v, got \n%+v", want, qrs)
	}
}
Ejemplo n.º 4
0
// ExecuteEntityIds excutes a non-streaming query based on given KeyspaceId map.
func (vtg *VTGate) ExecuteEntityIds(context context.Context, query *proto.EntityIdsQuery, reply *proto.QueryResult) (err error) {
	defer handlePanic(&err)

	startTime := time.Now()
	statsKey := []string{"ExecuteEntityIds", query.Keyspace, string(query.TabletType)}
	defer vtg.timings.Record(statsKey, startTime)

	x := vtg.inFlight.Add(1)
	defer vtg.inFlight.Add(-1)
	if 0 < vtg.maxInFlight && vtg.maxInFlight < x {
		return ErrTooManyInFlight
	}

	qr, err := vtg.resolver.ExecuteEntityIds(context, query)
	if err == nil {
		reply.Result = qr
	} else {
		reply.Error = err.Error()
		if strings.Contains(reply.Error, errDupKey) {
			infoErrors.Add("DupKey", 1)
		} else {
			normalErrors.Add(statsKey, 1)
			vtg.logExecuteEntityIds.Errorf("%v, query: %+v", err, query)
		}
	}
	reply.Session = query.Session
	return nil
}
Ejemplo n.º 5
0
// ExecuteShard executes a non-streaming query on the specified shards.
func (vtg *VTGate) ExecuteShard(ctx context.Context, query *proto.QueryShard, reply *proto.QueryResult) error {
	startTime := time.Now()
	statsKey := []string{"ExecuteShard", query.Keyspace, string(query.TabletType)}
	defer vtg.timings.Record(statsKey, startTime)

	x := vtg.inFlight.Add(1)
	defer vtg.inFlight.Add(-1)
	if 0 < vtg.maxInFlight && vtg.maxInFlight < x {
		return errTooManyInFlight
	}

	qr, err := vtg.resolver.Execute(
		ctx,
		query.Sql,
		query.BindVariables,
		query.Keyspace,
		query.TabletType,
		query.Session,
		func(keyspace string) (string, []string, error) {
			return query.Keyspace, query.Shards, nil
		},
		query.NotInTransaction,
	)
	if err == nil {
		reply.Result = qr
		vtg.rowsReturned.Add(statsKey, int64(len(qr.Rows)))
	} else {
		reply.Error = handleExecuteError(err, statsKey, query, vtg.logExecuteShard)
	}
	reply.Session = query.Session
	return nil
}
Ejemplo n.º 6
0
// ExecuteShard executes a non-streaming query on the specified shards.
func (vtg *VTGate) ExecuteShard(context context.Context, query *proto.QueryShard, reply *proto.QueryResult) error {
	startTime := time.Now()
	statsKey := []string{"ExecuteShard", query.Keyspace, string(query.TabletType)}
	defer vtg.timings.Record(statsKey, startTime)

	qr, err := vtg.resolver.Execute(
		context,
		query.Sql,
		query.BindVariables,
		query.Keyspace,
		query.TabletType,
		query.Session,
		func(keyspace string) (string, []string, error) {
			return query.Keyspace, query.Shards, nil
		},
	)
	if err == nil {
		reply.Result = qr
	} else {
		reply.Error = err.Error()
		if strings.Contains(reply.Error, errDupKey) {
			vtg.infoErrors.Add("DupKey", 1)
		} else {
			vtg.errors.Add(statsKey, 1)
			vtg.logExecuteShard.Errorf("%v, query: %+v", err, query)
		}
	}
	reply.Session = query.Session
	return nil
}
Ejemplo n.º 7
0
// StreamExecuteKeyRanges executes a streaming query on the specified KeyRanges.
// The KeyRanges are resolved to shards using the serving graph.
// This function currently temporarily enforces the restriction of executing on
// one shard since it cannot merge-sort the results to guarantee ordering of
// response which is needed for checkpointing.
// The api supports supplying multiple keyranges to make it future proof.
func (vtg *VTGate) StreamExecuteKeyRanges(ctx context.Context, query *proto.KeyRangeQuery, sendReply func(*proto.QueryResult) error) error {
	startTime := time.Now()
	statsKey := []string{"StreamExecuteKeyRanges", query.Keyspace, string(query.TabletType)}
	defer vtg.timings.Record(statsKey, startTime)

	x := vtg.inFlight.Add(1)
	defer vtg.inFlight.Add(-1)
	if 0 < vtg.maxInFlight && vtg.maxInFlight < x {
		return errTooManyInFlight
	}

	var rowCount int64
	err := vtg.resolver.StreamExecuteKeyRanges(
		ctx,
		query,
		func(mreply *mproto.QueryResult) error {
			reply := new(proto.QueryResult)
			reply.Result = mreply
			rowCount += int64(len(mreply.Rows))
			vtg.rowsReturned.Add(statsKey, int64(len(mreply.Rows)))
			// Note we don't populate reply.Session here,
			// as it may change incrementaly as responses are sent.
			return sendReply(reply)
		})

	if err != nil {
		normalErrors.Add(statsKey, 1)
		logError(err, query, vtg.logStreamExecuteKeyRanges)
	}
	// Now we can send the final Sessoin info.
	if query.Session != nil {
		sendReply(&proto.QueryResult{Session: query.Session})
	}
	return formatError(err)
}
Ejemplo n.º 8
0
// StreamExecuteShard executes a streaming query on the specified shards.
func (vtg *VTGate) StreamExecuteShard(context context.Context, query *proto.QueryShard, sendReply func(*proto.QueryResult) error) error {
	startTime := time.Now()
	statsKey := []string{"StreamExecuteShard", query.Keyspace, string(query.TabletType)}
	defer vtg.timings.Record(statsKey, startTime)

	err := vtg.resolver.StreamExecute(
		context,
		query.Sql,
		query.BindVariables,
		query.Keyspace,
		query.TabletType,
		query.Session,
		func(keyspace string) (string, []string, error) {
			return query.Keyspace, query.Shards, nil
		},
		func(mreply *mproto.QueryResult) error {
			reply := new(proto.QueryResult)
			reply.Result = mreply
			// Note we don't populate reply.Session here,
			// as it may change incrementaly as responses are sent.
			return sendReply(reply)
		})

	if err != nil {
		vtg.errors.Add(statsKey, 1)
		vtg.logStreamExecuteShard.Errorf("%v, query: %+v", err, query)
	}
	// Now we can send the final Sessoin info.
	if query.Session != nil {
		sendReply(&proto.QueryResult{Session: query.Session})
	}
	return err
}
Ejemplo n.º 9
0
// ExecuteEntityIds excutes a non-streaming query based on given KeyspaceId map.
func (vtg *VTGate) ExecuteEntityIds(ctx context.Context, sql string, bindVariables map[string]interface{}, keyspace string, entityColumnName string, entityKeyspaceIDs []*pbg.ExecuteEntityIdsRequest_EntityId, tabletType pb.TabletType, session *proto.Session, notInTransaction bool, reply *proto.QueryResult) error {
	startTime := time.Now()
	statsKey := []string{"ExecuteEntityIds", keyspace, strings.ToLower(tabletType.String())}
	defer vtg.timings.Record(statsKey, startTime)

	x := vtg.inFlight.Add(1)
	defer vtg.inFlight.Add(-1)
	if 0 < vtg.maxInFlight && vtg.maxInFlight < x {
		return errTooManyInFlight
	}

	qr, err := vtg.resolver.ExecuteEntityIds(ctx, sql, bindVariables, keyspace, entityColumnName, entityKeyspaceIDs, tabletType, session, notInTransaction)
	if err == nil {
		reply.Result = qr
		vtg.rowsReturned.Add(statsKey, int64(len(qr.Rows)))
	} else {
		query := map[string]interface{}{
			"Sql":               sql,
			"BindVariables":     bindVariables,
			"Keyspace":          keyspace,
			"EntityColumnName":  entityColumnName,
			"EntityKeyspaceIDs": entityKeyspaceIDs,
			"TabletType":        strings.ToLower(tabletType.String()),
			"Session":           session,
			"NotInTransaction":  notInTransaction,
		}
		reply.Error = handleExecuteError(err, statsKey, query, vtg.logExecuteEntityIds).Error()
		reply.Err = rpcErrFromVtGateError(err)
	}
	reply.Session = session
	return nil
}
Ejemplo n.º 10
0
func TestVTGateExecuteShard(t *testing.T) {
	sandbox := createSandbox("TestVTGateExecuteShard")
	sbc := &sandboxConn{}
	sandbox.MapTestConn("0", sbc)
	q := proto.QueryShard{
		Sql:        "query",
		Keyspace:   "TestVTGateExecuteShard",
		Shards:     []string{"0"},
		TabletType: topo.TYPE_REPLICA,
	}
	qr := new(proto.QueryResult)
	err := rpcVTGate.ExecuteShard(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.ExecuteShard(context.Background(), &q, qr)
	wantSession := &proto.Session{
		InTransaction: true,
		ShardSessions: []*proto.ShardSession{{
			Keyspace:      "TestVTGateExecuteShard",
			Shard:         "0",
			TabletType:    topo.TYPE_REPLICA,
			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 commitCount := sbc.CommitCount.Get(); commitCount != 1 {
		t.Errorf("want 1, got %d", commitCount)
	}

	q.Session = new(proto.Session)
	rpcVTGate.Begin(context.Background(), q.Session)
	rpcVTGate.ExecuteShard(context.Background(), &q, qr)
	rpcVTGate.Rollback(context.Background(), q.Session)
	/*
		// Flaky: This test should be run manually.
		runtime.Gosched()
		if sbc.RollbackCount != 1 {
			t.Errorf("want 1, got %d", sbc.RollbackCount)
		}
	*/
}
Ejemplo n.º 11
0
// Execute executes a non-streaming query by routing based on the values in the query.
func (vtg *VTGate) Execute(ctx context.Context, sql string, bindVariables map[string]interface{}, tabletType pb.TabletType, session *proto.Session, notInTransaction bool, reply *proto.QueryResult) error {
	startTime := time.Now()
	statsKey := []string{"Execute", "Any", strings.ToLower(tabletType.String())}
	defer vtg.timings.Record(statsKey, startTime)

	x := vtg.inFlight.Add(1)
	defer vtg.inFlight.Add(-1)
	if 0 < vtg.maxInFlight && vtg.maxInFlight < x {
		return errTooManyInFlight
	}

	qr, err := vtg.router.Execute(ctx, sql, bindVariables, tabletType, session, notInTransaction)
	if err == nil {
		reply.Result = qr
		vtg.rowsReturned.Add(statsKey, int64(len(qr.Rows)))
	} else {
		query := map[string]interface{}{
			"Sql":              sql,
			"BindVariables":    bindVariables,
			"TabletType":       strings.ToLower(tabletType.String()),
			"Session":          session,
			"NotInTransaction": notInTransaction,
		}
		reply.Error = handleExecuteError(err, statsKey, query, vtg.logExecute)
	}
	reply.Session = session
	return nil
}
Ejemplo n.º 12
0
func TestVTGateStreamExecuteShard(t *testing.T) {
	s := createSandbox("TestVTGateStreamExecuteShard")
	sbc := &sandboxConn{}
	s.MapTestConn("0", sbc)
	q := proto.QueryShard{
		Sql:        "query",
		Keyspace:   "TestVTGateStreamExecuteShard",
		Shards:     []string{"0"},
		TabletType: topo.TYPE_MASTER,
	}
	// Test for successful execution
	var qrs []*proto.QueryResult
	err := RpcVTGate.StreamExecuteShard(&context.DummyContext{}, &q, func(r *proto.QueryResult) error {
		qrs = append(qrs, r)
		return nil
	})
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}
	row := new(proto.QueryResult)
	row.Result = singleRowResult
	want := []*proto.QueryResult{row}
	if !reflect.DeepEqual(want, qrs) {
		t.Errorf("want \n%+v, got \n%+v", want, qrs)
	}

	q.Session = new(proto.Session)
	qrs = nil
	RpcVTGate.Begin(&context.DummyContext{}, q.Session)
	err = RpcVTGate.StreamExecuteShard(&context.DummyContext{}, &q, func(r *proto.QueryResult) error {
		qrs = append(qrs, r)
		return nil
	})
	want = []*proto.QueryResult{
		row,
		&proto.QueryResult{
			Session: &proto.Session{
				InTransaction: true,
				ShardSessions: []*proto.ShardSession{{
					Keyspace:      "TestVTGateStreamExecuteShard",
					Shard:         "0",
					TransactionId: 1,
					TabletType:    topo.TYPE_MASTER,
				}},
			},
		},
	}
	if !reflect.DeepEqual(want, qrs) {
		t.Errorf("want \n%+v, got \n%+v", want, qrs)
	}

}
Ejemplo n.º 13
0
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)
}
Ejemplo n.º 14
0
func (c *echoClient) Execute(ctx context.Context, sql string, bindVariables map[string]interface{}, tabletType pb.TabletType, session *proto.Session, notInTransaction bool, reply *proto.QueryResult) error {
	if strings.HasPrefix(sql, EchoPrefix) {
		reply.Result = echoQueryResult(map[string]interface{}{
			"callerId":         callerid.EffectiveCallerIDFromContext(ctx),
			"query":            sql,
			"bindVars":         bindVariables,
			"tabletType":       tabletType,
			"session":          session,
			"notInTransaction": notInTransaction,
		})
		reply.Session = session
		return nil
	}
	return c.fallbackClient.Execute(ctx, sql, bindVariables, tabletType, session, notInTransaction, reply)
}
Ejemplo n.º 15
0
func TestVTGateStreamExecuteKeyRanges(t *testing.T) {
	s := createSandbox("TestVTGateStreamExecuteKeyRanges")
	sbc := &sandboxConn{}
	s.MapTestConn("-20", sbc)
	sbc1 := &sandboxConn{}
	s.MapTestConn("20-40", sbc1)
	kr, err := key.ParseKeyRangeParts("", "20")
	// Test for successful execution
	var qrs []*proto.QueryResult
	err = rpcVTGate.StreamExecuteKeyRanges(context.Background(),
		"query",
		nil,
		"TestVTGateStreamExecuteKeyRanges",
		[]key.KeyRange{kr},
		pb.TabletType_MASTER,
		func(r *proto.QueryResult) error {
			qrs = append(qrs, r)
			return nil
		})
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}
	row := new(proto.QueryResult)
	row.Result = singleRowResult
	want := []*proto.QueryResult{row}
	if !reflect.DeepEqual(want, qrs) {
		t.Errorf("want \n%+v, got \n%+v", want, qrs)
	}

	// Test for successful execution - multiple shards
	kr, err = key.ParseKeyRangeParts("10", "40")
	err = rpcVTGate.StreamExecuteKeyRanges(context.Background(),
		"query",
		nil,
		"TestVTGateStreamExecuteKeyRanges",
		[]key.KeyRange{kr},
		pb.TabletType_MASTER,
		func(r *proto.QueryResult) error {
			qrs = append(qrs, r)
			return nil
		})
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}
}
Ejemplo n.º 16
0
func (c *echoClient) ExecuteEntityIds(ctx context.Context, sql string, bindVariables map[string]interface{}, keyspace string, entityColumnName string, entityKeyspaceIDs []*pbg.ExecuteEntityIdsRequest_EntityId, tabletType pb.TabletType, session *pbg.Session, notInTransaction bool, reply *proto.QueryResult) error {
	if strings.HasPrefix(sql, EchoPrefix) {
		reply.Result = echoQueryResult(map[string]interface{}{
			"callerId":         callerid.EffectiveCallerIDFromContext(ctx),
			"query":            sql,
			"bindVars":         bindVariables,
			"keyspace":         keyspace,
			"entityColumnName": entityColumnName,
			"entityIds":        entityKeyspaceIDs,
			"tabletType":       tabletType,
			"session":          session,
			"notInTransaction": notInTransaction,
		})
		reply.Session = session
		return nil
	}
	return c.fallbackClient.ExecuteEntityIds(ctx, sql, bindVariables, keyspace, entityColumnName, entityKeyspaceIDs, tabletType, session, notInTransaction, reply)
}
Ejemplo n.º 17
0
// StreamExecuteShards executes a streaming query on the specified shards.
func (vtg *VTGate) StreamExecuteShards(ctx context.Context, sql string, bindVariables map[string]interface{}, keyspace string, shards []string, tabletType pb.TabletType, sendReply func(*proto.QueryResult) error) error {
	startTime := time.Now()
	statsKey := []string{"StreamExecuteShards", keyspace, strings.ToLower(tabletType.String())}
	defer vtg.timings.Record(statsKey, startTime)

	x := vtg.inFlight.Add(1)
	defer vtg.inFlight.Add(-1)
	if 0 < vtg.maxInFlight && vtg.maxInFlight < x {
		return errTooManyInFlight
	}

	var rowCount int64
	err := vtg.resolver.StreamExecute(
		ctx,
		sql,
		bindVariables,
		keyspace,
		tabletType,
		func(keyspace string) (string, []string, error) {
			return keyspace, shards, nil
		},
		func(mreply *mproto.QueryResult) error {
			reply := new(proto.QueryResult)
			reply.Result = mreply
			rowCount += int64(len(mreply.Rows))
			vtg.rowsReturned.Add(statsKey, int64(len(mreply.Rows)))
			// Note we don't populate reply.Session here,
			// as it may change incrementaly as responses are sent.
			return sendReply(reply)
		})

	if err != nil {
		normalErrors.Add(statsKey, 1)
		query := map[string]interface{}{
			"Sql":           sql,
			"BindVariables": bindVariables,
			"Keyspace":      keyspace,
			"Shards":        shards,
			"TabletType":    strings.ToLower(tabletType.String()),
		}
		logError(err, query, vtg.logStreamExecuteShards)
	}
	return formatError(err)
}
Ejemplo n.º 18
0
// ExecuteEntityIds excutes a non-streaming query based on given KeyspaceId map.
func (vtg *VTGate) ExecuteEntityIds(context context.Context, query *proto.EntityIdsQuery, reply *proto.QueryResult) error {
	startTime := time.Now()
	statsKey := []string{"ExecuteEntityIds", query.Keyspace, string(query.TabletType)}
	defer vtg.timings.Record(statsKey, startTime)

	qr, err := vtg.resolver.ExecuteEntityIds(context, query)
	if err == nil {
		reply.Result = qr
	} else {
		reply.Error = err.Error()
		if strings.Contains(reply.Error, errDupKey) {
			vtg.infoErrors.Add("DupKey", 1)
		} else {
			vtg.errors.Add(statsKey, 1)
			vtg.logExecuteEntityIds.Errorf("%v, query: %+v", err, query)
		}
	}
	reply.Session = query.Session
	return nil
}
Ejemplo n.º 19
0
// ExecuteEntityIds excutes a non-streaming query based on given KeyspaceId map.
func (vtg *VTGate) ExecuteEntityIds(ctx context.Context, query *proto.EntityIdsQuery, reply *proto.QueryResult) error {
	startTime := time.Now()
	statsKey := []string{"ExecuteEntityIds", query.Keyspace, string(query.TabletType)}
	defer vtg.timings.Record(statsKey, startTime)

	x := vtg.inFlight.Add(1)
	defer vtg.inFlight.Add(-1)
	if 0 < vtg.maxInFlight && vtg.maxInFlight < x {
		return errTooManyInFlight
	}

	qr, err := vtg.resolver.ExecuteEntityIds(ctx, query)
	if err == nil {
		reply.Result = qr
		vtg.rowsReturned.Add(statsKey, int64(len(qr.Rows)))
	} else {
		reply.Error = handleExecuteError(err, statsKey, query, vtg.logExecuteEntityIds)
	}
	reply.Session = query.Session
	return nil
}
Ejemplo n.º 20
0
func TestVTGateStreamExecute(t *testing.T) {
	sandbox := createSandbox(KsTestUnsharded)
	sbc := &sandboxConn{}
	sandbox.MapTestConn("0", sbc)
	var qrs []*proto.QueryResult
	err := rpcVTGate.StreamExecute(context.Background(),
		"select * from t1",
		nil,
		topodatapb.TabletType_MASTER,
		func(r *proto.QueryResult) error {
			qrs = append(qrs, r)
			return nil
		})
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}
	row := new(proto.QueryResult)
	row.Result = singleRowResult
	want := []*proto.QueryResult{row}
	if !reflect.DeepEqual(want, qrs) {
		t.Errorf("want \n%+v, got \n%+v", want, qrs)
	}
}
Ejemplo n.º 21
0
func TestVTGateExecuteKeyspaceIds(t *testing.T) {
	s := createSandbox("TestVTGateExecuteKeyspaceIds")
	sbc1 := &sandboxConn{}
	sbc2 := &sandboxConn{}
	s.MapTestConn("-20", sbc1)
	s.MapTestConn("20-40", sbc2)
	kid10, err := key.HexKeyspaceId("10").Unhex()
	if err != nil {
		t.Errorf("want nil, got %+v", err)
	}
	q := proto.KeyspaceIdQuery{
		Sql:         "query",
		Keyspace:    "TestVTGateExecuteKeyspaceIds",
		KeyspaceIds: []key.KeyspaceId{kid10},
		TabletType:  topo.TYPE_MASTER,
	}
	// Test for successful execution
	qr := new(proto.QueryResult)
	err = RpcVTGate.ExecuteKeyspaceIds(&context.DummyContext{}, &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)
	}
	if sbc1.ExecCount != 1 {
		t.Errorf("want 1, got %v\n", sbc1.ExecCount)
	}
	// Test for successful execution in transaction
	q.Session = new(proto.Session)
	RpcVTGate.Begin(&context.DummyContext{}, q.Session)
	if !q.Session.InTransaction {
		t.Errorf("want true, got false")
	}
	RpcVTGate.ExecuteKeyspaceIds(&context.DummyContext{}, &q, qr)
	wantSession := &proto.Session{
		InTransaction: true,
		ShardSessions: []*proto.ShardSession{{
			Keyspace:      "TestVTGateExecuteKeyspaceIds",
			Shard:         "-20",
			TransactionId: 1,
			TabletType:    topo.TYPE_MASTER,
		}},
	}
	if !reflect.DeepEqual(wantSession, q.Session) {
		t.Errorf("want \n%+v, got \n%+v", wantSession, q.Session)
	}
	RpcVTGate.Commit(&context.DummyContext{}, q.Session)
	if sbc1.CommitCount.Get() != 1 {
		t.Errorf("want 1, got %d", sbc1.CommitCount.Get())
	}
	// Test for multiple shards
	kid30, err := key.HexKeyspaceId("30").Unhex()
	if err != nil {
		t.Errorf("want nil, got %+v", err)
	}
	q.KeyspaceIds = []key.KeyspaceId{kid10, kid30}
	RpcVTGate.ExecuteKeyspaceIds(&context.DummyContext{}, &q, qr)
	if qr.Result.RowsAffected != 2 {
		t.Errorf("want 2, got %v", qr.Result.RowsAffected)
	}
}
Ejemplo n.º 22
0
func TestVTGateExecuteKeyRanges(t *testing.T) {
	s := createSandbox("TestVTGateExecuteKeyRanges")
	sbc1 := &sandboxConn{}
	sbc2 := &sandboxConn{}
	s.MapTestConn("-20", sbc1)
	s.MapTestConn("20-40", sbc2)
	// Test for successful execution
	qr := new(proto.QueryResult)
	err := rpcVTGate.ExecuteKeyRanges(context.Background(),
		"query",
		nil,
		"TestVTGateExecuteKeyRanges",
		[]*pb.KeyRange{&pb.KeyRange{End: []byte{0x20}}},
		pb.TabletType_MASTER,
		nil,
		false,
		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)
	}
	if execCount := sbc1.ExecCount.Get(); execCount != 1 {
		t.Errorf("want 1, got %v\n", execCount)
	}
	// Test for successful execution in transaction
	session := new(proto.Session)
	rpcVTGate.Begin(context.Background(), session)
	if !session.InTransaction {
		t.Errorf("want true, got false")
	}
	err = rpcVTGate.ExecuteKeyRanges(context.Background(),
		"query",
		nil,
		"TestVTGateExecuteKeyRanges",
		[]*pb.KeyRange{&pb.KeyRange{End: []byte{0x20}}},
		pb.TabletType_MASTER,
		session,
		false,
		qr)
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}
	wantSession := &proto.Session{
		InTransaction: true,
		ShardSessions: []*proto.ShardSession{{
			Keyspace:      "TestVTGateExecuteKeyRanges",
			Shard:         "-20",
			TransactionId: 1,
			TabletType:    topo.TYPE_MASTER,
		}},
	}
	if !reflect.DeepEqual(wantSession, session) {
		t.Errorf("want \n%+v, got \n%+v", wantSession, session)
	}
	rpcVTGate.Commit(context.Background(), session)
	if commitCount := sbc1.CommitCount.Get(); commitCount != 1 {
		t.Errorf("want 1, got %v", commitCount)
	}
	// Test for multiple shards
	rpcVTGate.ExecuteKeyRanges(context.Background(), "query",
		nil,
		"TestVTGateExecuteKeyRanges",
		[]*pb.KeyRange{&pb.KeyRange{Start: []byte{0x10}, End: []byte{0x30}}},
		pb.TabletType_MASTER,
		nil,
		false, qr)
	if qr.Result.RowsAffected != 2 {
		t.Errorf("want 2, got %v", qr.Result.RowsAffected)
	}
}
Ejemplo n.º 23
0
func TestVTGateExecuteKeyRanges(t *testing.T) {
	s := createSandbox("TestVTGateExecuteKeyRanges")
	sbc1 := &sandboxConn{}
	sbc2 := &sandboxConn{}
	s.MapTestConn("-20", sbc1)
	s.MapTestConn("20-40", sbc2)
	kr, err := key.ParseKeyRangeParts("", "20")
	q := proto.KeyRangeQuery{
		Sql:        "query",
		Keyspace:   "TestVTGateExecuteKeyRanges",
		KeyRanges:  []key.KeyRange{kr},
		TabletType: topo.TYPE_MASTER,
	}
	// Test for successful execution
	qr := new(proto.QueryResult)
	err = rpcVTGate.ExecuteKeyRanges(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)
	}
	if sbc1.ExecCount != 1 {
		t.Errorf("want 1, got %v\n", sbc1.ExecCount)
	}
	// Test for successful execution in transaction
	q.Session = new(proto.Session)
	rpcVTGate.Begin(context.Background(), q.Session)
	if !q.Session.InTransaction {
		t.Errorf("want true, got false")
	}
	err = rpcVTGate.ExecuteKeyRanges(context.Background(), &q, qr)
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}
	wantSession := &proto.Session{
		InTransaction: true,
		ShardSessions: []*proto.ShardSession{{
			Keyspace:      "TestVTGateExecuteKeyRanges",
			Shard:         "-20",
			TransactionId: 1,
			TabletType:    topo.TYPE_MASTER,
		}},
	}
	if !reflect.DeepEqual(wantSession, q.Session) {
		t.Errorf("want \n%+v, got \n%+v", wantSession, q.Session)
	}
	rpcVTGate.Commit(context.Background(), q.Session)
	if sbc1.CommitCount.Get() != 1 {
		t.Errorf("want 1, got %v", sbc1.CommitCount.Get())
	}
	// Test for multiple shards
	kr, err = key.ParseKeyRangeParts("10", "30")
	q.KeyRanges = []key.KeyRange{kr}
	rpcVTGate.ExecuteKeyRanges(context.Background(), &q, qr)
	if qr.Result.RowsAffected != 2 {
		t.Errorf("want 2, got %v", qr.Result.RowsAffected)
	}
}
Ejemplo n.º 24
0
func TestVTGateExecuteEntityIds(t *testing.T) {
	s := createSandbox("TestVTGateExecuteEntityIds")
	sbc1 := &sandboxConn{}
	sbc2 := &sandboxConn{}
	s.MapTestConn("-20", sbc1)
	s.MapTestConn("20-40", sbc2)
	kid10, err := key.HexKeyspaceId("10").Unhex()
	if err != nil {
		t.Errorf("want nil, got %+v", err)
	}
	q := proto.EntityIdsQuery{
		Sql:              "query",
		Keyspace:         "TestVTGateExecuteEntityIds",
		EntityColumnName: "kid",
		EntityKeyspaceIDs: []proto.EntityId{
			proto.EntityId{
				ExternalID: "id1",
				KeyspaceID: kid10,
			},
		},
		TabletType: topo.TYPE_MASTER,
	}
	// Test for successful execution
	qr := new(proto.QueryResult)
	err = rpcVTGate.ExecuteEntityIds(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)
	}
	if sbc1.ExecCount != 1 {
		t.Errorf("want 1, got %v\n", sbc1.ExecCount)
	}
	// Test for successful execution in transaction
	q.Session = new(proto.Session)
	rpcVTGate.Begin(context.Background(), q.Session)
	if !q.Session.InTransaction {
		t.Errorf("want true, got false")
	}
	rpcVTGate.ExecuteEntityIds(context.Background(), &q, qr)
	wantSession := &proto.Session{
		InTransaction: true,
		ShardSessions: []*proto.ShardSession{{
			Keyspace:      "TestVTGateExecuteEntityIds",
			Shard:         "-20",
			TransactionId: 1,
			TabletType:    topo.TYPE_MASTER,
		}},
	}
	if !reflect.DeepEqual(wantSession, q.Session) {
		t.Errorf("want \n%+v, got \n%+v", wantSession, q.Session)
	}
	rpcVTGate.Commit(context.Background(), q.Session)
	if sbc1.CommitCount.Get() != 1 {
		t.Errorf("want 1, got %d", sbc1.CommitCount.Get())
	}
	// Test for multiple shards
	kid30, err := key.HexKeyspaceId("30").Unhex()
	if err != nil {
		t.Errorf("want nil, got %+v", err)
	}
	q.EntityKeyspaceIDs = append(q.EntityKeyspaceIDs, proto.EntityId{ExternalID: "id2", KeyspaceID: kid30})
	rpcVTGate.ExecuteEntityIds(context.Background(), &q, qr)
	if qr.Result.RowsAffected != 2 {
		t.Errorf("want 2, got %v", qr.Result.RowsAffected)
	}
}
Ejemplo n.º 25
0
func TestVTGateStreamExecuteKeyspaceIds(t *testing.T) {
	s := createSandbox("TestVTGateStreamExecuteKeyspaceIds")
	sbc := &sandboxConn{}
	s.MapTestConn("-20", sbc)
	sbc1 := &sandboxConn{}
	s.MapTestConn("20-40", sbc1)
	kid10, err := key.HexKeyspaceId("10").Unhex()
	if err != nil {
		t.Errorf("want nil, got %+v", err)
	}
	sq := proto.KeyspaceIdQuery{
		Sql:         "query",
		Keyspace:    "TestVTGateStreamExecuteKeyspaceIds",
		KeyspaceIds: []key.KeyspaceId{kid10},
		TabletType:  topo.TYPE_MASTER,
	}
	// Test for successful execution
	var qrs []*proto.QueryResult
	err = rpcVTGate.StreamExecuteKeyspaceIds(context.Background(), &sq, func(r *proto.QueryResult) error {
		qrs = append(qrs, r)
		return nil
	})
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}
	row := new(proto.QueryResult)
	row.Result = singleRowResult
	want := []*proto.QueryResult{row}
	if !reflect.DeepEqual(want, qrs) {
		t.Errorf("want \n%+v, got \n%+v", want, qrs)
	}

	// Test for successful execution in transaction
	sq.Session = new(proto.Session)
	qrs = nil
	rpcVTGate.Begin(context.Background(), sq.Session)
	err = rpcVTGate.StreamExecuteKeyspaceIds(context.Background(), &sq, func(r *proto.QueryResult) error {
		qrs = append(qrs, r)
		return nil
	})
	want = []*proto.QueryResult{
		row,
		&proto.QueryResult{
			Session: &proto.Session{
				InTransaction: true,
				ShardSessions: []*proto.ShardSession{{
					Keyspace:      "TestVTGateStreamExecuteKeyspaceIds",
					Shard:         "-20",
					TransactionId: 1,
					TabletType:    topo.TYPE_MASTER,
				}},
			},
		},
	}
	if !reflect.DeepEqual(want, qrs) {
		t.Errorf("want\n%#v\ngot\n%#v", want, qrs)
	}
	rpcVTGate.Commit(context.Background(), sq.Session)
	if sbc.CommitCount.Get() != 1 {
		t.Errorf("want 1, got %d", sbc.CommitCount.Get())
	}
	// Test for successful execution - multiple keyspaceids in single shard
	sq.Session = nil
	qrs = nil
	kid15, err := key.HexKeyspaceId("15").Unhex()
	if err != nil {
		t.Errorf("want nil, got %+v", err)
	}
	sq.KeyspaceIds = []key.KeyspaceId{kid10, kid15}
	err = rpcVTGate.StreamExecuteKeyspaceIds(context.Background(), &sq, func(r *proto.QueryResult) error {
		qrs = append(qrs, r)
		return nil
	})
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}
	row = new(proto.QueryResult)
	row.Result = singleRowResult
	want = []*proto.QueryResult{row}
	if !reflect.DeepEqual(want, qrs) {
		t.Errorf("want \n%+v, got \n%+v", want, qrs)
	}
	// Test for successful execution - multiple keyspaceids in multiple shards
	kid30, err := key.HexKeyspaceId("30").Unhex()
	if err != nil {
		t.Errorf("want nil, got %+v", err)
	}
	sq.KeyspaceIds = []key.KeyspaceId{kid10, kid30}
	err = rpcVTGate.StreamExecuteKeyspaceIds(context.Background(), &sq, func(r *proto.QueryResult) error {
		qrs = append(qrs, r)
		return nil
	})
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}
}
Ejemplo n.º 26
0
func TestVTGateStreamExecuteKeyspaceIds(t *testing.T) {
	s := createSandbox("TestVTGateStreamExecuteKeyspaceIds")
	sbc := &sandboxConn{}
	s.MapTestConn("-20", sbc)
	sbc1 := &sandboxConn{}
	s.MapTestConn("20-40", sbc1)
	// Test for successful execution
	var qrs []*proto.QueryResult
	err := rpcVTGate.StreamExecuteKeyspaceIds(context.Background(),
		"query",
		nil,
		"TestVTGateStreamExecuteKeyspaceIds",
		[][]byte{[]byte{0x10}},
		topodatapb.TabletType_MASTER,
		func(r *proto.QueryResult) error {
			qrs = append(qrs, r)
			return nil
		})
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}
	row := new(proto.QueryResult)
	row.Result = singleRowResult
	want := []*proto.QueryResult{row}
	if !reflect.DeepEqual(want, qrs) {
		t.Errorf("want \n%+v, got \n%+v", want, qrs)
	}

	// Test for successful execution - multiple keyspaceids in single shard
	qrs = nil
	err = rpcVTGate.StreamExecuteKeyspaceIds(context.Background(),
		"query",
		nil,
		"TestVTGateStreamExecuteKeyspaceIds",
		[][]byte{[]byte{0x10}, []byte{0x15}},
		topodatapb.TabletType_MASTER,
		func(r *proto.QueryResult) error {
			qrs = append(qrs, r)
			return nil
		})
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}
	row = new(proto.QueryResult)
	row.Result = singleRowResult
	want = []*proto.QueryResult{row}
	if !reflect.DeepEqual(want, qrs) {
		t.Errorf("want \n%+v, got \n%+v", want, qrs)
	}
	// Test for successful execution - multiple keyspaceids in multiple shards
	err = rpcVTGate.StreamExecuteKeyspaceIds(context.Background(),
		"query",
		nil,
		"TestVTGateStreamExecuteKeyspaceIds",
		[][]byte{[]byte{0x10}, []byte{0x30}},
		topodatapb.TabletType_MASTER,
		func(r *proto.QueryResult) error {
			qrs = append(qrs, r)
			return nil
		})
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}
}
Ejemplo n.º 27
0
func TestVTGateExecute(t *testing.T) {
	sandbox := createSandbox(KsTestUnsharded)
	sbc := &sandboxConn{}
	sandbox.MapTestConn("0", sbc)
	qr := new(proto.QueryResult)
	err := rpcVTGate.Execute(context.Background(),
		"select * from t1",
		nil,
		topodatapb.TabletType_MASTER,
		nil,
		false,
		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)
	}

	session := new(pbg.Session)
	rpcVTGate.Begin(context.Background(), session)
	if !session.InTransaction {
		t.Errorf("want true, got false")
	}
	rpcVTGate.Execute(context.Background(),
		"select * from t1",
		nil,
		topodatapb.TabletType_MASTER,
		session,
		false,
		qr)
	wantSession := &pbg.Session{
		InTransaction: true,
		ShardSessions: []*pbg.Session_ShardSession{{
			Target: &querypb.Target{
				Keyspace:   KsTestUnsharded,
				Shard:      "0",
				TabletType: topodatapb.TabletType_MASTER,
			},
			TransactionId: 1,
		}},
	}
	if !reflect.DeepEqual(wantSession, session) {
		t.Errorf("want \n%+v, got \n%+v", wantSession, session)
	}

	rpcVTGate.Commit(context.Background(), session)
	if commitCount := sbc.CommitCount.Get(); commitCount != 1 {
		t.Errorf("want 1, got %d", commitCount)
	}

	session = new(pbg.Session)
	rpcVTGate.Begin(context.Background(), session)
	rpcVTGate.Execute(context.Background(),
		"select * from t1",
		nil,
		topodatapb.TabletType_MASTER,
		session,
		false,
		qr)
	rpcVTGate.Rollback(context.Background(), session)
}
Ejemplo n.º 28
0
func TestVTGateExecuteEntityIds(t *testing.T) {
	s := createSandbox("TestVTGateExecuteEntityIds")
	sbc1 := &sandboxConn{}
	sbc2 := &sandboxConn{}
	s.MapTestConn("-20", sbc1)
	s.MapTestConn("20-40", sbc2)
	// Test for successful execution
	qr := new(proto.QueryResult)
	err := rpcVTGate.ExecuteEntityIds(context.Background(),
		"query",
		nil,
		"TestVTGateExecuteEntityIds",
		"kid",
		[]*pbg.ExecuteEntityIdsRequest_EntityId{
			&pbg.ExecuteEntityIdsRequest_EntityId{
				XidType:    sqltypes.VarBinary,
				XidValue:   []byte("id1"),
				KeyspaceId: []byte{0x10},
			},
		},
		topodatapb.TabletType_MASTER,
		nil,
		false,
		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)
	}
	if execCount := sbc1.ExecCount.Get(); execCount != 1 {
		t.Errorf("want 1, got %v\n", execCount)
	}
	// Test for successful execution in transaction
	session := new(pbg.Session)
	rpcVTGate.Begin(context.Background(), session)
	if !session.InTransaction {
		t.Errorf("want true, got false")
	}
	rpcVTGate.ExecuteEntityIds(context.Background(),
		"query",
		nil,
		"TestVTGateExecuteEntityIds",
		"kid",
		[]*pbg.ExecuteEntityIdsRequest_EntityId{
			&pbg.ExecuteEntityIdsRequest_EntityId{
				XidType:    sqltypes.VarBinary,
				XidValue:   []byte("id1"),
				KeyspaceId: []byte{0x10},
			},
		},
		topodatapb.TabletType_MASTER,
		session,
		false,
		qr)
	wantSession := &pbg.Session{
		InTransaction: true,
		ShardSessions: []*pbg.Session_ShardSession{{
			Target: &querypb.Target{
				Keyspace:   "TestVTGateExecuteEntityIds",
				Shard:      "-20",
				TabletType: topodatapb.TabletType_MASTER,
			},
			TransactionId: 1,
		}},
	}
	if !reflect.DeepEqual(wantSession, session) {
		t.Errorf("want \n%+v, got \n%+v", wantSession, session)
	}
	rpcVTGate.Commit(context.Background(), session)
	if commitCount := sbc1.CommitCount.Get(); commitCount != 1 {
		t.Errorf("want 1, got %d", commitCount)
	}

	// Test for multiple shards
	rpcVTGate.ExecuteEntityIds(context.Background(), "query",
		nil,
		"TestVTGateExecuteEntityIds",
		"kid",
		[]*pbg.ExecuteEntityIdsRequest_EntityId{
			&pbg.ExecuteEntityIdsRequest_EntityId{
				XidType:    sqltypes.VarBinary,
				XidValue:   []byte("id1"),
				KeyspaceId: []byte{0x10},
			},
			&pbg.ExecuteEntityIdsRequest_EntityId{
				XidType:    sqltypes.VarBinary,
				XidValue:   []byte("id2"),
				KeyspaceId: []byte{0x30},
			},
		},
		topodatapb.TabletType_MASTER,
		nil,
		false,
		qr)
	if qr.Result.RowsAffected != 2 {
		t.Errorf("want 2, got %v", qr.Result.RowsAffected)
	}
}
Ejemplo n.º 29
0
func TestVTGateStreamExecuteKeyRanges(t *testing.T) {
	s := createSandbox("TestVTGateStreamExecuteKeyRanges")
	sbc := &sandboxConn{}
	s.MapTestConn("-20", sbc)
	sbc1 := &sandboxConn{}
	s.MapTestConn("20-40", sbc1)
	kr, err := key.ParseKeyRangeParts("", "20")
	sq := proto.KeyRangeQuery{
		Sql:        "query",
		Keyspace:   "TestVTGateStreamExecuteKeyRanges",
		KeyRanges:  []key.KeyRange{kr},
		TabletType: topo.TYPE_MASTER,
	}
	// Test for successful execution
	var qrs []*proto.QueryResult
	err = rpcVTGate.StreamExecuteKeyRanges(context.Background(), &sq, func(r *proto.QueryResult) error {
		qrs = append(qrs, r)
		return nil
	})
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}
	row := new(proto.QueryResult)
	row.Result = singleRowResult
	want := []*proto.QueryResult{row}
	if !reflect.DeepEqual(want, qrs) {
		t.Errorf("want \n%+v, got \n%+v", want, qrs)
	}

	sq.Session = new(proto.Session)
	qrs = nil
	rpcVTGate.Begin(context.Background(), sq.Session)
	err = rpcVTGate.StreamExecuteKeyRanges(context.Background(), &sq, func(r *proto.QueryResult) error {
		qrs = append(qrs, r)
		return nil
	})
	want = []*proto.QueryResult{
		row,
		&proto.QueryResult{
			Session: &proto.Session{
				InTransaction: true,
				ShardSessions: []*proto.ShardSession{{
					Keyspace:      "TestVTGateStreamExecuteKeyRanges",
					Shard:         "-20",
					TransactionId: 1,
					TabletType:    topo.TYPE_MASTER,
				}},
			},
		},
	}
	if !reflect.DeepEqual(want, qrs) {
		t.Errorf("want \n%+v, got \n%+v", want, qrs)
	}

	// Test for successful execution - multiple shards
	kr, err = key.ParseKeyRangeParts("10", "40")
	sq.KeyRanges = []key.KeyRange{kr}
	err = rpcVTGate.StreamExecuteKeyRanges(context.Background(), &sq, func(r *proto.QueryResult) error {
		qrs = append(qrs, r)
		return nil
	})
	if err != nil {
		t.Errorf("want nil, got %v", err)
	}
}