add table configuration

Signed-off-by: 童剑 <1045931706@qq.com>

pkg/schedule/schedulers/balance_range.go

@@ -68,10 +68,11 @@ type balanceRangeSchedulerConfig struct {
 }
 
 type balanceRangeSchedulerParam struct {
-	Role    string          `json:"role"`
-	Engine  string          `json:"engine"`
-	Timeout time.Duration   `json:"timeout"`
-	Ranges  []core.KeyRange `json:"ranges"`
+	Role      string          `json:"role"`
+	Engine    string          `json:"engine"`
+	Timeout   time.Duration   `json:"timeout"`
+	Ranges    []core.KeyRange `json:"ranges"`
+	TableName string          `json:"table-name"`
 }
 
 func (conf *balanceRangeSchedulerConfig) clone() *balanceRangeSchedulerParam {
@@ -80,10 +81,11 @@ func (conf *balanceRangeSchedulerConfig) clone() *balanceRangeSchedulerParam {
 	ranges := make([]core.KeyRange, len(conf.Ranges))
 	copy(ranges, conf.Ranges)
 	return &balanceRangeSchedulerParam{
-		Ranges:  ranges,
-		Role:    conf.Role,
-		Engine:  conf.Engine,
-		Timeout: conf.Timeout,
+		Ranges:    ranges,
+		Role:      conf.Role,
+		Engine:    conf.Engine,
+		Timeout:   conf.Timeout,
+		TableName: conf.TableName,
 	}
 }
 

pkg/schedule/schedulers/init.go

@@ -575,14 +575,19 @@ func schedulersRegister() {
 			if err != nil {
 				return errs.ErrURLParse.Wrap(err)
 			}
-			ranges, err := getKeyRanges(args[3:])
+			tableName, err := url.QueryUnescape(args[3])
+			if err != nil {
+				return errs.ErrURLParse.Wrap(err)
+			}
+			ranges, err := getKeyRanges(args[4:])
 			if err != nil {
 				return err
 			}
 			conf.Ranges = ranges
 			conf.Engine = engine
 			conf.Role = role
 			conf.Timeout = duration
+			conf.TableName = tableName
 			return nil
 		}
 	})

server/api/scheduler.go

@@ -113,6 +113,7 @@ func (h *schedulerHandler) CreateScheduler(w http.ResponseWriter, r *http.Reques
 			h.r.JSON(w, http.StatusInternalServerError, err.Error())
 			return
 		}
