wip: set destination in the right places

Signed-off-by: Nick Van Wiggeren <[email protected]>

Author: nickvanw

go/test/endtoend/vtgate/misc_test.go

@@ -835,8 +835,8 @@ func TestDDLTargeted(t *testing.T) {
 	utils.AssertMatches(t, conn, `select id from ddl_targeted`, `[[INT64(1)]]`)
 }
 
-// TestTabletTargeting tests tablet-specific routing with USE keyspace@tablet-alias syntax.
-// In a sharded keyspace, this validates that tablet-specific routing overrides normal hash-based routing.
+// TestTabletTargeting tests tablet-specific routing with USE keyspace:shard@tablet-alias syntax.
+// When shard is specified, tablet-specific routing bypasses vindex-based shard resolution.
 func TestTabletTargeting(t *testing.T) {
 	ctx := context.Background()
 	conn, err := mysql.Connect(ctx, &vtParams)
@@ -865,41 +865,46 @@ func TestTabletTargeting(t *testing.T) {
 	require.NotEmpty(t, primaryShard80Minus, "no PRIMARY tablet found for -80 shard")
 	require.NotEmpty(t, primaryShard80Plus, "no PRIMARY tablet found for 80- shard")
 
-	// Test: Target the -80 shard's PRIMARY and insert rows that normally hash to -80
-	// id1=1 and id1=2 both hash to -80 shard based on the hash vindex
-	utils.Exec(t, conn, fmt.Sprintf("USE ks@%s", primaryShard80Minus))
+	// Test: Target a specific tablet in -80 shard and insert rows
+	// id1=1 and id1=2 both hash to -80 shard - this validates normal behavior
+	utils.Exec(t, conn, fmt.Sprintf("USE ks:-80@%s", primaryShard80Minus))
 	utils.Exec(t, conn, "insert into t1(id1, id2) values(1, 100), (2, 200)")
 	utils.AssertMatches(t, conn, "select id1 from t1 where id1 in (1, 2) order by id1", "[[INT64(1)] [INT64(2)]]")
 
-	// Test: Verify the other shard (80-) did not receive these writes
-	utils.Exec(t, conn, fmt.Sprintf("USE ks@%s", primaryShard80Plus))
-	utils.AssertIsEmpty(t, conn, "select id1 from t1 where id1 in (1, 2)")
+	// Test: Insert data that would normally hash to 80- shard, but goes to -80 because of shard targeting
+	// id1=4 hashes to 80-, but we're targeting -80 shard explicitly
+	utils.Exec(t, conn, fmt.Sprintf("USE ks:-80@%s", primaryShard80Minus))
+	utils.Exec(t, conn, "insert into t1(id1, id2) values(4, 400)")
+
+	// Verify the data went to -80 shard (not where vindex would have put it)
+	utils.Exec(t, conn, "USE ks:-80")
+	utils.AssertMatches(t, conn, "select id1 from t1 where id1=4", "[[INT64(4)]]")
+
+	// Verify the data did NOT go to 80- shard (where vindex says it should be)
+	utils.Exec(t, conn, "USE ks:80-")
+	utils.AssertIsEmpty(t, conn, "select id1 from t1 where id1=4")
 
 	// Test: Transaction with tablet-specific routing
-	utils.Exec(t, conn, fmt.Sprintf("USE ks@%s", primaryShard80Minus))
+	utils.Exec(t, conn, fmt.Sprintf("USE ks:-80@%s", primaryShard80Minus))
 	utils.Exec(t, conn, "begin")
 	utils.Exec(t, conn, "insert into t1(id1, id2) values(10, 300)")
 	utils.Exec(t, conn, "commit")
 	utils.AssertMatches(t, conn, "select id1 from t1 where id1=10", "[[INT64(10)]]")
 
-	// Test: Verify the other shard still has no data for our test rows
-	utils.Exec(t, conn, fmt.Sprintf("USE ks@%s", primaryShard80Plus))
-	utils.AssertIsEmpty(t, conn, "select id1 from t1 where id1 in (1, 2, 10)")
-
 	// Test: Rollback with tablet-specific routing
-	utils.Exec(t, conn, fmt.Sprintf("USE ks@%s", primaryShard80Minus))
+	utils.Exec(t, conn, fmt.Sprintf("USE ks:-80@%s", primaryShard80Minus))
 	utils.Exec(t, conn, "begin")
-	utils.Exec(t, conn, "insert into t1(id1, id2) values(20, 400)")
+	utils.Exec(t, conn, "insert into t1(id1, id2) values(20, 500)")
 	utils.Exec(t, conn, "rollback")
 	// 20 should not exist
 	utils.AssertIsEmpty(t, conn, "select id1 from t1 where id1=20")
 
 	// Test: Clear tablet targeting returns to normal routing
 	utils.Exec(t, conn, "USE ks")
-	utils.AssertMatches(t, conn, "select id1 from t1 where id1 in (1, 2, 10) order by id1", "[[INT64(1)] [INT64(2)] [INT64(10)]]")
+	utils.AssertMatches(t, conn, "select id1 from t1 where id1 in (1, 2, 4, 10) order by id1", "[[INT64(1)] [INT64(2)] [INT64(4)] [INT64(10)]]")
 
 	// Cleanup
-	utils.Exec(t, conn, "delete from t1 where id1 in (1, 2, 10)")
+	utils.Exec(t, conn, "delete from t1 where id1 in (1, 2, 4, 10)")
 }
 
 // TestDynamicConfig tests the dynamic configurations.

go/vt/vtgate/executor.go

@@ -1104,7 +1104,7 @@ func (e *Executor) SaveVSchema(vschema *vindexes.VSchema, stats *VSchemaStats) {
 }
 
 // ParseDestinationTarget parses destination target string and sets default keyspace if possible.
-func (e *Executor) ParseDestinationTarget(targetString string) (string, topodatapb.TabletType, key.ShardDestination, error) {
+func (e *Executor) ParseDestinationTarget(targetString string) (string, topodatapb.TabletType, key.ShardDestination, *topodatapb.TabletAlias, error) {
 	return econtext.ParseDestinationTarget(targetString, defaultTabletType, e.VSchema())
 }
 

go/vt/vtgate/executor_test.go

@@ -1923,7 +1923,7 @@ func TestParseEmptyTargetSingleKeyspace(t *testing.T) {
 	}
 	r.vschema = altVSchema
 
-	destKeyspace, destTabletType, _, _ := r.ParseDestinationTarget("")
+	destKeyspace, destTabletType, _, _, _ := r.ParseDestinationTarget("")
 	if destKeyspace != KsTestUnsharded || destTabletType != topodatapb.TabletType_PRIMARY {
 		t.Errorf(
 			"parseDestinationTarget(%s): got (%v, %v), want (%v, %v)",
@@ -1947,7 +1947,7 @@ func TestParseEmptyTargetMultiKeyspace(t *testing.T) {
 	}
 	r.vschema = altVSchema
 
-	destKeyspace, destTabletType, _, _ := r.ParseDestinationTarget("")
+	destKeyspace, destTabletType, _, _, _ := r.ParseDestinationTarget("")
 	if destKeyspace != "" || destTabletType != topodatapb.TabletType_PRIMARY {
 		t.Errorf(
 			"parseDestinationTarget(%s): got (%v, %v), want (%v, %v)",
@@ -1970,7 +1970,7 @@ func TestParseTargetSingleKeyspace(t *testing.T) {
 	}
 	r.vschema = altVSchema
 
-	destKeyspace, destTabletType, _, _ := r.ParseDestinationTarget("@replica")
+	destKeyspace, destTabletType, _, _, _ := r.ParseDestinationTarget("@replica")
 	if destKeyspace != KsTestUnsharded || destTabletType != topodatapb.TabletType_REPLICA {
 		t.Errorf(
 			"parseDestinationTarget(%s): got (%v, %v), want (%v, %v)",

go/vt/vtgate/executorcontext/vcursor_impl.go

@@ -208,11 +208,15 @@ func NewVCursorImpl(
 	cfg VCursorConfig,
 	metrics Metrics,
 ) (*VCursorImpl, error) {
-	keyspace, tabletType, destination, err := ParseDestinationTarget(safeSession.TargetString, cfg.DefaultTabletType, vschema)
+	keyspace, tabletType, destination, tabletAlias, err := ParseDestinationTarget(safeSession.TargetString, cfg.DefaultTabletType, vschema)
 	if err != nil {
 		return nil, err
 	}
 
+	// Store tablet alias from target string into session
+	// This ensures tablet-specific routing persists across queries
+	safeSession.SetTargetTabletAlias(tabletAlias)
+
 	var ts *topo.Server
 	// We don't have access to the underlying TopoServer if this vtgate is
 	// filtering keyspaces because we don't have an accurate view of the topo.
@@ -533,20 +537,20 @@ func (vc *VCursorImpl) FindViewTarget(name sqlparser.TableName) (*vindexes.Keysp
 }
 
 func (vc *VCursorImpl) parseDestinationTarget(targetString string) (string, topodatapb.TabletType, key.ShardDestination, error) {
-	return ParseDestinationTarget(targetString, vc.tabletType, vc.vschema)
+	keyspace, tabletType, dest, _, err := ParseDestinationTarget(targetString, vc.tabletType, vc.vschema)
+	return keyspace, tabletType, dest, err
 }
 
 // ParseDestinationTarget parses destination target string and provides a keyspace if possible.
-func ParseDestinationTarget(targetString string, tablet topodatapb.TabletType, vschema *vindexes.VSchema) (string, topodatapb.TabletType, key.ShardDestination, error) {
-	destKeyspace, destTabletType, dest, _, err := topoprotopb.ParseDestination(targetString, tablet)
-	// Note: We ignore the tablet alias here as it's handled separately in SetTarget
+func ParseDestinationTarget(targetString string, tablet topodatapb.TabletType, vschema *vindexes.VSchema) (string, topodatapb.TabletType, key.ShardDestination, *topodatapb.TabletAlias, error) {
+	destKeyspace, destTabletType, dest, tabletAlias, err := topoprotopb.ParseDestination(targetString, tablet)
 	// If the keyspace is not specified, and there is only one keyspace in the VSchema, use that.
 	if destKeyspace == "" && len(vschema.Keyspaces) == 1 {
 		for k := range vschema.Keyspaces {
 			destKeyspace = k
 		}
 	}
-	return destKeyspace, destTabletType, dest, err
+	return destKeyspace, destTabletType, dest, tabletAlias, err
 }
 
 func (vc *VCursorImpl) getDualTable() (*vindexes.BaseTable, vindexes.Vindex, string, topodatapb.TabletType, key.ShardDestination, error) {
@@ -1025,7 +1029,7 @@ func (vc *VCursorImpl) Session() engine.SessionActions {
 }
 
 func (vc *VCursorImpl) SetTarget(target string) error {
-	keyspace, tabletType, _, tabletAlias, err := topoprotopb.ParseDestination(target, vc.config.DefaultTabletType)
+	keyspace, tabletType, destination, tabletAlias, err := topoprotopb.ParseDestination(target, vc.config.DefaultTabletType)
 	if err != nil {
 		return err
 	}
@@ -1052,6 +1056,7 @@ func (vc *VCursorImpl) SetTarget(target string) error {
 	}
 	vc.SafeSession.SetTargetString(target)
 	vc.keyspace = keyspace
+	vc.destination = destination
 	vc.tabletType = tabletType
 	return nil
 }

go/vt/vtgate/executorcontext/vcursor_impl_test.go

@@ -393,7 +393,7 @@ func (f fakeExecutor) SetVitessMetadata(ctx context.Context, name, value string)
 	panic("implement me")
 }
 
-func (f fakeExecutor) ParseDestinationTarget(targetString string) (string, topodatapb.TabletType, key.ShardDestination, error) {
+func (f fakeExecutor) ParseDestinationTarget(targetString string) (string, topodatapb.TabletType, key.ShardDestination, *topodatapb.TabletAlias, error) {
 	// TODO implement me
 	panic("implement me")
 }

go/vt/vtgate/vtgate.go

@@ -558,7 +558,7 @@ func (vtg *VTGate) Execute(
 	prepared bool,
 ) (newSession *vtgatepb.Session, qr *sqltypes.Result, err error) {
 	// In this context, we don't care if we can't fully parse destination
-	destKeyspace, destTabletType, _, _ := vtg.executor.ParseDestinationTarget(session.TargetString)
+	destKeyspace, destTabletType, _, _, _ := vtg.executor.ParseDestinationTarget(session.TargetString)
 	statsKey := []string{"Execute", destKeyspace, topoproto.TabletTypeLString(destTabletType)}
 	defer vtg.timings.Record(statsKey, time.Now())
 
@@ -620,7 +620,7 @@ func (vtg *VTGate) ExecuteMulti(
 // ExecuteBatch executes a batch of queries.
 func (vtg *VTGate) ExecuteBatch(ctx context.Context, session *vtgatepb.Session, sqlList []string, bindVariablesList []map[string]*querypb.BindVariable) (*vtgatepb.Session, []sqltypes.QueryResponse, error) {
 	// In this context, we don't care if we can't fully parse destination
-	destKeyspace, destTabletType, _, _ := vtg.executor.ParseDestinationTarget(session.TargetString)
+	destKeyspace, destTabletType, _, _, _ := vtg.executor.ParseDestinationTarget(session.TargetString)
 	statsKey := []string{"ExecuteBatch", destKeyspace, topoproto.TabletTypeLString(destTabletType)}
 	defer vtg.timings.Record(statsKey, time.Now())
 
@@ -649,7 +649,7 @@ func (vtg *VTGate) ExecuteBatch(ctx context.Context, session *vtgatepb.Session,
 // Note we guarantee the callback will not be called concurrently by multiple go routines.
 func (vtg *VTGate) StreamExecute(ctx context.Context, mysqlCtx vtgateservice.MySQLConnection, session *vtgatepb.Session, sql string, bindVariables map[string]*querypb.BindVariable, callback func(*sqltypes.Result) error) (*vtgatepb.Session, error) {
 	// In this context, we don't care if we can't fully parse destination
-	destKeyspace, destTabletType, _, _ := vtg.executor.ParseDestinationTarget(session.TargetString)
+	destKeyspace, destTabletType, _, _, _ := vtg.executor.ParseDestinationTarget(session.TargetString)
 	statsKey := []string{"StreamExecute", destKeyspace, topoproto.TabletTypeLString(destTabletType)}
 
 	defer vtg.timings.Record(statsKey, time.Now())
@@ -735,7 +735,7 @@ func (vtg *VTGate) CloseSession(ctx context.Context, session *vtgatepb.Session)
 // Prepare supports non-streaming prepare statement query with multi shards
 func (vtg *VTGate) Prepare(ctx context.Context, session *vtgatepb.Session, sql string) (newSession *vtgatepb.Session, fld []*querypb.Field, paramsCount uint16, err error) {
 	// In this context, we don't care if we can't fully parse destination
-	destKeyspace, destTabletType, _, _ := vtg.executor.ParseDestinationTarget(session.TargetString)
+	destKeyspace, destTabletType, _, _, _ := vtg.executor.ParseDestinationTarget(session.TargetString)
 	statsKey := []string{"Prepare", destKeyspace, topoproto.TabletTypeLString(destTabletType)}
 	defer vtg.timings.Record(statsKey, time.Now())