pd-ctl, tests: add the keyspace group commands

server/apiv2/handlers/tso_keyspace_group.go

@@ -70,6 +70,10 @@ func CreateKeyspaceGroups(c *gin.Context) {
 
 	svr := c.MustGet(middlewares.ServerContextKey).(*server.Server)
 	manager := svr.GetKeyspaceGroupManager()
+	if manager == nil {
+		c.AbortWithStatusJSON(http.StatusInternalServerError, "keyspace group manager is not initialized")
+		return
+	}
 	err = manager.CreateKeyspaceGroups(createParams.KeyspaceGroups)
 	if err != nil {
 		c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error())
@@ -89,6 +93,10 @@ func GetKeyspaceGroups(c *gin.Context) {
 
 	svr := c.MustGet(middlewares.ServerContextKey).(*server.Server)
 	manager := svr.GetKeyspaceGroupManager()
+	if manager == nil {
+		c.AbortWithStatusJSON(http.StatusInternalServerError, "keyspace group manager is not initialized")
+		return
+	}
 	keyspaceGroups, err := manager.GetKeyspaceGroups(scanStart, scanLimit)
 	if err != nil {
 		c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error())
@@ -108,6 +116,10 @@ func GetKeyspaceGroupByID(c *gin.Context) {
 
 	svr := c.MustGet(middlewares.ServerContextKey).(*server.Server)
 	manager := svr.GetKeyspaceGroupManager()
+	if manager == nil {
+		c.AbortWithStatusJSON(http.StatusInternalServerError, "keyspace group manager is not initialized")
+		return
+	}
 	kg, err := manager.GetKeyspaceGroupByID(id)
 	if err != nil {
 		c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error())
@@ -127,6 +139,10 @@ func DeleteKeyspaceGroupByID(c *gin.Context) {
 
 	svr := c.MustGet(middlewares.ServerContextKey).(*server.Server)
 	manager := svr.GetKeyspaceGroupManager()
+	if manager == nil {
+		c.AbortWithStatusJSON(http.StatusInternalServerError, "keyspace group manager is not initialized")
+		return
+	}
 	kg, err := manager.DeleteKeyspaceGroupByID(id)
 	if err != nil {
 		c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error())
@@ -174,6 +190,10 @@ func SplitKeyspaceGroupByID(c *gin.Context) {
 	if !patrolKeyspaceAssignmentState.patrolled {
 		// Patrol keyspace assignment before splitting keyspace group.
 		manager := svr.GetKeyspaceManager()
+		if manager == nil {
+			c.AbortWithStatusJSON(http.StatusInternalServerError, "keyspace manager is not initialized")
+			return
+		}
 		err = manager.PatrolKeyspaceAssignment()
 		if err != nil {
 			c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error())
@@ -185,6 +205,10 @@ func SplitKeyspaceGroupByID(c *gin.Context) {
 	patrolKeyspaceAssignmentState.Unlock()
 	// Split keyspace group.
 	groupManager := svr.GetKeyspaceGroupManager()
+	if groupManager == nil {
+		c.AbortWithStatusJSON(http.StatusInternalServerError, "keyspace group manager is not initialized")
+		return
+	}
 	err = groupManager.SplitKeyspaceGroupByID(id, splitParams.NewID, splitParams.Keyspaces)
 	if err != nil {
 		c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error())
@@ -225,6 +249,10 @@ func AllocNodesForKeyspaceGroup(c *gin.Context) {
 	}
 	svr := c.MustGet(middlewares.ServerContextKey).(*server.Server)
 	manager := svr.GetKeyspaceGroupManager()
+	if manager == nil {
+		c.AbortWithStatusJSON(http.StatusInternalServerError, "keyspace group manager is not initialized")
+		return
+	}
 	allocParams := &AllocNodesForKeyspaceGroupParams{}
 	err = c.BindJSON(allocParams)
 	if err != nil {
@@ -268,6 +296,10 @@ func SetNodesForKeyspaceGroup(c *gin.Context) {
 	}
 	svr := c.MustGet(middlewares.ServerContextKey).(*server.Server)
 	manager := svr.GetKeyspaceGroupManager()
+	if manager == nil {
+		c.AbortWithStatusJSON(http.StatusInternalServerError, "keyspace group manager is not initialized")
+		return
+	}
 	setParams := &SetNodesForKeyspaceGroupParams{}
 	err = c.BindJSON(setParams)
 	if err != nil {

tests/pdctl/keyspace/keyspace_group_test.go

@@ -0,0 +1,84 @@
+// Copyright 2023 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 keyspace_test
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"github.com/tikv/pd/pkg/mcs/utils"
+	"github.com/tikv/pd/pkg/storage/endpoint"
+	"github.com/tikv/pd/server/apiv2/handlers"
+	"github.com/tikv/pd/tests"
+	"github.com/tikv/pd/tests/pdctl"
+	handlersutil "github.com/tikv/pd/tests/server/apiv2/handlers"
+	pdctlCmd "github.com/tikv/pd/tools/pd-ctl/pdctl"
+)
+
+func TestKeyspaceGroup(t *testing.T) {
+	re := require.New(t)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	tc, err := tests.NewTestAPICluster(ctx, 1)
+	re.NoError(err)
+	err = tc.RunInitialServers()
+	re.NoError(err)
+	tc.WaitLeader()
+	leaderServer := tc.GetServer(tc.GetLeader())
+	re.NoError(leaderServer.BootstrapCluster())
+	pdAddr := tc.GetConfig().GetClientURL()
+	cmd := pdctlCmd.GetRootCmd()
+
+	// Show keyspace group information.
+	defaultKeyspaceGroupID := fmt.Sprintf("%d", utils.DefaultKeyspaceGroupID)
+	args := []string{"-u", pdAddr, "keyspace-group"}
+	output, err := pdctl.ExecuteCommand(cmd, append(args, defaultKeyspaceGroupID)...)
+	re.NoError(err)
+	var keyspaceGroup endpoint.KeyspaceGroup
+	err = json.Unmarshal(output, &keyspaceGroup)
+	re.NoError(err)
+	re.Equal(utils.DefaultKeyspaceGroupID, keyspaceGroup.ID)
+	re.Contains(keyspaceGroup.Keyspaces, utils.DefaultKeyspaceID)
+	// Split keyspace group.
+	handlersutil.MustCreateKeyspaceGroup(re, leaderServer, &handlers.CreateKeyspaceGroupParams{
+		KeyspaceGroups: []*endpoint.KeyspaceGroup{
+			{
+				ID:        1,
+				UserKind:  endpoint.Standard.String(),
+				Members:   make([]endpoint.KeyspaceGroupMember, utils.KeyspaceGroupDefaultReplicaCount),
+				Keyspaces: []uint32{111, 222, 333},
+			},
+		},
+	})
+	_, err = pdctl.ExecuteCommand(cmd, append(args, "split", "1", "2", "222", "333")...)
+	re.NoError(err)
+	output, err = pdctl.ExecuteCommand(cmd, append(args, "1")...)
+	re.NoError(err)
+	keyspaceGroup = endpoint.KeyspaceGroup{}
+	err = json.Unmarshal(output, &keyspaceGroup)
+	re.NoError(err)
+	re.Equal(uint32(1), keyspaceGroup.ID)
+	re.Equal(keyspaceGroup.Keyspaces, []uint32{111})
+	output, err = pdctl.ExecuteCommand(cmd, append(args, "2")...)
+	re.NoError(err)
+	keyspaceGroup = endpoint.KeyspaceGroup{}
+	err = json.Unmarshal(output, &keyspaceGroup)
+	re.NoError(err)
+	re.Equal(uint32(2), keyspaceGroup.ID)
+	re.Equal(keyspaceGroup.Keyspaces, []uint32{222, 333})
+}

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

@@ -20,8 +20,10 @@ import (
 	"github.com/spf13/cobra"
 )
 
-const clusterPrefix = "pd/api/v1/cluster"
-const clusterStatusPrefix = "pd/api/v1/cluster/status"
+const (
+	clusterPrefix       = "pd/api/v1/cluster"
+	clusterStatusPrefix = "pd/api/v1/cluster/status"
+)
 
 // NewClusterCommand return a cluster subcommand of rootCmd
 func NewClusterCommand() *cobra.Command {

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

@@ -0,0 +1,83 @@
+// Copyright 2023 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 (
+	"fmt"
+	"net/http"
+	"strconv"
+
+	"github.com/spf13/cobra"
+)
+
+const keyspaceGroupsPrefix = "pd/api/v2/tso/keyspace-groups"
+
+// NewKeyspaceGroupCommand return a keyspace group subcommand of rootCmd
+func NewKeyspaceGroupCommand() *cobra.Command {
+	cmd := &cobra.Command{
+		Use:   "keyspace-group <keyspace_group_id>",
+		Short: "show keyspace group information with the given ID",
+		Run:   showKeyspaceGroupCommandFunc,
+	}
+	cmd.AddCommand(newSplitKeyspaceGroupCommand())
+	return cmd
+}
+
+func newSplitKeyspaceGroupCommand() *cobra.Command {
+	r := &cobra.Command{
+		Use:   "split <keyspace_group_id> <new_keyspace_group_id> [<keyspace_id>]",
+		Short: "split the keyspace group with the given ID and transfer the keyspaces into the newly split one",
+		Run:   splitKeyspaceGroupCommandFunc,
+	}
+	return r
+}
+
+func showKeyspaceGroupCommandFunc(cmd *cobra.Command, args []string) {
+	if len(args) < 1 {
+		cmd.Usage()
+		return
+	}
+	r, err := doRequest(cmd, fmt.Sprintf("%s/%s", keyspaceGroupsPrefix, args[0]), http.MethodGet, http.Header{})
+	if err != nil {
+		cmd.Printf("Failed to get the keyspace groups information: %s\n", err)
+		return
+	}
+	cmd.Println(r)
+}
+
+func splitKeyspaceGroupCommandFunc(cmd *cobra.Command, args []string) {
+	if len(args) < 3 {
+		cmd.Usage()
+		return
+	}
+	newID, err := strconv.ParseUint(args[1], 10, 32)
+	if err != nil {
+		cmd.Printf("Failed to parse the new keyspace group ID: %s\n", err)
+		return
+	}
+	keyspaces := make([]uint32, 0, len(args)-2)
+	for _, arg := range args[2:] {
+		id, err := strconv.ParseUint(arg, 10, 32)
+		if err != nil {
+			cmd.Printf("Failed to parse the keyspace ID: %s\n", err)
+			return
+		}
+		keyspaces = append(keyspaces, uint32(id))
+	}
+	postJSON(cmd, fmt.Sprintf("%s/%s/split", keyspaceGroupsPrefix, args[0]), map[string]interface{}{
+		"new-id":    uint32(newID),
+		"keyspaces": keyspaces,
+	})
+}

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

@@ -21,7 +21,7 @@ import (
 	"github.com/tikv/pd/pkg/utils/tsoutil"
 )
 
-// NewTSOCommand return a ping subcommand of rootCmd
+// NewTSOCommand return a TSO subcommand of rootCmd
 func NewTSOCommand() *cobra.Command {
 	cmd := &cobra.Command{
 		Use:   "tso <timestamp>",

tools/pd-ctl/pdctl/ctl.go

@@ -65,6 +65,7 @@ func GetRootCmd() *cobra.Command {
 		command.NewMinResolvedTSCommand(),
 		command.NewCompletionCommand(),
 		command.NewUnsafeCommand(),
+		command.NewKeyspaceGroupCommand(),
 	)
 
 	rootCmd.Flags().ParseErrorsWhitelist.UnknownFlags = true