+
 		defaultTimeout := "1h"
 		if err := apiutil.CollectStringOption("timeout", input, collector); err != nil {
 			if errors.ErrorEqual(err, errs.ErrOptionNotExist) {
@@ -123,12 +124,17 @@ func (h *schedulerHandler) CreateScheduler(w http.ResponseWriter, r *http.Reques
 			}
 		}
 
-		if err := apiutil.CollectEscapeStringOption("start_key", input, collector); err != nil {
+		if err := apiutil.CollectStringOption("table-name", input, collector); err != nil {
 			h.r.JSON(w, http.StatusInternalServerError, err.Error())
 			return
 		}
 
-		if err := apiutil.CollectEscapeStringOption("end_key", input, collector); err != nil {
+		if err := apiutil.CollectEscapeStringOption("start-key", input, collector); err != nil {
+			h.r.JSON(w, http.StatusInternalServerError, err.Error())
+			return
+		}
+
+		if err := apiutil.CollectEscapeStringOption("end-key", input, collector); err != nil {
 			h.r.JSON(w, http.StatusInternalServerError, err.Error())
 			return
 		}

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

@@ -378,7 +378,7 @@ func NewBalanceWitnessSchedulerCommand() *cobra.Command {
 // NewBalanceRangeSchedulerCommand returns a command to add a balance-range-scheduler.
 func NewBalanceRangeSchedulerCommand() *cobra.Command {
 	c := &cobra.Command{
-		Use:        "balance-range-scheduler [--format=raw|encode|hex] <engine> <role> <start_key> <end_key>",
+		Use:        "balance-range-scheduler [--format=raw|encode|hex] <engine> <role> <table-name> <start_key> <end_key>",
 		Short:      "add a scheduler to balance region for given range",
 		Run:        addSchedulerForBalanceRangeCommandFunc,
 		Deprecated: "balance-range will be deprecated in the future, please use sql instead",
@@ -426,16 +426,16 @@ func addSchedulerForGrantHotRegionCommandFunc(cmd *cobra.Command, args []string)
 }
 
 func addSchedulerForBalanceRangeCommandFunc(cmd *cobra.Command, args []string) {
-	if len(args) != 4 {
+	if len(args) != 5 {
 		cmd.Println(cmd.UsageString())
 		return
 	}
-	startKey, err := parseKey(cmd.Flags(), args[2])
+	startKey, err := parseKey(cmd.Flags(), args[3])
 	if err != nil {
 		cmd.Println("Error: ", err)
 		return
 	}
-	endKey, err := parseKey(cmd.Flags(), args[3])
+	endKey, err := parseKey(cmd.Flags(), args[4])
 	if err != nil {
 		cmd.Println("Error: ", err)
 		return
@@ -445,8 +445,9 @@ func addSchedulerForBalanceRangeCommandFunc(cmd *cobra.Command, args []string) {
 	input["name"] = cmd.Name()
 	input["engine"] = args[0]
 	input["role"] = args[1]
-	input["start_key"] = url.QueryEscape(startKey)
-	input["end_key"] = url.QueryEscape(endKey)
+	input["table-name"] = args[2]
+	input["start-key"] = url.QueryEscape(startKey)
+	input["end-key"] = url.QueryEscape(endKey)
 
 	postJSON(cmd, schedulersPrefix, input)
 }

tools/pd-ctl/tests/scheduler/scheduler_test.go

@@ -544,19 +544,19 @@ func (suite *schedulerTestSuite) checkSchedulerConfig(cluster *pdTests.TestClust
 	// test balance key range scheduler
 	echo = mustExec(re, cmd, []string{"-u", pdAddr, "scheduler", "add", "balance-range-scheduler"}, nil)
 	re.NotContains(echo, "Success!")
-	echo = mustExec(re, cmd, []string{"-u", pdAddr, "scheduler", "add", "balance-range-scheduler", "--format=raw", "tiflash", "learner", "a", "b"}, nil)
+	echo = mustExec(re, cmd, []string{"-u", pdAddr, "scheduler", "add", "balance-range-scheduler", "--format=raw", "tiflash", "learner", "test", "a", "b"}, nil)
 	re.Contains(echo, "Success!")
 	conf = make(map[string]any)
 	testutil.Eventually(re, func() bool {
 		mightExec(re, cmd, []string{"-u", pdAddr, "scheduler", "config", "balance-range-scheduler"}, &conf)
-		return conf["role"] == "learner" && conf["engine"] == "tiflash"
+		return conf["role"] == "learner" && conf["engine"] == "tiflash" && conf["table-name"] == "test"
 	})
 	re.Equal(float64(time.Hour.Nanoseconds()), conf["timeout"])
 	ranges := conf["ranges"].([]any)[0].(map[string]any)
 	re.Equal(core.HexRegionKeyStr([]byte("a")), ranges["start-key"])
 	re.Equal(core.HexRegionKeyStr([]byte("b")), ranges["end-key"])
 
-	echo = mustExec(re, cmd, []string{"-u", pdAddr, "scheduler", "add", "balance-range-scheduler", "--format=raw", "tiflash", "learner", "a", "b"}, nil)
+	echo = mustExec(re, cmd, []string{"-u", pdAddr, "scheduler", "add", "balance-range-scheduler", "--format=raw", "tiflash", "learner", "learner", "a", "b"}, nil)
 	re.Contains(echo, "400")
 	re.Contains(echo, "scheduler already exists")
 	echo = mustExec(re, cmd, []string{"-u", pdAddr, "scheduler", "remove", "balance-range-scheduler"}, nil)