add tests

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

Author: nickvanw

go/test/endtoend/vtgate/misc_test.go

@@ -835,6 +835,56 @@ 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.
+func TestTabletTargeting(t *testing.T) {
+	ctx := context.Background()
+	conn, err := mysql.Connect(ctx, &vtParams)
+	require.NoError(t, err)
+	defer conn.Close()
+
+	// Get tablet aliases from show vitess_tablets
+	qr := utils.Exec(t, conn, "show vitess_tablets")
+	require.Greater(t, len(qr.Rows), 0, "no tablets found")
+
+	// Find a PRIMARY tablet
+	var primaryAlias string
+	for _, row := range qr.Rows {
+		tabletType := row[3].ToString()
+		if tabletType == "PRIMARY" {
+			primaryAlias = row[0].ToString()
+			break
+		}
+	}
+	require.NotEmpty(t, primaryAlias, "no PRIMARY tablet found")
+
+	// Test: Basic tablet-specific SELECT
+	utils.Exec(t, conn, fmt.Sprintf("USE ks@%s", primaryAlias))
+	utils.Exec(t, conn, "insert into t1(id1, id2) values(9001, 1)")
+	utils.AssertMatches(t, conn, "select id1 from t1 where id1=9001", "[[INT64(9001)]]")
+
+	// Test: Clear tablet targeting with plain USE
+	utils.Exec(t, conn, "USE ks")
+	utils.AssertMatches(t, conn, "select id1 from t1 where id1=9001", "[[INT64(9001)]]")
+
+	// Test: Transaction with tablet-specific routing
+	utils.Exec(t, conn, fmt.Sprintf("USE ks@%s", primaryAlias))
+	utils.Exec(t, conn, "begin")
+	utils.Exec(t, conn, "insert into t1(id1, id2) values(9002, 1)")
+	utils.Exec(t, conn, "commit")
+	utils.AssertMatches(t, conn, "select id1 from t1 where id1 in (9001, 9002) order by id1", "[[INT64(9001)] [INT64(9002)]]")
+
+	// Test: Rollback with tablet-specific routing
+	utils.Exec(t, conn, fmt.Sprintf("USE ks@%s", primaryAlias))
+	utils.Exec(t, conn, "begin")
+	utils.Exec(t, conn, "insert into t1(id1, id2) values(9003, 1)")
+	utils.Exec(t, conn, "rollback")
+	// 9003 should not exist
+	utils.AssertIsEmpty(t, conn, "select id1 from t1 where id1=9003")
+
+	// Cleanup
+	utils.Exec(t, conn, "delete from t1 where id1 >= 9001 and id1 <= 9003")
+}
+
 // TestDynamicConfig tests the dynamic configurations.
 func TestDynamicConfig(t *testing.T) {
 	t.Run("DiscoveryLowReplicationLag", func(t *testing.T) {

go/vt/topo/topoproto/destination_test.go

@@ -35,6 +35,7 @@ func TestParseDestination(t *testing.T) {
 		dest         key.ShardDestination
 		keyspace     string
 		tabletType   topodatapb.TabletType
+		tabletAlias  *topodatapb.TabletAlias
 	}{{
 		targetString: "ks[10-20]@primary",
 		keyspace:     "ks",
@@ -87,18 +88,36 @@ func TestParseDestination(t *testing.T) {
 		keyspace:     "ks",
 		dest:         key.DestinationShard("-80"),
 		tabletType:   topodatapb.TabletType_PRIMARY,
+	}, {
+		targetString: "ks@zone1-0000000100",
+		keyspace:     "ks",
+		tabletType:   topodatapb.TabletType_PRIMARY,
+		tabletAlias:  &topodatapb.TabletAlias{Cell: "zone1", Uid: 100},
+	}, {
+		targetString: "ks:-80@zone1-0000000100",
+		keyspace:     "ks",
+		dest:         key.DestinationShard("-80"),
+		tabletType:   topodatapb.TabletType_PRIMARY,
+		tabletAlias:  &topodatapb.TabletAlias{Cell: "zone1", Uid: 100},
+	}, {
+		targetString: "@zone1-0000000200",
+		keyspace:     "",
+		tabletType:   topodatapb.TabletType_PRIMARY,
+		tabletAlias:  &topodatapb.TabletAlias{Cell: "zone1", Uid: 200},
 	}}
 
 	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 {
-			t.Errorf("ParseDestination(%s) - got: (%v, %v, %v), want (%v, %v, %v)",
+		if targetKeyspace, targetTabletType, targetDest, targetAlias, _ := ParseDestination(tcase.targetString, topodatapb.TabletType_PRIMARY); !reflect.DeepEqual(targetDest, tcase.dest) || targetKeyspace != tcase.keyspace || targetTabletType != tcase.tabletType || !reflect.DeepEqual(targetAlias, tcase.tabletAlias) {
+			t.Errorf("ParseDestination(%s) - got: (%v, %v, %v, %v), want (%v, %v, %v, %v)",
 				tcase.targetString,
-				targetDest,
 				targetKeyspace,
 				targetTabletType,
-				tcase.dest,
+				targetDest,
+				targetAlias,
 				tcase.keyspace,
 				tcase.tabletType,
+				tcase.dest,
+				tcase.tabletAlias,
 			)
 		}
 	}

go/vt/vtgate/executorcontext/safe_session_test.go

@@ -205,3 +205,37 @@ func TestTimeZone(t *testing.T) {
 		})
 	}
 }
+
+// TestTargetTabletAlias tests the SetTargetTabletAlias and GetTargetTabletAlias methods.
+func TestTargetTabletAlias(t *testing.T) {
+	session := NewSafeSession(&vtgatepb.Session{})
+
+	// Test: initially nil
+	assert.Nil(t, session.GetTargetTabletAlias())
+
+	// Test: Set and get
+	alias := &topodatapb.TabletAlias{Cell: "zone1", Uid: 100}
+	session.SetTargetTabletAlias(alias)
+	got := session.GetTargetTabletAlias()
+	assert.Equal(t, alias, got)
+
+	// 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/scatter_conn_test.go

@@ -632,3 +632,68 @@ func TestIsConnClosed(t *testing.T) {
 		})
 	}
 }
