tools: add show region by keyspace-id or table-id

tools/pd-ctl/pdctl/command/region_command.go

@@ -17,6 +17,7 @@
 import (
 	"bytes"
 	"context"
+	"encoding/binary"
 	"encoding/hex"
 	"encoding/json"
 	"fmt"
@@ -535,28 +536,55 @@
 		Short: "show region information of the given keyspace",
 	}
 	r.AddCommand(&cobra.Command{
-		Use:   "id <keyspace_id> <limit>",
+		Use:   "id <keyspace_id> [table-id <table_id>] [<limit>]",
 		Short: "show region information for the given keyspace id",
 		Run:   showRegionWithKeyspaceCommandFunc,
 	})
 	return r
 }
 
 func showRegionWithKeyspaceCommandFunc(cmd *cobra.Command, args []string) {
-	if len(args) < 1 || len(args) > 2 {
+	if len(args) < 1 || len(args) > 4 {
 		cmd.Println(cmd.UsageString())
 		return
 	}
 
-	keyspaceID := args[0]
-	prefix := regionsKeyspacePrefix + "/id/" + keyspaceID
+	if _, err := strconv.ParseUint(args[0], 10, 32); err != nil {
+		cmd.Println("keyspace_id should be a number")
+		return
+	}
+
+	prefix := regionsKeyspacePrefix + "/id/" + args[0]
 	if len(args) == 2 {
 		if _, err := strconv.Atoi(args[1]); err != nil {
 			cmd.Println("limit should be a number")
 			return
 		}
 		prefix += "?limit=" + args[1]
+	} else if len(args) >= 3 {
+		if args[1] != "table-id" {
+			cmd.Println("the second argument should be table-id")
+			return
+		}
+		tableID, err := strconv.ParseInt(args[2], 10, 64)
+		if err != nil {
+			cmd.Println("table-id should be a number")
+			return
+		}
+		query := make(url.Values)
+		startKey, endKey := makeTableRangeInKeyspace(args[0], tableID)
+		query.Set("key", url.QueryEscape(string(startKey)))
+		query.Set("end_key", url.QueryEscape(string(endKey)))
+		if len(args) == 4 {
+			if _, err := strconv.Atoi(args[3]); err != nil {
+				cmd.Println("limit should be a number")
+				return
+			}
+			query.Set("limit", args[3])
+		}
+		prefix = regionsKeyPrefix + "?" + query.Encode()
 	}
+
 	r, err := doRequest(cmd, prefix, http.MethodGet, http.Header{})
 	if err != nil {
 		cmd.Printf("Failed to get regions with the given keyspace: %s\n", err)
@@ -565,6 +593,17 @@
 	cmd.Println(r)
 }
 
+func makeTableRangeInKeyspace(keyspaceID string, tableID int64) ([]byte, []byte) {
+	keyspaceIDUint64, _ := strconv.ParseUint(keyspaceID, 10, 32)
+	keyspaceIDBytes := make([]byte, 4)
+	binary.BigEndian.PutUint32(keyspaceIDBytes, uint32(keyspaceIDUint64))
+
+	keyPrefix := append([]byte{'x'}, keyspaceIDBytes[1:]...)
+	startKey := codec.EncodeBytes(nil, append(keyPrefix, append([]byte{'t'}, codec.EncodeInt(nil, tableID)...)...))
+	endKey := codec.EncodeBytes(nil, append(keyPrefix, append([]byte{'t'}, codec.EncodeInt(nil, tableID+1)...)...))
+	return startKey, endKey
+}
+
 const (
 	rangeHolesLongDesc = `There are some cases that the region range is not continuous, for example, the region doesn't send the heartbeat to PD after a splitting.
 This command will output all empty ranges without any region info.`

tools/pd-ctl/pdctl/command/region_command_test.go

@@ -0,0 +1,205 @@
+// Copyright 2025 TiKV Project Authors.
+//
+// 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 command
+
+import (
+	"bytes"
+	"encoding/binary"
+	"io"
+	"net/http"
+	"net/url"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/tikv/pd/tools/pd-ctl/helper/tidb/codec"
+)
+
+type captureRegionRoundTripper struct {
+	path     string
+	rawQuery string
+	body     string
+	err      error
+}
+
+func (m *captureRegionRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
+	m.path = req.URL.Path
+	m.rawQuery = req.URL.RawQuery
+	if m.err != nil {
+		return nil, m.err
+	}
+	return &http.Response{
+		StatusCode: http.StatusOK,
+		Body:       io.NopCloser(strings.NewReader(m.body)),
+	}, nil
+}
+
+func TestRegionKeyspaceIDPath(t *testing.T) {
+	re := require.New(t)
+	rt := &captureRegionRoundTripper{body: `{"ok":true}`}
+	oldClient := dialClient
+	dialClient = &http.Client{Transport: rt}
+	defer func() { dialClient = oldClient }()
+
+	cmd := NewRegionWithKeyspaceCommand()
+	cmd.PersistentFlags().String("pd", "http://mock-pd:2379", "")
+	cmd.SetArgs([]string{"id", "1"})
+	var out bytes.Buffer
+	cmd.SetOut(&out)
+	cmd.SetErr(&out)
+
+	re.NoError(cmd.Execute())
+	re.Equal("/pd/api/v1/regions/keyspace/id/1", rt.path)
+	re.Empty(rt.rawQuery)
+	re.Contains(out.String(), `{"ok":true}`)
+}
+
+func TestRegionKeyspaceIDTableIDPath(t *testing.T) {
+	re := require.New(t)
+	rt := &captureRegionRoundTripper{body: `{"ok":true}`}
+	oldClient := dialClient
+	dialClient = &http.Client{Transport: rt}
+	defer func() { dialClient = oldClient }()
+
+	startKey, endKey := expectedTableRangeInKeyspace(1, 100)
+	query := make(url.Values)
+	query.Set("key", url.QueryEscape(string(startKey)))
+	query.Set("end_key", url.QueryEscape(string(endKey)))
+
+	cmd := NewRegionWithKeyspaceCommand()
+	cmd.PersistentFlags().String("pd", "http://mock-pd:2379", "")
+	cmd.SetArgs([]string{"id", "1", "table-id", "100"})
+	var out bytes.Buffer
+	cmd.SetOut(&out)
+	cmd.SetErr(&out)
+
+	re.NoError(cmd.Execute())
+	re.Equal("/pd/api/v1/regions/key", rt.path)
+	re.Equal(query.Encode(), rt.rawQuery)
+	re.Contains(out.String(), `{"ok":true}`)
+}
+
+func TestRegionKeyspaceIDPathWithLimit(t *testing.T) {
+	re := require.New(t)
+	rt := &captureRegionRoundTripper{body: `{"ok":true}`}
+	oldClient := dialClient
+	dialClient = &http.Client{Transport: rt}
+	defer func() { dialClient = oldClient }()
+
+	cmd := NewRegionWithKeyspaceCommand()
+	cmd.PersistentFlags().String("pd", "http://mock-pd:2379", "")
+	cmd.SetArgs([]string{"id", "1", "16"})
+	var out bytes.Buffer
+	cmd.SetOut(&out)
+	cmd.SetErr(&out)
+
+	re.NoError(cmd.Execute())
+	re.Equal("/pd/api/v1/regions/keyspace/id/1", rt.path)
+	re.Equal("limit=16", rt.rawQuery)
+	re.Contains(out.String(), `{"ok":true}`)
+}
+
+func TestRegionKeyspaceIDTableIDPathWithLimit(t *testing.T) {
+	re := require.New(t)
+	rt := &captureRegionRoundTripper{body: `{"ok":true}`}
+	oldClient := dialClient
+	dialClient = &http.Client{Transport: rt}
+	defer func() { dialClient = oldClient }()
+
+	startKey, endKey := expectedTableRangeInKeyspace(1, 100)
+	query := make(url.Values)
+	query.Set("key", url.QueryEscape(string(startKey)))
+	query.Set("end_key", url.QueryEscape(string(endKey)))
+	query.Set("limit", "16")
+
+	cmd := NewRegionWithKeyspaceCommand()
+	cmd.PersistentFlags().String("pd", "http://mock-pd:2379", "")
+	cmd.SetArgs([]string{"id", "1", "table-id", "100", "16"})
+	var out bytes.Buffer
+	cmd.SetOut(&out)
+	cmd.SetErr(&out)
+
+	re.NoError(cmd.Execute())
+	re.Equal("/pd/api/v1/regions/key", rt.path)
+	re.Equal(query.Encode(), rt.rawQuery)
+	re.Contains(out.String(), `{"ok":true}`)
+}
+
+func TestRegionKeyspaceIDInvalidKeyspaceID(t *testing.T) {
+	re := require.New(t)
+
+	cmd := NewRegionWithKeyspaceCommand()
+	cmd.PersistentFlags().String("pd", "http://mock-pd:2379", "")
+	cmd.SetArgs([]string{"id", "invalid"})
+	var out bytes.Buffer
+	cmd.SetOut(&out)
+	cmd.SetErr(&out)
+
+	re.NoError(cmd.Execute())
+	re.Contains(out.String(), "keyspace_id should be a number")
+}
+
+func TestRegionKeyspaceIDInvalidTableID(t *testing.T) {
+	re := require.New(t)
+
+	cmd := NewRegionWithKeyspaceCommand()
+	cmd.PersistentFlags().String("pd", "http://mock-pd:2379", "")
+	cmd.SetArgs([]string{"id", "1", "table-id", "invalid"})
+	var out bytes.Buffer
+	cmd.SetOut(&out)
+	cmd.SetErr(&out)
+
+	re.NoError(cmd.Execute())
+	re.Contains(out.String(), "table-id should be a number")
+}
+
+func TestRegionKeyspaceIDInvalidLimit(t *testing.T) {
+	re := require.New(t)
+
+	cmd := NewRegionWithKeyspaceCommand()
+	cmd.PersistentFlags().String("pd", "http://mock-pd:2379", "")
+	cmd.SetArgs([]string{"id", "1", "table-id", "100", "invalid"})
+	var out bytes.Buffer
+	cmd.SetOut(&out)
+	cmd.SetErr(&out)
+
+	re.NoError(cmd.Execute())
+	re.Contains(out.String(), "limit should be a number")
+}
+
+func TestRegionKeyspaceIDWrongTableIDLiteral(t *testing.T) {
+	re := require.New(t)
+
+	cmd := NewRegionWithKeyspaceCommand()
+	cmd.PersistentFlags().String("pd", "http://mock-pd:2379", "")
+	cmd.SetArgs([]string{"id", "1", "table", "100"})
+	var out bytes.Buffer
+	cmd.SetOut(&out)
+	cmd.SetErr(&out)
+
+	re.NoError(cmd.Execute())
+	re.Contains(out.String(), "the second argument should be table-id")
+}
+
+func expectedTableRangeInKeyspace(keyspaceID uint32, tableID int64) ([]byte, []byte) {
+	keyspaceIDBytes := make([]byte, 4)
+	binary.BigEndian.PutUint32(keyspaceIDBytes, keyspaceID)
+
+	keyPrefix := append([]byte{'x'}, keyspaceIDBytes[1:]...)
+	startKey := codec.EncodeBytes(nil, append(keyPrefix, append([]byte{'t'}, codec.EncodeInt(nil, tableID)...)...))
+	endKey := codec.EncodeBytes(nil, append(keyPrefix, append([]byte{'t'}, codec.EncodeInt(nil, tableID+1)...)...))
+	return startKey, endKey
+}