diff --git a/tools/pd-ctl/pdctl/command/region_command.go b/tools/pd-ctl/pdctl/command/region_command.go index 5d8e89309ee..7eaa5a8886d 100644 --- a/tools/pd-ctl/pdctl/command/region_command.go +++ b/tools/pd-ctl/pdctl/command/region_command.go @@ -17,6 +17,7 @@ package command import ( "bytes" "context" + "encoding/binary" "encoding/hex" "encoding/json" "fmt" @@ -36,6 +37,7 @@ import ( "github.com/pingcap/failpoint" "github.com/tikv/pd/client/clients/router" + "github.com/tikv/pd/client/constants" pd "github.com/tikv/pd/client/http" "github.com/tikv/pd/pkg/utils/typeutil" "github.com/tikv/pd/tools/pd-ctl/helper/mok" @@ -535,28 +537,62 @@ func NewRegionWithKeyspaceCommand() *cobra.Command { Short: "show region information of the given keyspace", } r.AddCommand(&cobra.Command{ - Use: "id ", - Short: "show region information for the given keyspace id", + Use: "id [table-id ] []", + Short: "show region information for the given keyspace id, optionally filtered by table 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 + keyspaceIDUint64, err := strconv.ParseUint(args[0], 10, 32) + if err != nil { + cmd.Println("keyspace_id should be a number") + return + } + keyspaceID := uint32(keyspaceIDUint64) + if keyspaceID < constants.DefaultKeyspaceID || keyspaceID > constants.MaxKeyspaceID { + cmd.Printf("invalid keyspace id %d. It must be in the range of [%d, %d]\n", + keyspaceID, constants.DefaultKeyspaceID, constants.MaxKeyspaceID) + 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(keyspaceID, tableID) + query.Set("key", string(startKey)) + query.Set("end_key", 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 +601,16 @@ func showRegionWithKeyspaceCommandFunc(cmd *cobra.Command, args []string) { cmd.Println(r) } +func makeTableRangeInKeyspace(keyspaceID uint32, tableID int64) (startKey, endKey []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 +} + 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.` diff --git a/tools/pd-ctl/pdctl/command/region_command_test.go b/tools/pd-ctl/pdctl/command/region_command_test.go new file mode 100644 index 00000000000..e1557ec2213 --- /dev/null +++ b/tools/pd-ctl/pdctl/command/region_command_test.go @@ -0,0 +1,219 @@ +// Copyright 2026 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", string(startKey)) + query.Set("end_key", 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", string(startKey)) + query.Set("end_key", 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 TestRegionKeyspaceIDOutOfRangeKeyspaceID(t *testing.T) { + re := require.New(t) + + cmd := NewRegionWithKeyspaceCommand() + cmd.PersistentFlags().String("pd", "http://mock-pd:2379", "") + cmd.SetArgs([]string{"id", "16777216"}) + var out bytes.Buffer + cmd.SetOut(&out) + cmd.SetErr(&out) + + re.NoError(cmd.Execute()) + re.Contains(out.String(), "invalid keyspace id 16777216. It must be in the range of [0, 16777215]") +} + +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) (startKey, endKey []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 +}