Skip to content

Commit 6248525

Browse files
rleungxrleungx
committed
pd-ctl: extend region keyspace lookup to table ranges
Reuse the existing region keyspace command instead of adding a parallel entry point for table-scoped lookups. Add request-shape tests so the keyspace and table range queries stay aligned with the current CLI behavior. Signed-off-by: Ryan Leung <rleungx@gmail.com>
1 parent 3eb99ae commit 6248525

File tree

2 files changed

+248
-4
lines changed

2 files changed

+248
-4
lines changed

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

Lines changed: 43 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ package command
1717
import (
1818
"bytes"
1919
"context"
20+
"encoding/binary"
2021
"encoding/hex"
2122
"encoding/json"
2223
"fmt"
@@ -535,28 +536,55 @@ func NewRegionWithKeyspaceCommand() *cobra.Command {
535536
Short: "show region information of the given keyspace",
536537
}
537538
r.AddCommand(&cobra.Command{
538-
Use: "id <keyspace_id> <limit>",
539+
Use: "id <keyspace_id> [table-id <table_id>] [<limit>]",
539540
Short: "show region information for the given keyspace id",
540541
Run: showRegionWithKeyspaceCommandFunc,
541542
})
542543
return r
543544
}
544545

545546
func showRegionWithKeyspaceCommandFunc(cmd *cobra.Command, args []string) {
546-
if len(args) < 1 || len(args) > 2 {
547+
if len(args) < 1 || len(args) > 4 {
547548
cmd.Println(cmd.UsageString())
548549
return
549550
}
550551

551-
keyspaceID := args[0]
552-
prefix := regionsKeyspacePrefix + "/id/" + keyspaceID
552+
if _, err := strconv.ParseUint(args[0], 10, 32); err != nil {
553+
cmd.Println("keyspace_id should be a number")
554+
return
555+
}
556+
557+
prefix := regionsKeyspacePrefix + "/id/" + args[0]
553558
if len(args) == 2 {
554559
if _, err := strconv.Atoi(args[1]); err != nil {
555560
cmd.Println("limit should be a number")
556561
return
557562
}
558563
prefix += "?limit=" + args[1]
564+
} else if len(args) >= 3 {
565+
if args[1] != "table-id" {
566+
cmd.Println("the second argument should be table-id")
567+
return
568+
}
569+
tableID, err := strconv.ParseInt(args[2], 10, 64)
570+
if err != nil {
571+
cmd.Println("table-id should be a number")
572+
return
573+
}
574+
query := make(url.Values)
575+
startKey, endKey := makeTableRangeInKeyspace(args[0], tableID)
576+
query.Set("key", url.QueryEscape(string(startKey)))
577+
query.Set("end_key", url.QueryEscape(string(endKey)))
578+
if len(args) == 4 {
579+
if _, err := strconv.Atoi(args[3]); err != nil {
580+
cmd.Println("limit should be a number")
581+
return
582+
}
583+
query.Set("limit", args[3])
584+
}
585+
prefix = regionsKeyPrefix + "?" + query.Encode()
559586
}
587+
560588
r, err := doRequest(cmd, prefix, http.MethodGet, http.Header{})
561589
if err != nil {
562590
cmd.Printf("Failed to get regions with the given keyspace: %s\n", err)
@@ -565,6 +593,17 @@ func showRegionWithKeyspaceCommandFunc(cmd *cobra.Command, args []string) {
565593
cmd.Println(r)
566594
}
567595

596+
func makeTableRangeInKeyspace(keyspaceID string, tableID int64) ([]byte, []byte) {
597+
keyspaceIDUint64, _ := strconv.ParseUint(keyspaceID, 10, 32)
598+
keyspaceIDBytes := make([]byte, 4)
599+
binary.BigEndian.PutUint32(keyspaceIDBytes, uint32(keyspaceIDUint64))
600+
601+
keyPrefix := append([]byte{'x'}, keyspaceIDBytes[1:]...)
602+
startKey := codec.EncodeBytes(nil, append(keyPrefix, append([]byte{'t'}, codec.EncodeInt(nil, tableID)...)...))
603+
endKey := codec.EncodeBytes(nil, append(keyPrefix, append([]byte{'t'}, codec.EncodeInt(nil, tableID+1)...)...))
604+
return startKey, endKey
605+
}
606+
568607
const (
569608
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.
570609
This command will output all empty ranges without any region info.`

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

Lines changed: 205 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,205 @@
1+
// Copyright 2025 TiKV Project Authors.
2+
//
3+
// Licensed under the Apache License, Version 2.0 (the "License");
4+
// you may not use this file except in compliance with the License.
5+
// You may obtain a copy of the License at
6+
//
7+
// http://www.apache.org/licenses/LICENSE-2.0
8+
//
9+
// Unless required by applicable law or agreed to in writing, software
10+
// distributed under the License is distributed on an "AS IS" BASIS,
11+
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+
// See the License for the specific language governing permissions and
13+
// limitations under the License.
14+
15+
package command
16+
17+
import (
18+
"bytes"
19+
"encoding/binary"
20+
"io"
21+
"net/http"
22+
"net/url"
23+
"strings"
24+
"testing"
25+
26+
"github.com/stretchr/testify/require"
27+
28+
"github.com/tikv/pd/tools/pd-ctl/helper/tidb/codec"
29+
)
30+
31+
type captureRegionRoundTripper struct {
32+
path string
33+
rawQuery string
34+
body string
35+
err error
36+
}
37+
38+
func (m *captureRegionRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
39+
m.path = req.URL.Path
40+
m.rawQuery = req.URL.RawQuery
41+
if m.err != nil {
42+
return nil, m.err
43+
}
44+
return &http.Response{
45+
StatusCode: http.StatusOK,
46+
Body: io.NopCloser(strings.NewReader(m.body)),
47+
}, nil
48+
}
49+
50+
func TestRegionKeyspaceIDPath(t *testing.T) {
51+
re := require.New(t)
52+
rt := &captureRegionRoundTripper{body: `{"ok":true}`}
53+
oldClient := dialClient
54+
dialClient = &http.Client{Transport: rt}
55+
defer func() { dialClient = oldClient }()
56+
57+
cmd := NewRegionWithKeyspaceCommand()
58+
cmd.PersistentFlags().String("pd", "http://mock-pd:2379", "")
59+
cmd.SetArgs([]string{"id", "1"})
60+
var out bytes.Buffer
61+
cmd.SetOut(&out)
62+
cmd.SetErr(&out)
63+
64+
re.NoError(cmd.Execute())
65+
re.Equal("/pd/api/v1/regions/keyspace/id/1", rt.path)
66+
re.Empty(rt.rawQuery)
67+
re.Contains(out.String(), `{"ok":true}`)
68+
}
69+
70+
func TestRegionKeyspaceIDTableIDPath(t *testing.T) {
71+
re := require.New(t)
72+
rt := &captureRegionRoundTripper{body: `{"ok":true}`}
73+
oldClient := dialClient
74+
dialClient = &http.Client{Transport: rt}
75+
defer func() { dialClient = oldClient }()
76+
77+
startKey, endKey := expectedTableRangeInKeyspace(1, 100)
78+
query := make(url.Values)
79+
query.Set("key", url.QueryEscape(string(startKey)))
80+
query.Set("end_key", url.QueryEscape(string(endKey)))
81+
82+
cmd := NewRegionWithKeyspaceCommand()
83+
cmd.PersistentFlags().String("pd", "http://mock-pd:2379", "")
84+
cmd.SetArgs([]string{"id", "1", "table-id", "100"})
85+
var out bytes.Buffer
86+
cmd.SetOut(&out)
87+
cmd.SetErr(&out)
88+
89+
re.NoError(cmd.Execute())
90+
re.Equal("/pd/api/v1/regions/key", rt.path)
91+
re.Equal(query.Encode(), rt.rawQuery)
92+
re.Contains(out.String(), `{"ok":true}`)
93+
}
94+
95+
func TestRegionKeyspaceIDPathWithLimit(t *testing.T) {
96+
re := require.New(t)
97+
rt := &captureRegionRoundTripper{body: `{"ok":true}`}
98+
oldClient := dialClient
99+
dialClient = &http.Client{Transport: rt}
100+
defer func() { dialClient = oldClient }()
101+
102+
cmd := NewRegionWithKeyspaceCommand()
103+
cmd.PersistentFlags().String("pd", "http://mock-pd:2379", "")
104+
cmd.SetArgs([]string{"id", "1", "16"})
105+
var out bytes.Buffer
106+
cmd.SetOut(&out)
107+
cmd.SetErr(&out)
108+
109+
re.NoError(cmd.Execute())
110+
re.Equal("/pd/api/v1/regions/keyspace/id/1", rt.path)
111+
re.Equal("limit=16", rt.rawQuery)
112+
re.Contains(out.String(), `{"ok":true}`)
113+
}
114+
115+
func TestRegionKeyspaceIDTableIDPathWithLimit(t *testing.T) {
116+
re := require.New(t)
117+
rt := &captureRegionRoundTripper{body: `{"ok":true}`}
118+
oldClient := dialClient
119+
dialClient = &http.Client{Transport: rt}
120+
defer func() { dialClient = oldClient }()
121+
122+
startKey, endKey := expectedTableRangeInKeyspace(1, 100)
123+
query := make(url.Values)
124+
query.Set("key", url.QueryEscape(string(startKey)))
125+
query.Set("end_key", url.QueryEscape(string(endKey)))
126+
query.Set("limit", "16")
127+
128+
cmd := NewRegionWithKeyspaceCommand()
129+
cmd.PersistentFlags().String("pd", "http://mock-pd:2379", "")
130+
cmd.SetArgs([]string{"id", "1", "table-id", "100", "16"})
131+
var out bytes.Buffer
132+
cmd.SetOut(&out)
133+
cmd.SetErr(&out)
134+
135+
re.NoError(cmd.Execute())
136+
re.Equal("/pd/api/v1/regions/key", rt.path)
137+
re.Equal(query.Encode(), rt.rawQuery)
138+
re.Contains(out.String(), `{"ok":true}`)
139+
}
140+
141+
func TestRegionKeyspaceIDInvalidKeyspaceID(t *testing.T) {
142+
re := require.New(t)
143+
144+
cmd := NewRegionWithKeyspaceCommand()
145+
cmd.PersistentFlags().String("pd", "http://mock-pd:2379", "")
146+
cmd.SetArgs([]string{"id", "invalid"})
147+
var out bytes.Buffer
148+
cmd.SetOut(&out)
149+
cmd.SetErr(&out)
150+
151+
re.NoError(cmd.Execute())
152+
re.Contains(out.String(), "keyspace_id should be a number")
153+
}
154+
155+
func TestRegionKeyspaceIDInvalidTableID(t *testing.T) {
156+
re := require.New(t)
157+
158+
cmd := NewRegionWithKeyspaceCommand()
159+
cmd.PersistentFlags().String("pd", "http://mock-pd:2379", "")
160+
cmd.SetArgs([]string{"id", "1", "table-id", "invalid"})
161+
var out bytes.Buffer
162+
cmd.SetOut(&out)
163+
cmd.SetErr(&out)
164+
165+
re.NoError(cmd.Execute())
166+
re.Contains(out.String(), "table-id should be a number")
167+
}
168+
169+
func TestRegionKeyspaceIDInvalidLimit(t *testing.T) {
170+
re := require.New(t)
171+
172+
cmd := NewRegionWithKeyspaceCommand()
173+
cmd.PersistentFlags().String("pd", "http://mock-pd:2379", "")
174+
cmd.SetArgs([]string{"id", "1", "table-id", "100", "invalid"})
175+
var out bytes.Buffer
176+
cmd.SetOut(&out)
177+
cmd.SetErr(&out)
178+
179+
re.NoError(cmd.Execute())
180+
re.Contains(out.String(), "limit should be a number")
181+
}
182+
183+
func TestRegionKeyspaceIDWrongTableIDLiteral(t *testing.T) {
184+
re := require.New(t)
185+
186+
cmd := NewRegionWithKeyspaceCommand()
187+
cmd.PersistentFlags().String("pd", "http://mock-pd:2379", "")
188+
cmd.SetArgs([]string{"id", "1", "table", "100"})
189+
var out bytes.Buffer
190+
cmd.SetOut(&out)
191+
cmd.SetErr(&out)
192+
193+
re.NoError(cmd.Execute())
194+
re.Contains(out.String(), "the second argument should be table-id")
195+
}
196+
197+
func expectedTableRangeInKeyspace(keyspaceID uint32, tableID int64) ([]byte, []byte) {
198+
keyspaceIDBytes := make([]byte, 4)
199+
binary.BigEndian.PutUint32(keyspaceIDBytes, keyspaceID)
200+
201+
keyPrefix := append([]byte{'x'}, keyspaceIDBytes[1:]...)
202+
startKey := codec.EncodeBytes(nil, append(keyPrefix, append([]byte{'t'}, codec.EncodeInt(nil, tableID)...)...))
203+
endKey := codec.EncodeBytes(nil, append(keyPrefix, append([]byte{'t'}, codec.EncodeInt(nil, tableID+1)...)...))
204+
return startKey, endKey
205+
}

0 commit comments

Comments
 (0)
⚡