pd-ctl: split operator support use key (#6600)

close #6601 Signed-off-by: bufferflies <1045931706@qq.com>
tikv · Jun 19, 2023 · 80edddb · 80edddb
1 parent 06c6b52
commit 80edddb
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 8 deletions.
diff --git a/tests/pdctl/operator/operator_test.go b/tests/pdctl/operator/operator_test.go
@@ -16,6 +16,7 @@ package operator_test
 
 import (
 	"context"
+	"encoding/hex"
 	"strconv"
 	"strings"
 	"testing"
@@ -87,7 +88,7 @@ func TestOperator(t *testing.T) {
 		{Id: 1, StoreId: 1},
 		{Id: 2, StoreId: 2},
 	}))
-	pdctl.MustPutRegion(re, cluster, 3, 2, []byte("b"), []byte("c"), core.SetPeers([]*metapb.Peer{
+	pdctl.MustPutRegion(re, cluster, 3, 2, []byte("b"), []byte("d"), core.SetPeers([]*metapb.Peer{
 		{Id: 3, StoreId: 1},
 		{Id: 4, StoreId: 2},
 	}))
@@ -135,33 +136,41 @@ func TestOperator(t *testing.T) {
 			reset:  []string{"-u", pdAddr, "operator", "remove", "1"},
 		},
 		{
-			// operator add split-region <region_id> [--policy=scan|approximate]
+			// operator add split-region <region_id> [--policy=scan|approximate|usekey] [--keys=xxx(xxx is hex encoded string)]
 			cmd:    []string{"-u", pdAddr, "operator", "add", "split-region", "3", "--policy=scan"},
 			show:   []string{"-u", pdAddr, "operator", "show"},
 			expect: "split region with policy SCAN",
 			reset:  []string{"-u", pdAddr, "operator", "remove", "3"},
 		},
 		{
-			// operator add split-region <region_id> [--policy=scan|approximate]
+			// operator add split-region <region_id> [--policy=scan|approximate|usekey] [--keys=xxx(xxx is hex encoded string)]
 			cmd:    []string{"-u", pdAddr, "operator", "add", "split-region", "3", "--policy=approximate"},
 			show:   []string{"-u", pdAddr, "operator", "show"},
 			expect: "split region with policy APPROXIMATE",
 			reset:  []string{"-u", pdAddr, "operator", "remove", "3"},
 		},
 		{
-			// operator add split-region <region_id> [--policy=scan|approximate]
+			// operator add split-region <region_id> [--policy=scan|approximate|usekey] [--keys=xxx(xxx is hex encoded string)]
 			cmd:    []string{"-u", pdAddr, "operator", "add", "split-region", "3", "--policy=scan"},
 			show:   []string{"-u", pdAddr, "operator", "check", "3"},
 			expect: "split region with policy SCAN",
 			reset:  []string{"-u", pdAddr, "operator", "remove", "3"},
 		},
 		{
-			// operator add split-region <region_id> [--policy=scan|approximate]
+			// operator add split-region <region_id> [--policy=scan|approximate|usekey] [--keys=xxx(xxx is hex encoded string)]
 			cmd:    []string{"-u", pdAddr, "operator", "add", "split-region", "3", "--policy=approximate"},
 			show:   []string{"-u", pdAddr, "operator", "check", "3"},
 			expect: "status: RUNNING",
 			reset:  []string{"-u", pdAddr, "operator", "remove", "3"},
 		},
+		{
+			// operator add split-region <region_id> [--policy=scan|approximate|usekey] [--keys=xxx(xxx is hex encoded string)]
+			cmd: []string{"-u", pdAddr, "operator", "add", "split-region", "3", "--policy=usekey",
+				"--keys=" + hex.EncodeToString([]byte("c"))},
+			show:   []string{"-u", pdAddr, "operator", "show"},
+			expect: "split: region 3 use policy USEKEY and keys [" + hex.EncodeToString([]byte("c")) + "]",
+			reset:  []string{"-u", pdAddr, "operator", "remove", "3"},
+		},
 	}
 
 	for _, testCase := range testCases {

diff --git a/tools/pd-ctl/pdctl/command/operator.go b/tools/pd-ctl/pdctl/command/operator.go
@@ -351,16 +351,17 @@ func removePeerCommandFunc(cmd *cobra.Command, args []string) {
 // NewSplitRegionCommand returns a command to split a region.
 func NewSplitRegionCommand() *cobra.Command {
 	c := &cobra.Command{
-		Use:   "split-region <region_id> [--policy=scan|approximate]",
+		Use:   "split-region <region_id> [--policy=scan|approximate|usekey] [--keys]",
 		Short: "split a region",
 		Run:   splitRegionCommandFunc,
 	}
 	c.Flags().String("policy", "scan", "the policy to get region split key")
+	c.Flags().String("keys", "", "the split key, hex encoded")
 	return c
 }
 
 func splitRegionCommandFunc(cmd *cobra.Command, args []string) {
-	if len(args) != 1 {
+	if len(args) < 1 {
 		cmd.Println(cmd.UsageString())
 		return
 	}
@@ -373,7 +374,7 @@ func splitRegionCommandFunc(cmd *cobra.Command, args []string) {
 
 	policy := cmd.Flags().Lookup("policy").Value.String()
 	switch policy {
-	case "scan", "approximate":
+	case "scan", "approximate", "usekey":
 		break
 	default:
 		cmd.Println("Error: unknown policy")
@@ -384,6 +385,10 @@ func splitRegionCommandFunc(cmd *cobra.Command, args []string) {
 	input["name"] = cmd.Name()
 	input["region_id"] = ids[0]
 	input["policy"] = policy
+	keys := cmd.Flags().Lookup("keys").Value.String()
+	if len(keys) > 0 {
+		input["keys"] = []string{keys}
+	}
 	postJSON(cmd, operatorsPrefix, input)
 }