tools: isolate keyspace according to the given label

client/http/interface.go

@@ -83,6 +83,7 @@ type Client interface {
 	SetPlacementRuleInBatch(context.Context, []*RuleOp) error
 	SetPlacementRuleBundles(context.Context, []*GroupBundle, bool) error
 	DeletePlacementRule(context.Context, string, string) error
+	DeletePlacementRuleBundleByGroup(context.Context, string) error
 	GetAllPlacementRuleGroups(context.Context) ([]*RuleGroup, error)
 	GetPlacementRuleGroupByID(context.Context, string) (*RuleGroup, error)
 	SetPlacementRuleGroup(context.Context, *RuleGroup) error
@@ -727,6 +728,14 @@ func (c *client) DeletePlacementRule(ctx context.Context, group, id string) erro
 		WithMethod(http.MethodDelete))
 }
 
+// DeletePlacementRuleBundleByGroup deletes the placement rule bundle by group.
+func (c *client) DeletePlacementRuleBundleByGroup(ctx context.Context, group string) error {
+	return c.request(ctx, newRequestInfo().
+		WithName(deletePlacementRuleBundleByGroupName).
+		WithURI(PlacementRuleBundleByGroup(group)).
+		WithMethod(http.MethodDelete))
+}
+
 // GetAllPlacementRuleGroups gets all placement rule groups.
 func (c *client) GetAllPlacementRuleGroups(ctx context.Context) ([]*RuleGroup, error) {
 	var ruleGroups []*RuleGroup

client/http/request_info.go

@@ -70,6 +70,7 @@ const (
 	setPlacementRuleInBatchName             = "SetPlacementRuleInBatch"
 	setPlacementRuleBundlesName             = "SetPlacementRuleBundles"
 	deletePlacementRuleName                 = "DeletePlacementRule"
+	deletePlacementRuleBundleByGroupName    = "DeletePlacementRuleBundleByGroup"
 	getAllPlacementRuleGroupsName           = "GetAllPlacementRuleGroups"
 	getPlacementRuleGroupByIDName           = "GetPlacementRuleGroupByID"
 	setPlacementRuleGroupName               = "SetPlacementRuleGroup"

tests/server/keyspace/keyspace_test.go

@@ -16,7 +16,6 @@ package keyspace
 
 import (
 	"context"
-	"encoding/binary"
 	"encoding/hex"
 	"fmt"
 	"strconv"
@@ -30,7 +29,6 @@ import (
 	"github.com/pingcap/failpoint"
 	"github.com/pingcap/kvproto/pkg/keyspacepb"
 
-	"github.com/tikv/pd/pkg/codec"
 	"github.com/tikv/pd/pkg/keyspace"
 	"github.com/tikv/pd/pkg/keyspace/constant"
 	"github.com/tikv/pd/pkg/schedule/labeler"
@@ -113,19 +111,12 @@ func checkLabelRule(re *require.Assertions, id uint32, regionLabeler *labeler.Re
 	re.True(ok)
 	re.Len(rangeRule, 2)
 
-	keyspaceIDBytes := make([]byte, 4)
-	nextKeyspaceIDBytes := make([]byte, 4)
-	binary.BigEndian.PutUint32(keyspaceIDBytes, id)
-	binary.BigEndian.PutUint32(nextKeyspaceIDBytes, id+1)
-	rawLeftBound := hex.EncodeToString(codec.EncodeBytes(append([]byte{'r'}, keyspaceIDBytes[1:]...)))
-	rawRightBound := hex.EncodeToString(codec.EncodeBytes(append([]byte{'r'}, nextKeyspaceIDBytes[1:]...)))
-	txnLeftBound := hex.EncodeToString(codec.EncodeBytes(append([]byte{'x'}, keyspaceIDBytes[1:]...)))
-	txnRightBound := hex.EncodeToString(codec.EncodeBytes(append([]byte{'x'}, nextKeyspaceIDBytes[1:]...)))
-
-	re.Equal(rawLeftBound, rangeRule[0].StartKeyHex)
-	re.Equal(rawRightBound, rangeRule[0].EndKeyHex)
-	re.Equal(txnLeftBound, rangeRule[1].StartKeyHex)
-	re.Equal(txnRightBound, rangeRule[1].EndKeyHex)
+	bound := keyspace.MakeRegionBound(id)
+
+	re.Equal(hex.EncodeToString(bound.RawLeftBound), rangeRule[0].StartKeyHex)
+	re.Equal(hex.EncodeToString(bound.RawRightBound), rangeRule[0].EndKeyHex)
+	re.Equal(hex.EncodeToString(bound.TxnLeftBound), rangeRule[1].StartKeyHex)
+	re.Equal(hex.EncodeToString(bound.TxnRightBound), rangeRule[1].EndKeyHex)
 }
 
 func (suite *keyspaceTestSuite) TestPreAlloc() {

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

@@ -27,6 +27,7 @@ import (
 	"github.com/spf13/cobra"
 
 	"github.com/tikv/pd/client/constants"
+	pd "github.com/tikv/pd/client/http"
 	"github.com/tikv/pd/pkg/keyspace"
 	"github.com/tikv/pd/server/apiv2/handlers"
 )
@@ -45,15 +46,18 @@ const (
 // NewKeyspaceCommand returns a keyspace subcommand of rootCmd.
 func NewKeyspaceCommand() *cobra.Command {
 	cmd := &cobra.Command{
-		Use:   "keyspace <command> [flags]",
-		Short: "keyspace commands",
+		Use:               "keyspace <command> [flags]",
+		Short:             "keyspace commands",
+		PersistentPreRunE: requirePDClient,
 	}
 	cmd.AddCommand(newShowKeyspaceCommand())
 	cmd.AddCommand(newCreateKeyspaceCommand())
 	cmd.AddCommand(newUpdateKeyspaceConfigCommand())
 	cmd.AddCommand(newUpdateKeyspaceStateCommand())
 	cmd.AddCommand(newListKeyspaceCommand())
 	cmd.AddCommand(newShowKeyspaceRangeCommand())
+	cmd.AddCommand(newSetPlacementCommand())
+	cmd.AddCommand(newRevertPlacementCommand())
 	return cmd
 }
 
@@ -444,3 +448,150 @@ func showKeyspaceRangeByNameCommandFunc(cmd *cobra.Command, args []string) {
 
 	cmd.Println(string(output))
 }
+
+func newSetPlacementCommand() *cobra.Command {
+	r := &cobra.Command{
+		Use:   "set-placement <keyspace-id> <label-key>=<label-value> [<label-key>=<label-value>...]",
+		Short: "set keyspace placement rules with store label constraints",
+		Long: "Set placement rules for all regions of a keyspace to stores matching the specified labels.\n" +
+			"This creates a placement rule bundle that places the keyspace's regions on stores matching all the label constraints.\n" +
+			"Examples:\n" +
+			"  pd-ctl keyspace set-placement 1 zone=east\n" +
+			"  pd-ctl keyspace set-placement 1 zone=east disk=ssd",
+		Run: setPlacementCommandFunc,
+	}
+	return r
+}
+
+func setPlacementCommandFunc(cmd *cobra.Command, args []string) {
+	if len(args) < 2 {
+		cmd.Usage()
+		return
+	}
+
+	keyspaceIDStr := args[0]
+	labelPairs := args[1:]
+
+	// Parse keyspace ID
+	var keyspaceID uint64
+	if _, err := fmt.Sscanf(keyspaceIDStr, "%d", &keyspaceID); err != nil {
+		cmd.PrintErrf("Invalid keyspace ID: %v\n", err)
+		return
+	}
+	keyspaceID32 := uint32(keyspaceID)
+	if keyspaceID32 < constants.DefaultKeyspaceID || keyspaceID32 > constants.MaxKeyspaceID {
+		cmd.PrintErrf("Invalid keyspace ID %d. It must be in the range of [%d, %d]\n",
+			keyspaceID, constants.DefaultKeyspaceID, constants.MaxKeyspaceID)
+		return
+	}
+
+	// Parse all label key=value pairs
+	var labelConstraints []pd.LabelConstraint
+	for _, labelPair := range labelPairs {
+		parts := strings.SplitN(labelPair, "=", 2)
+		if len(parts) != 2 {
+			cmd.PrintErrf("Invalid label format, expected <label-key>=<label-value>, got: %s\n", labelPair)
+			return
+		}
+		labelKey := strings.TrimSpace(parts[0])
+		labelValue := strings.TrimSpace(parts[1])
+		if labelKey == "" || labelValue == "" {
+			cmd.PrintErrf("Label key and value cannot be empty\n")
+			return
+		}
+		labelConstraints = append(labelConstraints, pd.LabelConstraint{
+			Key:    labelKey,
+			Op:     pd.In,
+			Values: []string{labelValue},
+		})
+	}
+
+	// Generate key ranges for the keyspace
+	keyRanges := keyspace.MakeKeyRanges(keyspaceID32)
+
+	// Create placement rule bundle
+	groupID := fmt.Sprintf("keyspace-%d", keyspaceID)
+	bundle := &pd.GroupBundle{
+		ID:       groupID,
+		Index:    100,
+		Override: true,
+		Rules: []*pd.Rule{
+			{
+				GroupID: groupID,
+				ID:      fmt.Sprintf("keyspace-%d-rule", keyspaceID),
+				Role:    pd.Voter,
+				// TODO: make replica count configurable
+				Count:            3,
+				StartKeyHex:      keyRanges[0].(map[string]any)["start_key"].(string),
+				EndKeyHex:        keyRanges[0].(map[string]any)["end_key"].(string),
+				LabelConstraints: labelConstraints,
+			},
+			{
+				GroupID: groupID,
+				ID:      fmt.Sprintf("keyspace-%d-rule-txn", keyspaceID),
+				Role:    pd.Voter,
+				// TODO: make replica count configurable
+				Count:            3,
+				StartKeyHex:      keyRanges[1].(map[string]any)["start_key"].(string),
+				EndKeyHex:        keyRanges[1].(map[string]any)["end_key"].(string),
+				LabelConstraints: labelConstraints,
+			},
+		},
+	}
+
+	// Send to PD server using PDCli
+	// Use partial=true to allow this bundle to coexist with other bundles (like TiFlash rules)
+	if err := PDCli.SetPlacementRuleBundles(cmd.Context(), []*pd.GroupBundle{bundle}, true); err != nil {
+		cmd.PrintErrf("Failed to set placement rule: %v\n", err)
+		return
+	}
+
+	// Build label description for output
+	var labelDescs []string
+	for _, lc := range labelConstraints {
+		labelDescs = append(labelDescs, fmt.Sprintf("%s=%s", lc.Key, strings.Join(lc.Values, ",")))
+	}
+	cmd.Printf("Successfully set placement rules for keyspace %d with label constraints: %s\n", keyspaceID, strings.Join(labelDescs, ", "))
+}
+
+func newRevertPlacementCommand() *cobra.Command {
+	r := &cobra.Command{
+		Use:   "revert-placement <keyspace-id>",
+		Short: "revert keyspace placement rules",
+		Long: "Remove placement rules for a keyspace.\n" +
+			"This deletes the placement rule bundle that was created by set-placement command.",
+		Run: revertPlacementCommandFunc,
+	}
+	return r
+}
+
+func revertPlacementCommandFunc(cmd *cobra.Command, args []string) {
+	if len(args) != 1 {
+		cmd.Usage()
+		return
+	}
+
+	keyspaceIDStr := args[0]
+
+	// Parse keyspace ID
+	var keyspaceID uint64
+	if _, err := fmt.Sscanf(keyspaceIDStr, "%d", &keyspaceID); err != nil {
+		cmd.PrintErrf("Invalid keyspace ID: %v\n", err)
+		return
+	}
+	keyspaceID32 := uint32(keyspaceID)
+	if keyspaceID32 < constants.DefaultKeyspaceID || keyspaceID32 > constants.MaxKeyspaceID {
+		cmd.PrintErrf("Invalid keyspace ID %d. It must be in the range of [%d, %d]\n",
+			keyspaceID, constants.DefaultKeyspaceID, constants.MaxKeyspaceID)
+		return
+	}
+
+	// Delete the placement rule bundle using PDCli
+	groupID := fmt.Sprintf("keyspace-%d", keyspaceID)
+	if err := PDCli.DeletePlacementRuleBundleByGroup(cmd.Context(), groupID); err != nil {
+		cmd.PrintErrf("Failed to revert placement rule: %v\n", err)
+		return
+	}
+
+	cmd.Printf("Successfully reverted placement rules for keyspace %d\n", keyspaceID)
+}

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

@@ -0,0 +1,75 @@
+// Copyright 2025 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 (
+	"encoding/binary"
+	"encoding/hex"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/tikv/pd/pkg/codec"
+	"github.com/tikv/pd/pkg/keyspace"
+)
+
+func TestMakeKeyRanges(t *testing.T) {
+	re := require.New(t)
+
+	testCases := []struct {
+		keyspaceID uint32
+	}{
+		{keyspaceID: 0},
+		{keyspaceID: 1},
+		{keyspaceID: 10},
+		{keyspaceID: 100},
+		{keyspaceID: 16777215}, // max valid keyspace ID (2^24 - 1)
+	}
+
+	for _, tc := range testCases {
+		ranges := keyspace.MakeKeyRanges(tc.keyspaceID)
+		re.Len(ranges, 2, "should have 2 ranges (raw and txn)")
+
+		// Verify raw key range
+		rawRange := ranges[0].(map[string]any)
+		rawStart := rawRange["start_key"].(string)
+		rawEnd := rawRange["end_key"].(string)
+		re.NotEmpty(rawStart)
+		re.NotEmpty(rawEnd)
+
+		// Verify txn key range
+		txnRange := ranges[1].(map[string]any)
+		txnStart := txnRange["start_key"].(string)
+		txnEnd := txnRange["end_key"].(string)
+		re.NotEmpty(txnStart)
+		re.NotEmpty(txnEnd)
+
+		// Verify the encoding matches the expected format
+		keyspaceIDBytes := make([]byte, 4)
+		nextKeyspaceIDBytes := make([]byte, 4)
+		binary.BigEndian.PutUint32(keyspaceIDBytes, tc.keyspaceID)
+		binary.BigEndian.PutUint32(nextKeyspaceIDBytes, tc.keyspaceID+1)
+
+		expectedRawLeft := codec.EncodeBytes(append([]byte{'r'}, keyspaceIDBytes[1:]...))
+		expectedRawRight := codec.EncodeBytes(append([]byte{'r'}, nextKeyspaceIDBytes[1:]...))
+		expectedTxnLeft := codec.EncodeBytes(append([]byte{'x'}, keyspaceIDBytes[1:]...))
+		expectedTxnRight := codec.EncodeBytes(append([]byte{'x'}, nextKeyspaceIDBytes[1:]...))
+
+		re.Equal(hex.EncodeToString(expectedRawLeft), rawStart)
+		re.Equal(hex.EncodeToString(expectedRawRight), rawEnd)
+		re.Equal(hex.EncodeToString(expectedTxnLeft), txnStart)
+		re.Equal(hex.EncodeToString(expectedTxnRight), txnEnd)
+	}
+}