Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
56 changes: 51 additions & 5 deletions tools/pd-ctl/pdctl/command/region_command.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@ package command
import (
"bytes"
"context"
"encoding/binary"
"encoding/hex"
"encoding/json"
"fmt"
Expand All @@ -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"
Expand Down Expand Up @@ -535,28 +537,62 @@ func NewRegionWithKeyspaceCommand() *cobra.Command {
Short: "show region information of the given keyspace",
}
r.AddCommand(&cobra.Command{
Use: "id <keyspace_id> <limit>",
Short: "show region information for the given keyspace id",
Use: "id <keyspace_id> [table-id <table_id>] [<limit>]",
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)
Expand All @@ -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.`
Expand Down
219 changes: 219 additions & 0 deletions tools/pd-ctl/pdctl/command/region_command_test.go
Original file line number Diff line number Diff line change
@@ -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
}
Loading