planner: fix partition table getting error result when select `TIKV_R…

…EGION_STATUS` with `table_id`. (#37498)

close #37436

executor/infoschema_reader.go

@@ -1674,11 +1674,12 @@ func (e *memtableRetriever) setDataForTiKVRegionStatus(ctx sessionctx.Context) (
 	}
 	requestByTableRange := false
 	allRegionsInfo := helper.NewRegionsInfo()
+	is := ctx.GetDomainInfoSchema().(infoschema.InfoSchema)
 	if e.extractor != nil {
 		extractor, ok := e.extractor.(*plannercore.TiKVRegionStatusExtractor)
 		if ok && len(extractor.GetTablesID()) > 0 {
 			for _, tableID := range extractor.GetTablesID() {
-				regionsInfo, err := e.getRegionsInfoForSingleTable(tikvHelper, tableID)
+				regionsInfo, err := e.getRegionsInfoForTable(tikvHelper, is, tableID)
 				if err != nil {
 					return err
 				}
@@ -1693,8 +1694,7 @@ func (e *memtableRetriever) setDataForTiKVRegionStatus(ctx sessionctx.Context) (
 			return err
 		}
 	}
-	allSchemas := ctx.GetInfoSchema().(infoschema.InfoSchema).AllSchemas()
-	tableInfos := tikvHelper.GetRegionsTableInfo(allRegionsInfo, allSchemas)
+	tableInfos := tikvHelper.GetRegionsTableInfo(allRegionsInfo, is.AllSchemas())
 	for i := range allRegionsInfo.Regions {
 		tableList := tableInfos[allRegionsInfo.Regions[i].ID]
 		if len(tableList) == 0 {
@@ -1707,6 +1707,32 @@ func (e *memtableRetriever) setDataForTiKVRegionStatus(ctx sessionctx.Context) (
 	return nil
 }
 
+func (e *memtableRetriever) getRegionsInfoForTable(h *helper.Helper, is infoschema.InfoSchema, tableID int64) (*helper.RegionsInfo, error) {
+	tbl, _ := is.TableByID(tableID)
+	if tbl == nil {
+		return nil, infoschema.ErrTableExists.GenWithStackByArgs(tableID)
+	}
+
+	pt := tbl.Meta().GetPartitionInfo()
+	if pt == nil {
+		regionsInfo, err := e.getRegionsInfoForSingleTable(h, tableID)
+		if err != nil {
+			return nil, err
+		}
+		return regionsInfo, nil
+	}
+
+	allRegionsInfo := helper.NewRegionsInfo()
+	for _, def := range pt.Definitions {
+		regionsInfo, err := e.getRegionsInfoForSingleTable(h, def.ID)
+		if err != nil {
+			return nil, err
+		}
+		allRegionsInfo = allRegionsInfo.Merge(regionsInfo)
+	}
+	return allRegionsInfo, nil
+}
+
 func (e *memtableRetriever) getRegionsInfoForSingleTable(helper *helper.Helper, tableID int64) (*helper.RegionsInfo, error) {
 	sk, ek := tablecodec.GetTableHandleKeyRange(tableID)
 	sRegion, err := helper.GetRegionByKey(codec.EncodeBytes(nil, sk))

planner/core/memtable_predicate_extractor_test.go

@@ -16,8 +16,10 @@ package core_test
 
 import (
 	"context"
+	"fmt"
 	"regexp"
 	"sort"
+	"strconv"
 	"testing"
 	"time"
 
@@ -1702,6 +1704,20 @@ func TestTikvRegionStatusExtractor(t *testing.T) {
 	se, err := session.CreateSession4Test(store)
 	require.NoError(t, err)
 
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec(`CREATE TABLE p (id int(11), unique index(id))
+PARTITION BY RANGE COLUMNS ( id ) (
+		PARTITION p0 VALUES LESS THAN (6),
+		PARTITION p1 VALUES LESS THAN (11),
+		PARTITION p3 VALUES LESS THAN (21)
+)`)
+	res := tk.MustQuery("select * from information_schema.tables where table_name = 'p'")
+	idStr := res.Rows()[0][21]
+	id, err := strconv.Atoi(idStr.(string))
+	require.NoError(t, err)
+	sSQL := fmt.Sprintf("select * from information_schema.TIKV_REGION_STATUS where table_id = %v", id)
+
 	var cases = []struct {
 		sql      string
 		tableIDs []int64
@@ -1718,6 +1734,10 @@ func TestTikvRegionStatusExtractor(t *testing.T) {
 			sql:      "select * from information_schema.TIKV_REGION_STATUS where table_id in (1,2,3)",
 			tableIDs: []int64{1, 2, 3},
 		},
+		{
+			sql:      sSQL,
+			tableIDs: []int64{int64(id)},
+		},
 	}
 	parser := parser.New()
 	for _, ca := range cases {