ddl,executor: fix usages of getting regions from pd client (#63371)

pkg/ddl/index.go

@@ -73,7 +73,6 @@ import (
 	"github.com/pingcap/tidb/pkg/util"
 	"github.com/pingcap/tidb/pkg/util/backoff"
 	"github.com/pingcap/tidb/pkg/util/chunk"
-	"github.com/pingcap/tidb/pkg/util/codec"
 	"github.com/pingcap/tidb/pkg/util/dbterror"
 	"github.com/pingcap/tidb/pkg/util/generatedexpr"
 	"github.com/pingcap/tidb/pkg/util/intest"
@@ -2962,25 +2961,10 @@ func estimateRowSizeFromRegion(ctx context.Context, store kv.Storage, tbl table.
 	}
 	pid := tbl.Meta().ID
 	sk, ek := tablecodec.GetTableHandleKeyRange(pid)
-	sRegion, err := pdCli.GetRegionByKey(ctx, codec.EncodeBytes(nil, sk))
-	if err != nil {
-		return 0, err
-	}
-	eRegion, err := pdCli.GetRegionByKey(ctx, codec.EncodeBytes(nil, ek))
-	if err != nil {
-		return 0, err
-	}
-	sk, err = hex.DecodeString(sRegion.StartKey)
-	if err != nil {
-		return 0, err
-	}
-	ek, err = hex.DecodeString(eRegion.EndKey)
-	if err != nil {
-		return 0, err
-	}
+	start, end := hStore.GetCodec().EncodeRegionRange(sk, ek)
 	// We use the second region to prevent the influence of the front and back tables.
 	regionLimit := 3
-	regionInfos, err := pdCli.GetRegionsByKeyRange(ctx, pdHttp.NewKeyRange(sk, ek), regionLimit)
+	regionInfos, err := pdCli.GetRegionsByKeyRange(ctx, pdHttp.NewKeyRange(start, end), regionLimit)
 	if err != nil {
 		return 0, err
 	}

pkg/ddl/reorg_util.go

@@ -34,7 +34,6 @@ import (
 	"github.com/pingcap/tidb/pkg/store/helper"
 	"github.com/pingcap/tidb/pkg/table"
 	"github.com/pingcap/tidb/pkg/tablecodec"
-	"github.com/pingcap/tidb/pkg/util/codec"
 	"github.com/pingcap/tidb/pkg/util/dbterror"
 	pdhttp "github.com/tikv/pd/client/http"
 	"go.uber.org/zap"
@@ -172,7 +171,7 @@ func getTableSizeByID(ctx context.Context, store kv.Storage, tbl table.Table) in
 	}
 	var totalSize int64
 	for _, pid := range pids {
-		size, err := estimateTableSizeByID(ctx, pdCli, pid)
+		size, err := estimateTableSizeByID(ctx, pdCli, helperStore, pid)
 		if err != nil {
 			logutil.DDLLogger().Warn("failed to estimate table size for calculating concurrency",
 				zap.Int64("physicalID", pid),
@@ -193,28 +192,12 @@ func getTableSizeByID(ctx context.Context, store kv.Storage, tbl table.Table) in
 	return totalSize
 }
 
-func estimateTableSizeByID(ctx context.Context, pdCli pdhttp.Client, pid int64) (int64, error) {
+func estimateTableSizeByID(ctx context.Context, pdCli pdhttp.Client, store helper.Storage, pid int64) (int64, error) {
 	sk, ek := tablecodec.GetTableHandleKeyRange(pid)
-	sRegion, err := pdCli.GetRegionByKey(ctx, codec.EncodeBytes(nil, sk))
-	if err != nil {
-		return 0, err
-	}
-	eRegion, err := pdCli.GetRegionByKey(ctx, codec.EncodeBytes(nil, ek))
-	if err != nil {
-		return 0, err
-	}
-	sk, err = hex.DecodeString(sRegion.StartKey)
-	if err != nil {
-		return 0, err
-	}
-	ek, err = hex.DecodeString(eRegion.EndKey)
-	if err != nil {
-		return 0, err
-	}
-
+	start, end := store.GetCodec().EncodeRegionRange(sk, ek)
 	var totalSize int64
 	for {
-		regionInfos, err := pdCli.GetRegionsByKeyRange(ctx, pdhttp.NewKeyRange(sk, ek), 128)
+		regionInfos, err := pdCli.GetRegionsByKeyRange(ctx, pdhttp.NewKeyRange(start, end), 128)
 		if err != nil {
 			return 0, err
 		}
@@ -225,11 +208,11 @@ func estimateTableSizeByID(ctx context.Context, pdCli pdhttp.Client, pid int64)
 			totalSize += r.ApproximateSize * units.MiB
 		}
 		lastKey := regionInfos.Regions[len(regionInfos.Regions)-1].EndKey
-		sk, err = hex.DecodeString(lastKey)
+		start, err = hex.DecodeString(lastKey)
 		if err != nil {
 			return 0, err
 		}
-		if bytes.Compare(sk, ek) >= 0 {
+		if bytes.Compare(start, end) >= 0 {
 			break
 		}
 	}

pkg/executor/infoschema_reader.go

@@ -2137,23 +2137,8 @@ func (*memtableRetriever) getRegionsInfoForSingleTable(ctx context.Context, help
 		return nil, err
 	}
 	sk, ek := tablecodec.GetTableHandleKeyRange(tableID)
-	sRegion, err := pdCli.GetRegionByKey(ctx, codec.EncodeBytes(nil, sk))
-	if err != nil {
-		return nil, err
-	}
-	eRegion, err := pdCli.GetRegionByKey(ctx, codec.EncodeBytes(nil, ek))
-	if err != nil {
-		return nil, err
-	}
-	sk, err = hex.DecodeString(sRegion.StartKey)
-	if err != nil {
-		return nil, err
-	}
-	ek, err = hex.DecodeString(eRegion.EndKey)
-	if err != nil {
-		return nil, err
-	}
-	return pdCli.GetRegionsByKeyRange(ctx, pd.NewKeyRange(sk, ek), -1)
+	start, end := helper.Store.GetCodec().EncodeRegionRange(sk, ek)
+	return pdCli.GetRegionsByKeyRange(ctx, pd.NewKeyRange(start, end), -1)
 }
 
 func (e *memtableRetriever) setNewTiKVRegionStatusCol(region *pd.RegionInfo, table *helper.TableInfo) {

tests/realtikvtest/sessiontest/BUILD.bazel

@@ -4,6 +4,7 @@ go_test(
     name = "sessiontest_test",
     timeout = "moderate",
     srcs = [
+        "infoschema_test.go",
         "infoschema_v2_test.go",
         "main_test.go",
         "paging_test.go",

tests/realtikvtest/sessiontest/infoschema_test.go

@@ -0,0 +1,45 @@
+// Copyright 2025 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package sessiontest
+
+import (
+	"context"
+	"fmt"
+	"testing"
+
+	"github.com/pingcap/tidb/pkg/parser/ast"
+	"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)
+	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)
+
+	showRegions := tk.MustQuery("show table t regions").Rows()
+	t.Log(showRegions)
+	require.Equal(t, 1, len(showRegions), showRegions)
+	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])
+	tk.MustExec("drop table t;")
+}