code review feedback: test improvements & do not set target alias where it is not valid to set

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

Author: nickvanw

go/test/endtoend/vtgate/misc_test.go

@@ -960,40 +960,39 @@ func TestTabletTargeting(t *testing.T) {
 
 	// Query different replicas and verify different server UUIDs
 	// This proves we're actually hitting different physical tablets
-	for range 10 {
-		// Get server UUID from first replica
-		useStmt = fmt.Sprintf("USE `ks:-80@replica|%s`", instances["-80"]["replica"][0])
-		utils.Exec(t, conn, useStmt)
-		var uuid1 string
-		for i := range 5 {
-			result1 := utils.Exec(t, conn, "SELECT @@server_uuid")
-			require.NotNil(t, result1)
-			require.Greater(t, len(result1.Rows), 0)
-			if i > 0 {
-				// UUID should be the same across multiple queries to same tablet
-				require.Equal(t, uuid1, result1.Rows[0][0].ToString())
-			}
-			uuid1 = result1.Rows[0][0].ToString()
-		}
 
-		// Get server UUID from second replica
-		useStmt = fmt.Sprintf("USE `ks:-80@replica|%s`", instances["-80"]["replica"][1])
-		utils.Exec(t, conn, useStmt)
-		var uuid2 string
-		for i := range 5 {
-			result2 := utils.Exec(t, conn, "SELECT @@server_uuid")
-			require.NotNil(t, result2)
-			require.Greater(t, len(result2.Rows), 0)
-			if i > 0 {
-				// UUID should be the same across multiple queries to same tablet
-				require.Equal(t, uuid2, result2.Rows[0][0].ToString())
-			}
-			uuid2 = result2.Rows[0][0].ToString()
+	// Get server UUID from first replica
+	useStmt = fmt.Sprintf("USE `ks:-80@replica|%s`", instances["-80"]["replica"][0])
+	utils.Exec(t, conn, useStmt)
+	var uuid1 string
+	for i := range 5 {
+		result1 := utils.Exec(t, conn, "SELECT @@server_uuid")
+		require.NotNil(t, result1)
+		require.Greater(t, len(result1.Rows), 0)
+		if i > 0 {
+			// UUID should be the same across multiple queries to same tablet
+			require.Equal(t, uuid1, result1.Rows[0][0].ToString())
 		}
+		uuid1 = result1.Rows[0][0].ToString()
+	}
 
-		// Server UUIDs should be different, proving we're targeting different tablets
-		require.NotEqual(t, uuid1, uuid2, "different replicas should have different server UUIDs")
+	// Get server UUID from second replica
+	useStmt = fmt.Sprintf("USE `ks:-80@replica|%s`", instances["-80"]["replica"][1])
+	utils.Exec(t, conn, useStmt)
+	var uuid2 string
+	for i := range 5 {
+		result2 := utils.Exec(t, conn, "SELECT @@server_uuid")
+		require.NotNil(t, result2)
+		require.Greater(t, len(result2.Rows), 0)
+		if i > 0 {
+			// UUID should be the same across multiple queries to same tablet
+			require.Equal(t, uuid2, result2.Rows[0][0].ToString())
+		}
+		uuid2 = result2.Rows[0][0].ToString()
 	}
+
+	// Server UUIDs should be different, proving we're targeting different tablets
+	require.NotEqual(t, uuid1, uuid2, "different replicas should have different server UUIDs")
 }
 
 // TestDynamicConfig tests the dynamic configurations.

go/vt/vtgate/executorcontext/safe_session_test.go

@@ -222,20 +222,4 @@ func TestTargetTabletAlias(t *testing.T) {
 	// Test: Clear (set to nil)
 	session.SetTargetTabletAlias(nil)
 	assert.Nil(t, session.GetTargetTabletAlias())
-
-	// Test: Thread safety - concurrent access
-	done := make(chan bool)
-	for i := 0; i < 100; i++ {
-		go func(uid uint32) {
-			testAlias := &topodatapb.TabletAlias{Cell: "cell", Uid: uid}
-			session.SetTargetTabletAlias(testAlias)
-			_ = session.GetTargetTabletAlias()
-			done <- true
-		}(uint32(i))
-	}
-	for i := 0; i < 100; i++ {
-		<-done
-	}
-	// Just verify we didn't panic - the final value is non-deterministic
-	assert.NotNil(t, session.GetTargetTabletAlias())
 }

go/vt/vtgate/executorcontext/vcursor_impl.go

@@ -1041,13 +1041,18 @@ func (vc *VCursorImpl) SetTarget(target string) error {
 	if tabletAlias == nil {
 		vc.SafeSession.SetTargetTabletAlias(nil)
 	} else {
+		// Tablet targeting must be set before starting a transaction, not during.
+		// The feature is designed to pick a specific tablet and do work there,
+		// not to switch tablets mid-transaction.
+		if vc.SafeSession.InTransaction() {
+			return vterrors.Errorf(vtrpcpb.Code_FAILED_PRECONDITION,
+				"cannot set tablet target while in a transaction")
+		}
+
 		// 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 {

go/vt/vtgate/scatter_conn.go

@@ -903,9 +903,15 @@ 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
+	// Set tablet alias for routing if tablet-specific targeting is active
 	if targetAlias := session.GetTargetTabletAlias(); targetAlias != nil {
-		info.alias = targetAlias
+		if info.alias == nil {
+			info.alias = targetAlias
+		} else if !proto.Equal(info.alias, targetAlias) {
+			return nil, nil, vterrors.Errorf(vtrpcpb.Code_FAILED_PRECONDITION,
+				"cannot change tablet target mid-transaction: session has %v, target is %v",
+				topoproto.TabletAliasString(info.alias), topoproto.TabletAliasString(targetAlias))
+		}
 	}
 	return info, shardSession, nil
 }