infoschema: fix information_schema.tikv_region_status for nextgen (#67444)

pkg/executor/BUILD.bazel

@@ -413,6 +413,7 @@ go_test(
         "//pkg/config",
         "//pkg/config/kerneltype",
         "//pkg/ddl",
+        "//pkg/ddl/placement",
         "//pkg/ddl/util",
         "//pkg/distsql",
         "//pkg/distsql/context",

pkg/executor/infoschema_cluster_table_test.go

@@ -16,8 +16,12 @@ package executor_test
 
 import (
 	"crypto/tls"
+	"encoding/hex"
+	"encoding/json"
 	"fmt"
+	"hash/crc32"
 	"net"
+	"net/http"
 	"net/http/httptest"
 	"strconv"
 	"strings"
@@ -28,6 +32,7 @@ import (
 	"github.com/pingcap/failpoint"
 	"github.com/pingcap/fn"
 	"github.com/pingcap/tidb/pkg/config"
+	"github.com/pingcap/tidb/pkg/ddl/placement"
 	"github.com/pingcap/tidb/pkg/domain"
 	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/parser/auth"
@@ -129,14 +134,14 @@ func (s *infosSchemaClusterTableSuite) setUpMockPDHTTPServer() (*httptest.Server
 		}, nil
 	}))
 	// mock regions
-	router.Handle(pd.Regions, fn.Wrap(func() (*pd.RegionsInfo, error) {
+	regionsResp := func(regionID int64, startKey, endKey string) *pd.RegionsInfo {
 		return &pd.RegionsInfo{
 			Count: 1,
 			Regions: []pd.RegionInfo{
 				{
-					ID:       1,
-					StartKey: "",
-					EndKey:   "",
+					ID:       regionID,
+					StartKey: startKey,
+					EndKey:   endKey,
 					Epoch: pd.RegionEpoch{
 						ConfVer: 1,
 						Version: 2,
@@ -147,8 +152,31 @@ func (s *infosSchemaClusterTableSuite) setUpMockPDHTTPServer() (*httptest.Server
 					ApproximateKeys: 1000,
 				},
 			},
-		}, nil
+		}
+	}
+	router.Handle(pd.Regions, fn.Wrap(func() (*pd.RegionsInfo, error) {
+		return regionsResp(1, "", ""), nil
 	}))
+	router.HandleFunc("/pd/api/v1/regions/key", func(w http.ResponseWriter, r *http.Request) {
+		startKey := strings.ToUpper(hex.EncodeToString([]byte(r.URL.Query().Get("key"))))
+		endKey := strings.ToUpper(hex.EncodeToString([]byte(r.URL.Query().Get("end_key"))))
+		regionID := int64(crc32.ChecksumIEEE([]byte(startKey+"|"+endKey))) + 1
+		w.Header().Set("Content-Type", "application/json")
+		if err := json.NewEncoder(w).Encode(regionsResp(regionID, startKey, endKey)); err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+		}
+	})
+	router.HandleFunc(pd.MinResolvedTSPrefix, func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{"min_resolved_ts":1,"is_real_time":true}`))
+	})
+	router.HandleFunc(pd.PlacementRuleGroupByID(placement.TiFlashRuleGroupID), func(w http.ResponseWriter, _ *http.Request) {
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(fmt.Sprintf(`{"id":"%s","index":%d,"override":false}`, placement.TiFlashRuleGroupID, placement.RuleIndexTiFlash)))
+	})
+	router.HandleFunc("/pd/api/v1/config/rule_group", func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusOK)
+	})
 	// mock PD API
 	router.Handle(pd.Status, fn.Wrap(func() (any, error) {
 		return struct {
@@ -322,48 +350,61 @@ func TestTikvRegionStatus(t *testing.T) {
 	tk.MustExec("use test")
 	tk.MustExec("drop table if exists test_t1")
 	tk.MustExec(`CREATE TABLE test_t1 ( a int(11) DEFAULT NULL, b int(11) DEFAULT NULL, c int(11) DEFAULT NULL)`)
-	tk.MustQuery("select REGION_ID, DB_NAME, TABLE_NAME, IS_INDEX, INDEX_ID, INDEX_NAME, IS_PARTITION, PARTITION_NAME from information_schema.TIKV_REGION_STATUS where DB_NAME = 'test' and TABLE_NAME = 'test_t1'").Check(testkit.Rows(
-		"1 test test_t1 0 <nil> <nil> 0 <nil>",
+	tk.MustQuery("select DB_NAME, TABLE_NAME, IS_INDEX, INDEX_ID, INDEX_NAME, IS_PARTITION, PARTITION_NAME from information_schema.TIKV_REGION_STATUS where DB_NAME = 'test' and TABLE_NAME = 'test_t1'").Check(testkit.Rows(
+		"test test_t1 0 <nil> <nil> 0 <nil>",
 	))
 
 	tk.MustExec("alter table test_t1 add index p_a (a)")
-	tk.MustQuery("select REGION_ID, DB_NAME, TABLE_NAME, IS_INDEX, INDEX_NAME, IS_PARTITION, PARTITION_NAME from information_schema.TIKV_REGION_STATUS where DB_NAME = 'test' and TABLE_NAME = 'test_t1' order by IS_INDEX").Check(testkit.Rows(
-		"1 test test_t1 0 <nil> 0 <nil>",
-		"1 test test_t1 1 p_a 0 <nil>",
+	tk.MustQuery("select DB_NAME, TABLE_NAME, IS_INDEX, INDEX_NAME, IS_PARTITION, PARTITION_NAME from information_schema.TIKV_REGION_STATUS where DB_NAME = 'test' and TABLE_NAME = 'test_t1' order by IS_INDEX").Check(testkit.Rows(
+		"test test_t1 0 <nil> 0 <nil>",
+		"test test_t1 1 p_a 0 <nil>",
+	))
+	tableID := tk.MustQuery("select TIDB_TABLE_ID from information_schema.tables where TABLE_SCHEMA = 'test' and TABLE_NAME = 'test_t1'").Rows()[0][0]
+	tk.MustQuery(fmt.Sprintf("select DB_NAME, TABLE_NAME, IS_INDEX, INDEX_NAME, IS_PARTITION, PARTITION_NAME from information_schema.TIKV_REGION_STATUS where TABLE_ID = %v order by IS_INDEX", tableID)).Check(testkit.Rows(
+		"test test_t1 0 <nil> 0 <nil>",
+		"test test_t1 1 p_a 0 <nil>",
 	))
 
 	tk.MustExec("alter table test_t1 add unique p_b (b);")
-	tk.MustQuery("select REGION_ID, DB_NAME, TABLE_NAME, IS_INDEX, INDEX_NAME, IS_PARTITION, PARTITION_NAME from information_schema.TIKV_REGION_STATUS where DB_NAME = 'test' and TABLE_NAME = 'test_t1' order by IS_INDEX, INDEX_NAME").Check(testkit.Rows(
-		"1 test test_t1 0 <nil> 0 <nil>",
-		"1 test test_t1 1 p_a 0 <nil>",
-		"1 test test_t1 1 p_b 0 <nil>",
+	tk.MustQuery("select DB_NAME, TABLE_NAME, IS_INDEX, INDEX_NAME, IS_PARTITION, PARTITION_NAME from information_schema.TIKV_REGION_STATUS where DB_NAME = 'test' and TABLE_NAME = 'test_t1' order by IS_INDEX, INDEX_NAME").Check(testkit.Rows(
+		"test test_t1 0 <nil> 0 <nil>",
+		"test test_t1 1 p_a 0 <nil>",
+		"test test_t1 1 p_b 0 <nil>",
 	))
 
 	tk.MustExec("drop table if exists test_t2")
 	tk.MustExec(`CREATE TABLE test_t2 ( a int(11) DEFAULT NULL, b int(11) DEFAULT NULL, c int(11) DEFAULT NULL)
 		PARTITION BY RANGE (c) (
 		PARTITION p0 VALUES LESS THAN (10),
 		PARTITION p1 VALUES LESS THAN (MAXVALUE))`)
-	tk.MustQuery("select REGION_ID, DB_NAME, TABLE_NAME, IS_INDEX, INDEX_ID, INDEX_NAME, IS_PARTITION, PARTITION_NAME from information_schema.TIKV_REGION_STATUS where DB_NAME = 'test' and TABLE_NAME = 'test_t2' order by PARTITION_NAME").Check(testkit.Rows(
-		"1 test test_t2 0 <nil> <nil> 1 p0",
-		"1 test test_t2 0 <nil> <nil> 1 p1",
+	tk.MustQuery("select DB_NAME, TABLE_NAME, IS_INDEX, INDEX_ID, INDEX_NAME, IS_PARTITION, PARTITION_NAME from information_schema.TIKV_REGION_STATUS where DB_NAME = 'test' and TABLE_NAME = 'test_t2' order by PARTITION_NAME").Check(testkit.Rows(
+		"test test_t2 0 <nil> <nil> 1 p0",
+		"test test_t2 0 <nil> <nil> 1 p1",
 	))
 
 	tk.MustExec("alter table test_t2 add index p_a (a)")
-	tk.MustQuery("select REGION_ID, DB_NAME, TABLE_NAME, IS_INDEX, INDEX_NAME, IS_PARTITION, PARTITION_NAME from information_schema.TIKV_REGION_STATUS where DB_NAME = 'test' and TABLE_NAME = 'test_t2' order by IS_INDEX, PARTITION_NAME").Check(testkit.Rows(
-		"1 test test_t2 0 <nil> 1 p0",
-		"1 test test_t2 0 <nil> 1 p1",
-		"1 test test_t2 1 p_a 1 p0",
-		"1 test test_t2 1 p_a 1 p1",
+	tk.MustQuery("select DB_NAME, TABLE_NAME, IS_INDEX, INDEX_NAME, IS_PARTITION, PARTITION_NAME from information_schema.TIKV_REGION_STATUS where DB_NAME = 'test' and TABLE_NAME = 'test_t2' order by IS_INDEX, PARTITION_NAME").Check(testkit.Rows(
+		"test test_t2 0 <nil> 1 p0",
+		"test test_t2 0 <nil> 1 p1",
+		"test test_t2 1 p_a 1 p0",
+		"test test_t2 1 p_a 1 p1",
 	))
 
 	tk.MustExec("alter table test_t2 add unique p_b (b) global")
-	tk.MustQuery("select REGION_ID, DB_NAME, TABLE_NAME, IS_INDEX, INDEX_NAME, IS_PARTITION, PARTITION_NAME from information_schema.TIKV_REGION_STATUS where DB_NAME = 'test' and TABLE_NAME = 'test_t2' order by IS_INDEX, IS_PARTITION desc, PARTITION_NAME").Check(testkit.Rows(
-		"1 test test_t2 0 <nil> 1 p0",
-		"1 test test_t2 0 <nil> 1 p1",
-		"1 test test_t2 1 p_a 1 p0",
-		"1 test test_t2 1 p_a 1 p1",
-		"1 test test_t2 1 p_b 0 <nil>",
+	tk.MustQuery("select DB_NAME, TABLE_NAME, IS_INDEX, INDEX_NAME, IS_PARTITION, PARTITION_NAME from information_schema.TIKV_REGION_STATUS where DB_NAME = 'test' and TABLE_NAME = 'test_t2' order by IS_INDEX, IS_PARTITION desc, PARTITION_NAME").Check(testkit.Rows(
+		"test test_t2 0 <nil> 1 p0",
+		"test test_t2 0 <nil> 1 p1",
+		"test test_t2 1 p_a 1 p0",
+		"test test_t2 1 p_a 1 p1",
+		"test test_t2 1 p_b 0 <nil>",
+	))
+	tableID = tk.MustQuery("select TIDB_TABLE_ID from information_schema.tables where TABLE_SCHEMA = 'test' and TABLE_NAME = 'test_t2'").Rows()[0][0]
+	tk.MustQuery(fmt.Sprintf("select DB_NAME, TABLE_NAME, IS_INDEX, INDEX_NAME, IS_PARTITION, PARTITION_NAME from information_schema.TIKV_REGION_STATUS where TABLE_ID = %v order by IS_INDEX, IS_PARTITION desc, PARTITION_NAME", tableID)).Check(testkit.Rows(
+		"test test_t2 0 <nil> 1 p0",
+		"test test_t2 0 <nil> 1 p1",
+		"test test_t2 1 p_a 1 p0",
+		"test test_t2 1 p_a 1 p1",
+		"test test_t2 1 p_b 0 <nil>",
 	))
 
 	// Run the query to ensure virtual schemas are excluded and expect no rows to be returned

pkg/executor/infoschema_reader.go

@@ -2079,11 +2079,7 @@ func (e *memtableRetriever) setDataForTiKVRegionStatus(ctx context.Context, sctx
 		}
 	}
 	if !requestByTableRange {
-		pdCli, err := tikvHelper.TryGetPDHTTPClient()
-		if err != nil {
-			return err
-		}
-		allRegionsInfo, err = pdCli.GetRegions(ctx)
+		allRegionsInfo, err = tikvHelper.GetRegions(ctx)
 		if err != nil {
 			return err
 		}
@@ -2123,16 +2119,16 @@ func (e *memtableRetriever) getRegionsInfoForTable(ctx context.Context, h *helpe
 		return nil, infoschema.ErrTableNotExists.GenWithStackByArgs(tableID)
 	}
 
+	allRegionsInfo, err := e.getRegionsInfoForSingleTable(ctx, h, tableID)
+	if err != nil {
+		return nil, err
+	}
+
 	pt := tbl.Meta().GetPartitionInfo()
 	if pt == nil {
-		regionsInfo, err := e.getRegionsInfoForSingleTable(ctx, h, tableID)
-		if err != nil {
-			return nil, err
-		}
-		return regionsInfo, nil
+		return allRegionsInfo, nil
 	}
 
-	var allRegionsInfo *pd.RegionsInfo
 	for _, def := range pt.Definitions {
 		regionsInfo, err := e.getRegionsInfoForSingleTable(ctx, h, def.ID)
 		if err != nil {
@@ -2148,7 +2144,9 @@ func (*memtableRetriever) getRegionsInfoForSingleTable(ctx context.Context, help
 	if err != nil {
 		return nil, err
 	}
-	sk, ek := tablecodec.GetTableHandleKeyRange(tableID)
+	// Query the whole table prefix so both record and index regions are covered.
+	sk := tablecodec.EncodeTablePrefix(tableID)
+	ek := sk.PrefixNext()
 	start, end := helper.Store.GetCodec().EncodeRegionRange(sk, ek)
 	return pdCli.GetRegionsByKeyRange(ctx, pd.NewKeyRange(start, end), -1)
 }

pkg/executor/memtable_reader.go

@@ -917,7 +917,7 @@ func (e *tikvRegionPeersRetriever) retrieve(ctx context.Context, sctx sessionctx
 	storeMap := make(map[int64]struct{})
 
 	if len(e.extractor.StoreIDs) == 0 && len(e.extractor.RegionIDs) == 0 {
-		regionsInfo, err := pdCli.GetRegions(ctx)
+		regionsInfo, err := tikvHelper.GetRegions(ctx)
 		if err != nil {
 			return nil, err
 		}

pkg/store/helper/helper.go

@@ -114,6 +114,17 @@ func (h *Helper) TryGetPDHTTPClient() (pd.Client, error) {
 	return h.pdHTTPCli, nil
 }
 
+// GetRegions fetches regions for the current store. In keyspace-aware mode, it
+// restricts the scan to the current keyspace to avoid mixing regions from other keyspaces.
+func (h *Helper) GetRegions(ctx context.Context) (*pd.RegionsInfo, error) {
+	pdCli, err := h.TryGetPDHTTPClient()
+	if err != nil {
+		return nil, err
+	}
+	startKey, endKey := h.Store.GetCodec().EncodeRegionRange(nil, nil)
+	return pdCli.GetRegionsByKeyRange(ctx, pd.NewKeyRange(startKey, endKey), -1)
+}
+
 // MaxBackoffTimeoutForMvccGet is a derived value from previous implementation possible experiencing value 5000ms.
 const MaxBackoffTimeoutForMvccGet = 5000
 

tests/realtikvtest/sessiontest/BUILD.bazel

@@ -12,10 +12,12 @@ go_test(
     ],
     flaky = True,
     race = "on",
-    shard_count = 16,
+    shard_count = 17,
     deps = [
         "//pkg/config",
+        "//pkg/config/kerneltype",
         "//pkg/infoschema",
+        "//pkg/keyspace",
         "//pkg/meta",
         "//pkg/parser/ast",
         "//pkg/session",

tests/realtikvtest/sessiontest/infoschema_test.go

@@ -15,32 +15,59 @@
 package sessiontest
 
 import (
-	"context"
 	"fmt"
+	"slices"
+	"strings"
 	"testing"
 
-	"github.com/pingcap/tidb/pkg/parser/ast"
+	"github.com/pingcap/tidb/pkg/config/kerneltype"
+	"github.com/pingcap/tidb/pkg/keyspace"
 	"github.com/pingcap/tidb/pkg/testkit"
 	"github.com/pingcap/tidb/tests/realtikvtest"
 	"github.com/stretchr/testify/require"
 )
 
 func TestNextGenTiKVRegionStatus(t *testing.T) {
-	store, dom := realtikvtest.CreateMockStoreAndDomainAndSetup(t)
+	store := realtikvtest.CreateMockStoreAndSetup(t)
 	tk := testkit.NewTestKit(t, store)
 	tk.MustExec("use test")
-	tk.MustExec("create table t (a int);")
-	tbl, err := dom.InfoSchema().TableByName(context.Background(), ast.NewCIStr("test"), ast.NewCIStr("t"))
-	require.NoError(t, err)
+	tk.MustExec("drop table if exists t")
+	tk.MustExec("create table t (a int, key idx(a));")
+	tk.MustExec("split table t between (0) and (10000) regions 4")
+	tk.MustExec("split table t index idx between (0) and (10000) regions 4")
 
-	showRegions := tk.MustQuery("show table t regions").Rows()
-	t.Log(showRegions)
-	require.Equal(t, 1, len(showRegions), showRegions)
+	tableID := tk.MustQuery(`select tidb_table_id from information_schema.tables where table_schema = 'test' and table_name = 't'`).Rows()[0][0]
+	showRegions := uniqueSortedRegionIDs(tk.MustQuery("show table t regions").Rows())
+	showIndexRegions := uniqueSortedRegionIDs(tk.MustQuery("show table t index idx regions").Rows())
 	tikvRegions := tk.MustQuery(fmt.Sprintf(
-		"select region_id from information_schema.tikv_region_status where table_id = %d", tbl.Meta().ID)).Rows()
-	require.Equal(t, 1, len(tikvRegions), tikvRegions)
-	t.Log(tikvRegions)
-	require.Equal(t, showRegions[0][0], tikvRegions[0][0])
+		"select region_id from information_schema.tikv_region_status where table_id = %v", tableID)).Rows()
+	tikvIndexRegions := tk.MustQuery(fmt.Sprintf(
+		"select region_id from information_schema.tikv_region_status where table_id = %v and is_index = 1", tableID)).Rows()
+	require.Equal(t, showRegions, uniqueSortedRegionIDs(tikvRegions))
+	require.Equal(t, showIndexRegions, uniqueSortedRegionIDs(tikvIndexRegions))
+}
+
+func TestNextGenTiKVRegionStatusDoesNotMixOtherKeyspaces(t *testing.T) {
+	if kerneltype.IsClassic() {
+		t.Skip("only runs in nextgen kernel")
+	}
+
+	runtimes := realtikvtest.PrepareForCrossKSTest(t, "keyspace1")
+	sysTK := testkit.NewTestKit(t, runtimes[keyspace.System].Store)
+	sysTK.MustExec("create database if not exists sys_region_status")
+	sysTK.MustExec("use sys_region_status")
+	sysTK.MustExec("drop table if exists t")
+	sysTK.MustExec("create table t (a int, key idx(a))")
+	sysTK.MustExec("split table t between (0) and (10000) regions 4")
+	sysTK.MustExec("split table t index idx between (0) and (10000) regions 4")
+
+	systemRegionIDs := uniqueSortedRegionIDs(sysTK.MustQuery("show table t regions").Rows())
+	require.NotEmpty(t, systemRegionIDs)
+
+	userTK := testkit.NewTestKit(t, runtimes["keyspace1"].Store)
+	userTK.MustQuery(fmt.Sprintf(
+		"select count(*) from information_schema.tikv_region_status where region_id in (%s)",
+		strings.Join(systemRegionIDs, ","))).Check(testkit.Rows("0"))
 }
 
 func TestTableReaderWithSnapshot(t *testing.T) {
@@ -59,3 +86,18 @@ func TestTableReaderWithSnapshot(t *testing.T) {
 	tk.MustExec("set @@tidb_snapshot=@ts;")
 	tk.MustQuery("SELECT TABLE_NAME,TABLE_TYPE,AVG_ROW_LENGTH FROM INFORMATION_SCHEMA.TABLES WHERE TABLE_SCHEMA='test' AND (TABLE_TYPE='BASE TABLE')").Check(testkit.Rows("t BASE TABLE 0"))
 }
+
+func uniqueSortedRegionIDs(rows [][]any) []string {
+	seen := make(map[string]struct{}, len(rows))
+	ids := make([]string, 0, len(rows))
+	for _, row := range rows {
+		id := row[0].(string)
+		if _, ok := seen[id]; ok {
+			continue
+		}
+		seen[id] = struct{}{}
+		ids = append(ids, id)
+	}
+	slices.Sort(ids)
+	return ids
+}