+
+// TestActionInfoWithTabletAlias tests the actionInfo function with tablet-specific routing.
+func TestActionInfoWithTabletAlias(t *testing.T) {
+	ctx := utils.LeakCheckContext(t)
+	target := &querypb.Target{
+		Keyspace:   "ks",
+		Shard:      "-80",
+		TabletType: topodatapb.TabletType_PRIMARY,
+	}
+	tabletAlias := &topodatapb.TabletAlias{Cell: "zone1", Uid: 100}
+
+	t.Run("non-transactional with tablet alias", func(t *testing.T) {
+		session := econtext.NewSafeSession(&vtgatepb.Session{})
+		session.SetTargetTabletAlias(tabletAlias)
+
+		info, shardSession, err := actionInfo(ctx, target, session, false, vtgatepb.TransactionMode_MULTI)
+		require.NoError(t, err)
+		assert.Nil(t, shardSession)
+		assert.Equal(t, nothing, info.actionNeeded)
+		assert.Equal(t, tabletAlias, info.alias)
+	})
+
+	t.Run("transaction begin with tablet alias", func(t *testing.T) {
+		session := econtext.NewSafeSession(&vtgatepb.Session{
+			InTransaction: true,
+		})
+		session.SetTargetTabletAlias(tabletAlias)
+
+		info, shardSession, err := actionInfo(ctx, target, session, false, vtgatepb.TransactionMode_MULTI)
+		require.NoError(t, err)
+		assert.Nil(t, shardSession)
+		assert.Equal(t, begin, info.actionNeeded)
+		assert.Equal(t, tabletAlias, info.alias)
+	})
+
+	t.Run("existing transaction with tablet alias", func(t *testing.T) {
+		session := econtext.NewSafeSession(&vtgatepb.Session{
+			InTransaction: true,
+			ShardSessions: []*vtgatepb.Session_ShardSession{{
+				Target:        target,
+				TransactionId: 12345,
+				TabletAlias:   &topodatapb.TabletAlias{Cell: "zone1", Uid: 50},
+			}},
+		})
+		session.SetTargetTabletAlias(tabletAlias)
+
+		info, shardSession, err := actionInfo(ctx, target, session, false, vtgatepb.TransactionMode_MULTI)
+		require.NoError(t, err)
+		assert.NotNil(t, shardSession)
+		assert.Equal(t, int64(12345), info.transactionID)
+		assert.Equal(t, nothing, info.actionNeeded)
+		// Tablet alias should be overridden
+		assert.Equal(t, tabletAlias, info.alias)
+	})
+
+	t.Run("no tablet alias - existing behavior", func(t *testing.T) {
+		session := econtext.NewSafeSession(&vtgatepb.Session{})
+
+		info, shardSession, err := actionInfo(ctx, target, session, false, vtgatepb.TransactionMode_MULTI)
+		require.NoError(t, err)
+		assert.Nil(t, shardSession)
+		assert.Equal(t, nothing, info.actionNeeded)
+		assert.Nil(t, info.alias)
+	})
+}