initial commit: pick a tablet

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

Author: nickvanw

go/test/vschemawrapper/vschema_wrapper.go

@@ -236,7 +236,7 @@ func (vw *VSchemaWrapper) ShardDestination() key.ShardDestination {
 }
 
 func (vw *VSchemaWrapper) FindTable(tab sqlparser.TableName) (*vindexes.BaseTable, string, topodatapb.TabletType, key.ShardDestination, error) {
-	destKeyspace, destTabletType, destTarget, err := topoproto.ParseDestination(tab.Qualifier.String(), topodatapb.TabletType_PRIMARY)
+	destKeyspace, destTabletType, destTarget, _, err := topoproto.ParseDestination(tab.Qualifier.String(), topodatapb.TabletType_PRIMARY)
 	if err != nil {
 		return nil, destKeyspace, destTabletType, destTarget, err
 	}
@@ -248,15 +248,15 @@ func (vw *VSchemaWrapper) FindTable(tab sqlparser.TableName) (*vindexes.BaseTabl
 }
 
 func (vw *VSchemaWrapper) FindView(tab sqlparser.TableName) sqlparser.TableStatement {
-	destKeyspace, _, _, err := topoproto.ParseDestination(tab.Qualifier.String(), topodatapb.TabletType_PRIMARY)
+	destKeyspace, _, _, _, err := topoproto.ParseDestination(tab.Qualifier.String(), topodatapb.TabletType_PRIMARY)
 	if err != nil {
 		return nil
 	}
 	return vw.V.FindView(destKeyspace, tab.Name.String())
 }
 
 func (vw *VSchemaWrapper) FindViewTarget(name sqlparser.TableName) (*vindexes.Keyspace, error) {
-	destKeyspace, _, _, err := topoproto.ParseDestination(name.Qualifier.String(), topodatapb.TabletType_PRIMARY)
+	destKeyspace, _, _, _, err := topoproto.ParseDestination(name.Qualifier.String(), topodatapb.TabletType_PRIMARY)
 	if err != nil {
 		return nil, err
 	}
@@ -336,7 +336,7 @@ func (vw *VSchemaWrapper) IsViewsEnabled() bool {
 // FindMirrorRule finds the mirror rule for the requested keyspace, table
 // name, and the tablet type in the VSchema.
 func (vw *VSchemaWrapper) FindMirrorRule(tab sqlparser.TableName) (*vindexes.MirrorRule, error) {
-	destKeyspace, destTabletType, _, err := topoproto.ParseDestination(tab.Qualifier.String(), topodatapb.TabletType_PRIMARY)
+	destKeyspace, destTabletType, _, _, err := topoproto.ParseDestination(tab.Qualifier.String(), topodatapb.TabletType_PRIMARY)
 	if err != nil {
 		return nil, err
 	}

go/vt/topo/topoproto/destination.go

@@ -29,16 +29,37 @@ import (
 
 // ParseDestination parses the string representation of a ShardDestination
 // of the form keyspace:shard@tablet_type. You can use a / instead of a :.
-func ParseDestination(targetString string, defaultTabletType topodatapb.TabletType) (string, topodatapb.TabletType, key.ShardDestination, error) {
+// It also supports tablet-specific routing with keyspace@tablet-alias where
+// tablet-alias is in the format cell-uid (e.g., zone1-0000000100).
+func ParseDestination(targetString string, defaultTabletType topodatapb.TabletType) (string, topodatapb.TabletType, key.ShardDestination, *topodatapb.TabletAlias, error) {
 	var dest key.ShardDestination
 	var keyspace string
+	var tabletAlias *topodatapb.TabletAlias
 	tabletType := defaultTabletType
 
 	last := strings.LastIndexAny(targetString, "@")
 	if last != -1 {
-		// No need to check the error. UNKNOWN will be returned on
-		// error and it will fail downstream.
-		tabletType, _ = ParseTabletType(targetString[last+1:])
+		afterAt := targetString[last+1:]
+		// Try parsing as tablet type first (backward compatible)
+		parsedTabletType, err := ParseTabletType(afterAt)
+		// If tablet type parsing fails or returns UNKNOWN, try parsing as tablet alias
+		if err != nil || parsedTabletType == topodatapb.TabletType_UNKNOWN {
+			// Check if it looks like a tablet alias (contains a dash)
+			if strings.Contains(afterAt, "-") {
+				alias, aliasErr := ParseTabletAlias(afterAt)
+				if aliasErr == nil {
+					tabletAlias = alias
+					// Keep tabletType as defaultTabletType when using tablet alias
+				} else {
+					// If both tablet type and tablet alias parsing fail, keep the UNKNOWN tablet type
+					// which will cause appropriate error handling downstream
+					tabletType = topodatapb.TabletType_UNKNOWN
+				}
+			}
+		} else {
+			// Successfully parsed as tablet type
+			tabletType = parsedTabletType
+		}
 		targetString = targetString[:last]
 	}
 	last = strings.LastIndexAny(targetString, "/:")
@@ -51,29 +72,29 @@ func ParseDestination(targetString string, defaultTabletType topodatapb.TabletTy
 	if last != -1 {
 		rangeEnd := strings.LastIndexAny(targetString, "]")
 		if rangeEnd == -1 {
-			return keyspace, tabletType, dest, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "invalid key range provided. Couldn't find range end ']'")
+			return keyspace, tabletType, dest, nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "invalid key range provided. Couldn't find range end ']'")
 		}
 		rangeString := targetString[last+1 : rangeEnd]
 		if strings.Contains(rangeString, "-") {
 			// Parse as range
 			keyRange, err := key.ParseShardingSpec(rangeString)
 			if err != nil {
-				return keyspace, tabletType, dest, err
+				return keyspace, tabletType, dest, nil, err
 			}
 			if len(keyRange) != 1 {
-				return keyspace, tabletType, dest, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "single keyrange expected in %s", rangeString)
+				return keyspace, tabletType, dest, nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "single keyrange expected in %s", rangeString)
 			}
 			dest = key.DestinationExactKeyRange{KeyRange: keyRange[0]}
 		} else {
 			// Parse as keyspace id
 			destBytes, err := hex.DecodeString(rangeString)
 			if err != nil {
-				return keyspace, tabletType, dest, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "expected valid hex in keyspace id %s", rangeString)
+				return keyspace, tabletType, dest, nil, vterrors.Errorf(vtrpcpb.Code_INVALID_ARGUMENT, "expected valid hex in keyspace id %s", rangeString)
 			}
 			dest = key.DestinationKeyspaceID(destBytes)
 		}
 		targetString = targetString[:last]
 	}
 	keyspace = targetString
-	return keyspace, tabletType, dest, nil
+	return keyspace, tabletType, dest, tabletAlias, nil
 }

go/vt/topo/topoproto/destination_test.go

@@ -90,7 +90,7 @@ func TestParseDestination(t *testing.T) {
 	}}
 
 	for _, tcase := range testcases {
-		if targetKeyspace, targetTabletType, targetDest, _ := ParseDestination(tcase.targetString, topodatapb.TabletType_PRIMARY); !reflect.DeepEqual(targetDest, tcase.dest) || targetKeyspace != tcase.keyspace || targetTabletType != tcase.tabletType {
+		if targetKeyspace, targetTabletType, targetDest, _, _ := ParseDestination(tcase.targetString, topodatapb.TabletType_PRIMARY); !reflect.DeepEqual(targetDest, tcase.dest) || targetKeyspace != tcase.keyspace || targetTabletType != tcase.tabletType {
 			t.Errorf("ParseDestination(%s) - got: (%v, %v, %v), want (%v, %v, %v)",
 				tcase.targetString,
 				targetDest,
@@ -103,19 +103,19 @@ func TestParseDestination(t *testing.T) {
 		}
 	}
 
-	_, _, _, err := ParseDestination("ks[20-40-60]", topodatapb.TabletType_PRIMARY)
+	_, _, _, _, err := ParseDestination("ks[20-40-60]", topodatapb.TabletType_PRIMARY)
 	want := "single keyrange expected in 20-40-60"
 	if err == nil || err.Error() != want {
 		t.Errorf("executorExec error: %v, want %s", err, want)
 	}
 
-	_, _, _, err = ParseDestination("ks[--60]", topodatapb.TabletType_PRIMARY)
+	_, _, _, _, err = ParseDestination("ks[--60]", topodatapb.TabletType_PRIMARY)
 	want = "malformed spec: MinKey/MaxKey cannot be in the middle of the spec: \"--60\""
 	if err == nil || err.Error() != want {
 		t.Errorf("executorExec error: %v, want %s", err, want)
 	}
 
-	_, _, _, err = ParseDestination("ks[qrnqorrs]@primary", topodatapb.TabletType_PRIMARY)
+	_, _, _, _, err = ParseDestination("ks[qrnqorrs]@primary", topodatapb.TabletType_PRIMARY)
 	want = "expected valid hex in keyspace id qrnqorrs"
 	if err == nil || err.Error() != want {
 		t.Errorf("executorExec error: %v, want %s", err, want)

go/vt/vtgate/executorcontext/safe_session.go

@@ -65,6 +65,11 @@ type (
 
 		logging *ExecuteLogger
 
+		// targetTabletAlias is set when using tablet-specific routing via USE keyspace@tablet-alias.
+		// This causes all queries to route to the specified tablet until cleared.
+		// Note: This is stored in the Go wrapper, not in the protobuf Session.
+		targetTabletAlias *topodatapb.TabletAlias
+
 		*vtgatepb.Session
 	}
 
@@ -1143,3 +1148,19 @@ func (l *ExecuteLogger) GetLogs() []engine.ExecuteEntry {
 	copy(result, l.entries)
 	return result
 }
+
+// SetTargetTabletAlias sets the tablet alias for tablet-specific routing.
+// When set, all queries will route to the specified tablet until cleared.
+func (session *SafeSession) SetTargetTabletAlias(alias *topodatapb.TabletAlias) {
+	session.mu.Lock()
+	defer session.mu.Unlock()
+	session.targetTabletAlias = alias
+}
+
+// GetTargetTabletAlias returns the current tablet alias for tablet-specific routing,
+// or nil if not set.
+func (session *SafeSession) GetTargetTabletAlias() *topodatapb.TabletAlias {
+	session.mu.Lock()
+	defer session.mu.Unlock()
+	return session.targetTabletAlias
+}

go/vt/vtgate/executorcontext/vcursor_impl.go

@@ -538,7 +538,8 @@ func (vc *VCursorImpl) parseDestinationTarget(targetString string) (string, topo
 
 // 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)
+	destKeyspace, destTabletType, dest, _, err := topoprotopb.ParseDestination(targetString, tablet)
+	// Note: We ignore the tablet alias here as it's handled separately in SetTarget
 	// 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 {
@@ -1024,10 +1025,24 @@ func (vc *VCursorImpl) Session() engine.SessionActions {
 }
 
 func (vc *VCursorImpl) SetTarget(target string) error {
-	keyspace, tabletType, _, err := topoprotopb.ParseDestination(target, vc.config.DefaultTabletType)
+	keyspace, tabletType, _, tabletAlias, err := topoprotopb.ParseDestination(target, vc.config.DefaultTabletType)
 	if err != nil {
 		return err
 	}
+
+	// Clear any existing tablet-specific routing if not specified
+	if tabletAlias == nil {
+		vc.SafeSession.SetTargetTabletAlias(nil)
+	} else {
+		// Store tablet alias in session for routing
+		// Note: We don't validate the tablet exists here - if it doesn't exist or is
+		// unreachable, the query will fail during execution with a clear error message
+		vc.SafeSession.SetTargetTabletAlias(tabletAlias)
+
+		// Keep tabletType as determined by ParseDestination (defaultTabletType)
+		// The actual routing uses the tablet alias, so the type is only for VCursor state
+	}
+
 	if _, ok := vc.vschema.Keyspaces[keyspace]; !ignoreKeyspace(keyspace) && !ok {
 		return vterrors.VT05003(keyspace)
 	}

go/vt/vtgate/scatter_conn.go

@@ -862,6 +862,13 @@ func requireNewQS(err error, target *querypb.Target) bool {
 // actionInfo looks at the current session, and returns information about what needs to be done for this tablet
 func actionInfo(ctx context.Context, target *querypb.Target, session *econtext.SafeSession, autocommit bool, txMode vtgatepb.TransactionMode) (*shardActionInfo, *vtgatepb.Session_ShardSession, error) {
 	if !(session.InTransaction() || session.InReservedConn()) {
+		// Check for tablet-specific routing for non-transactional queries
+		if alias := session.GetTargetTabletAlias(); alias != nil {
+			return &shardActionInfo{
+				actionNeeded: nothing,
+				alias:        alias,
+			}, nil, nil
+		}
 		return &shardActionInfo{}, nil, nil
 	}
 	ignoreSession := ctx.Value(engine.IgnoreReserveTxn)
@@ -899,6 +906,10 @@ func actionInfo(ctx context.Context, target *querypb.Target, session *econtext.S
 		info.alias = shardSession.TabletAlias
 		info.rowsAffected = shardSession.RowsAffected
 	}
+	// Override alias if tablet-specific routing is set
+	if targetAlias := session.GetTargetTabletAlias(); targetAlias != nil {
+		info.alias = targetAlias
+	}
 	return info, shardSession, nil
 }
 

go/vtbench/client.go

@@ -124,7 +124,7 @@ func (c *grpcVttabletConn) connect(ctx context.Context, cp ConnParams) error {
 	})
 
 	// parse the "db" into the keyspace/shard target
-	keyspace, tabletType, dest, err := topoproto.ParseDestination(cp.DB, topodatapb.TabletType_PRIMARY)
+	keyspace, tabletType, dest, _, err := topoproto.ParseDestination(cp.DB, topodatapb.TabletType_PRIMARY)
 	if err != nil {
 		return err
 	}