From 2357ec41b9b7528819ab77c074af054e7971c0d4 Mon Sep 17 00:00:00 2001 From: David <8039876+AmoebaProtozoa@users.noreply.github.com> Date: Wed, 11 Oct 2023 21:32:48 -0500 Subject: [PATCH] pd-ctl: add keyspace commands (#7158) (#210) * pd-ctl: add keyspace commands (#7158) ref tikv/pd#4399 Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com> Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com> (cherry picked from commit 873212fc872270adcba3cb9bdd1f3a0131845415) * fix test Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com> --------- Signed-off-by: AmoebaProtozoa <8039876+AmoebaProtozoa@users.noreply.github.com> --- server/apiv2/handlers/keyspace.go | 6 +- tests/pdctl/keyspace/keyspace_test.go | 186 ++++++++++++- .../pd-ctl/pdctl/command/keyspace_command.go | 257 +++++++++++++++++- 3 files changed, 437 insertions(+), 12 deletions(-) diff --git a/server/apiv2/handlers/keyspace.go b/server/apiv2/handlers/keyspace.go index 5c2dd99bcfa..32270c42a2a 100644 --- a/server/apiv2/handlers/keyspace.go +++ b/server/apiv2/handlers/keyspace.go @@ -138,7 +138,7 @@ func LoadKeyspace(c *gin.Context) { // @Router /keyspaces/id/{id} [get] func LoadKeyspaceByID(c *gin.Context) { id, err := strconv.ParseUint(c.Param("id"), 10, 64) - if err != nil || id == 0 { + if err != nil { c.AbortWithStatusJSON(http.StatusInternalServerError, "invalid keyspace id") return } @@ -158,7 +158,7 @@ func LoadKeyspaceByID(c *gin.Context) { // parseLoadAllQuery parses LoadAllKeyspaces'/GetKeyspaceGroups' query parameters. // page_token: -// The keyspace/keyspace group id of the scan start. If not set, scan from keyspace/keyspace group with id 1. +// The keyspace/keyspace group id of the scan start. If not set, scan from keyspace/keyspace group with id 0. // It's string of ID of the previous scan result's last element (next_page_token). // limit: // The maximum number of keyspace metas/keyspace groups to return. If not set, no limit is posed. @@ -167,7 +167,7 @@ func LoadKeyspaceByID(c *gin.Context) { func parseLoadAllQuery(c *gin.Context) (scanStart uint32, scanLimit int, err error) { pageToken, set := c.GetQuery("page_token") if !set || pageToken == "" { - // If pageToken is empty or unset, then scan from ID of 1. + // If pageToken is empty or unset, then scan from ID of 0. scanStart = 0 } else { scanStart64, err := strconv.ParseUint(pageToken, 10, 32) diff --git a/tests/pdctl/keyspace/keyspace_test.go b/tests/pdctl/keyspace/keyspace_test.go index a0bab4114df..bf59185a41a 100644 --- a/tests/pdctl/keyspace/keyspace_test.go +++ b/tests/pdctl/keyspace/keyspace_test.go @@ -18,11 +18,14 @@ import ( "context" "encoding/json" "fmt" + "strconv" "strings" "testing" "github.com/pingcap/failpoint" + "github.com/pingcap/kvproto/pkg/keyspacepb" "github.com/stretchr/testify/require" + "github.com/stretchr/testify/suite" "github.com/tikv/pd/pkg/keyspace" "github.com/tikv/pd/pkg/mcs/utils" "github.com/tikv/pd/pkg/utils/testutil" @@ -65,7 +68,7 @@ func TestKeyspace(t *testing.T) { var k api.KeyspaceMeta keyspaceName := "keyspace_1" testutil.Eventually(re, func() bool { - args := []string{"-u", pdAddr, "keyspace", keyspaceName} + args := []string{"-u", pdAddr, "keyspace", "show", "name", keyspaceName} output, err := pdctl.ExecuteCommand(cmd, args...) re.NoError(err) re.NoError(json.Unmarshal(output, &k)) @@ -85,7 +88,7 @@ func TestKeyspace(t *testing.T) { // check keyspace group in keyspace whether changed. testutil.Eventually(re, func() bool { - args := []string{"-u", pdAddr, "keyspace", keyspaceName} + args := []string{"-u", pdAddr, "keyspace", "show", "name", keyspaceName} output, err := pdctl.ExecuteCommand(cmd, args...) re.NoError(err) re.NoError(json.Unmarshal(output, &k)) @@ -93,7 +96,7 @@ func TestKeyspace(t *testing.T) { }) // test error name - args := []string{"-u", pdAddr, "keyspace", "error_name"} + args := []string{"-u", pdAddr, "keyspace", "show", "name", "error_name"} output, err := pdctl.ExecuteCommand(cmd, args...) re.NoError(err) re.Contains(string(output), "Fail") @@ -101,3 +104,180 @@ func TestKeyspace(t *testing.T) { re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/tso/fastGroupSplitPatroller")) re.NoError(failpoint.Disable("github.com/tikv/pd/server/delayStartServerLoop")) } + +type keyspaceTestSuite struct { + suite.Suite + ctx context.Context + cancel context.CancelFunc + cluster *tests.TestCluster + pdAddr string +} + +func TestKeyspaceTestSuite(t *testing.T) { + suite.Run(t, new(keyspaceTestSuite)) +} + +func (suite *keyspaceTestSuite) SetupTest() { + suite.ctx, suite.cancel = context.WithCancel(context.Background()) + suite.NoError(failpoint.Enable("github.com/tikv/pd/server/delayStartServerLoop", `return(true)`)) + suite.NoError(failpoint.Enable("github.com/tikv/pd/pkg/keyspace/skipSplitRegion", "return(true)")) + tc, err := tests.NewTestAPICluster(suite.ctx, 1) + suite.NoError(err) + suite.NoError(tc.RunInitialServers()) + tc.WaitLeader() + leaderServer := tc.GetServer(tc.GetLeader()) + suite.NoError(leaderServer.BootstrapCluster()) + suite.cluster = tc + suite.pdAddr = tc.GetConfig().GetClientURL() +} + +func (suite *keyspaceTestSuite) TearDownTest() { + suite.NoError(failpoint.Disable("github.com/tikv/pd/server/delayStartServerLoop")) + suite.NoError(failpoint.Disable("github.com/tikv/pd/pkg/keyspace/skipSplitRegion")) + suite.cancel() +} + +func (suite *keyspaceTestSuite) TestShowKeyspace() { + re := suite.Require() + keyspaceName := "DEFAULT" + keyspaceID := uint32(0) + var k1, k2 api.KeyspaceMeta + // Show by name. + args := []string{"-u", suite.pdAddr, "keyspace", "show", "name", keyspaceName} + output, err := pdctl.ExecuteCommand(pdctlCmd.GetRootCmd(), args...) + re.NoError(err) + re.NoError(json.Unmarshal(output, &k1)) + re.Equal(keyspaceName, k1.GetName()) + re.Equal(keyspaceID, k1.GetId()) + // Show by ID. + args = []string{"-u", suite.pdAddr, "keyspace", "show", "id", strconv.Itoa(int(keyspaceID))} + output, err = pdctl.ExecuteCommand(pdctlCmd.GetRootCmd(), args...) + re.NoError(err) + re.NoError(json.Unmarshal(output, &k2)) + re.Equal(k1, k2) +} + +func mustCreateKeyspace(suite *keyspaceTestSuite, param api.CreateKeyspaceParams) api.KeyspaceMeta { + re := suite.Require() + var meta api.KeyspaceMeta + args := []string{"-u", suite.pdAddr, "keyspace", "create", param.Name} + for k, v := range param.Config { + args = append(args, "--config", fmt.Sprintf("%s=%s", k, v)) + } + output, err := pdctl.ExecuteCommand(pdctlCmd.GetRootCmd(), args...) + re.NoError(err) + re.NoError(json.Unmarshal(output, &meta)) + return meta +} + +func (suite *keyspaceTestSuite) TestCreateKeyspace() { + re := suite.Require() + param := api.CreateKeyspaceParams{ + Name: "test_keyspace", + Config: map[string]string{ + "foo": "bar", + "foo2": "bar2", + }, + } + meta := mustCreateKeyspace(suite, param) + re.Equal(param.Name, meta.GetName()) + for k, v := range param.Config { + re.Equal(v, meta.Config[k]) + } +} + +func (suite *keyspaceTestSuite) TestUpdateKeyspaceConfig() { + re := suite.Require() + param := api.CreateKeyspaceParams{ + Name: "test_keyspace", + Config: map[string]string{"foo": "1"}, + } + meta := mustCreateKeyspace(suite, param) + re.Equal("1", meta.Config["foo"]) + + // Update one existing config and add a new config, resulting in config: {foo: 2, foo2: 1}. + args := []string{"-u", suite.pdAddr, "keyspace", "update-config", param.Name, "--update", "foo=2,foo2=1"} + output, err := pdctl.ExecuteCommand(pdctlCmd.GetRootCmd(), args...) + re.NoError(err) + re.NoError(json.Unmarshal(output, &meta)) + re.Equal("test_keyspace", meta.GetName()) + re.Equal("2", meta.Config["foo"]) + re.Equal("1", meta.Config["foo2"]) + // Update one existing config and remove a config, resulting in config: {foo: 3}. + args = []string{"-u", suite.pdAddr, "keyspace", "update-config", param.Name, "--update", "foo=3", "--remove", "foo2"} + output, err = pdctl.ExecuteCommand(pdctlCmd.GetRootCmd(), args...) + re.NoError(err) + re.NoError(json.Unmarshal(output, &meta)) + re.Equal("test_keyspace", meta.GetName()) + re.Equal("3", meta.Config["foo"]) + re.NotContains(meta.GetConfig(), "foo2") + // Error if a key is specified in both --update and --remove list. + args = []string{"-u", suite.pdAddr, "keyspace", "update-config", param.Name, "--update", "foo=4", "--remove", "foo"} + output, err = pdctl.ExecuteCommand(pdctlCmd.GetRootCmd(), args...) + re.NoError(err) + re.Contains(string(output), "Fail") + // Error if a key is specified multiple times. + args = []string{"-u", suite.pdAddr, "keyspace", "update-config", param.Name, "--update", "foo=4,foo=5"} + output, err = pdctl.ExecuteCommand(pdctlCmd.GetRootCmd(), args...) + re.NoError(err) + re.Contains(string(output), "Fail") +} + +func (suite *keyspaceTestSuite) TestUpdateKeyspaceState() { + re := suite.Require() + param := api.CreateKeyspaceParams{ + Name: "test_keyspace", + } + meta := mustCreateKeyspace(suite, param) + re.Equal(keyspacepb.KeyspaceState_ENABLED, meta.State) + // Disable the keyspace, capitalization shouldn't matter. + args := []string{"-u", suite.pdAddr, "keyspace", "update-state", param.Name, "DiSAbleD"} + output, err := pdctl.ExecuteCommand(pdctlCmd.GetRootCmd(), args...) + re.NoError(err) + re.NoError(json.Unmarshal(output, &meta)) + re.Equal(keyspacepb.KeyspaceState_DISABLED, meta.State) + // Tombstone the keyspace without archiving should fail. + args = []string{"-u", suite.pdAddr, "keyspace", "update-state", param.Name, "TOMBSTONE"} + output, err = pdctl.ExecuteCommand(pdctlCmd.GetRootCmd(), args...) + re.NoError(err) + re.Contains(string(output), "Fail") +} + +func (suite *keyspaceTestSuite) TestListKeyspace() { + re := suite.Require() + var param api.CreateKeyspaceParams + for i := 0; i < 10; i++ { + param = api.CreateKeyspaceParams{ + Name: fmt.Sprintf("test_keyspace_%d", i), + Config: map[string]string{ + "foo": fmt.Sprintf("bar_%d", i), + }, + } + mustCreateKeyspace(suite, param) + } + // List all keyspaces, there should be 11 of them (default + 10 created above). + args := []string{"-u", suite.pdAddr, "keyspace", "list"} + output, err := pdctl.ExecuteCommand(pdctlCmd.GetRootCmd(), args...) + re.NoError(err) + var resp api.LoadAllKeyspacesResponse + re.NoError(json.Unmarshal(output, &resp)) + re.Len(resp.Keyspaces, 11) + re.Equal("", resp.NextPageToken) // No next page token since we load them all. + re.Equal("DEFAULT", resp.Keyspaces[0].GetName()) + for i, meta := range resp.Keyspaces[1:] { + re.Equal(fmt.Sprintf("test_keyspace_%d", i), meta.GetName()) + re.Equal(fmt.Sprintf("bar_%d", i), meta.Config["foo"]) + } + // List 3 keyspaces staring with keyspace id 3, should results in keyspace id 3, 4, 5 and next page token 6. + args = []string{"-u", suite.pdAddr, "keyspace", "list", "--limit", "3", "--page_token", "3"} + output, err = pdctl.ExecuteCommand(pdctlCmd.GetRootCmd(), args...) + re.NoError(err) + re.NoError(json.Unmarshal(output, &resp)) + re.Len(resp.Keyspaces, 3) + for i, meta := range resp.Keyspaces { + re.Equal(uint32(i+3), meta.GetId()) + re.Equal(fmt.Sprintf("test_keyspace_%d", i+2), meta.GetName()) + re.Equal(fmt.Sprintf("bar_%d", i+2), meta.Config["foo"]) + } + re.Equal("6", resp.NextPageToken) +} diff --git a/tools/pd-ctl/pdctl/command/keyspace_command.go b/tools/pd-ctl/pdctl/command/keyspace_command.go index a68e2f05a80..7c0d3d78bf6 100644 --- a/tools/pd-ctl/pdctl/command/keyspace_command.go +++ b/tools/pd-ctl/pdctl/command/keyspace_command.go @@ -15,33 +15,278 @@ package command import ( + "bytes" + "encoding/json" "fmt" "net/http" + "strings" "github.com/spf13/cobra" + "github.com/tikv/pd/server/apiv2/handlers" ) -const keyspacePrefix = "pd/api/v2/keyspaces" +const ( + keyspacePrefix = "pd/api/v2/keyspaces" + // flags + nmConfig = "config" + nmLimit = "limit" + nmPageToken = "page_token" + nmRemove = "remove" + nmUpdate = "update" +) // NewKeyspaceCommand returns a keyspace subcommand of rootCmd. func NewKeyspaceCommand() *cobra.Command { cmd := &cobra.Command{ - Use: "keyspace [command] [flags]", - Short: "show keyspace information", - Run: showKeyspaceCommandFunc, + Use: "keyspace [flags]", + Short: "keyspace commands", } + cmd.AddCommand(newShowKeyspaceCommand()) + cmd.AddCommand(newCreateKeyspaceCommand()) + cmd.AddCommand(newUpdateKeyspaceConfigCommand()) + cmd.AddCommand(newUpdateKeyspaceStateCommand()) + cmd.AddCommand(newListKeyspaceCommand()) return cmd } -func showKeyspaceCommandFunc(cmd *cobra.Command, args []string) { +func newShowKeyspaceCommand() *cobra.Command { + r := &cobra.Command{ + Use: "show", + Short: "show keyspace metadata", + } + showByID := &cobra.Command{ + Use: "id ", + Short: "show keyspace metadata specified by keyspace id", + Run: showKeyspaceIDCommandFunc, + } + showByName := &cobra.Command{ + Use: "name ", + Short: "show keyspace metadata specified by keyspace name", + Run: showKeyspaceNameCommandFunc, + } + r.AddCommand(showByID) + r.AddCommand(showByName) + return r +} + +func showKeyspaceIDCommandFunc(cmd *cobra.Command, args []string) { if len(args) != 1 { cmd.Usage() return } + resp, err := doRequest(cmd, fmt.Sprintf("%s/id/%s", keyspacePrefix, args[0]), http.MethodGet, http.Header{}) + if err != nil { + cmd.PrintErrln("Failed to get the keyspace information: ", err) + return + } + cmd.Println(resp) +} +func showKeyspaceNameCommandFunc(cmd *cobra.Command, args []string) { + if len(args) != 1 { + cmd.Usage() + return + } resp, err := doRequest(cmd, fmt.Sprintf("%s/%s?force_refresh_group_id=true", keyspacePrefix, args[0]), http.MethodGet, http.Header{}) if err != nil { - cmd.Printf("Failed to get the keyspace information: %s\n", err) + cmd.PrintErrln("Failed to get the keyspace information: ", err) + return + } + cmd.Println(resp) +} + +func newCreateKeyspaceCommand() *cobra.Command { + r := &cobra.Command{ + Use: "create [flags]", + Short: "create a keyspace", + Run: createKeyspaceCommandFunc, + } + r.Flags().StringSlice(nmConfig, nil, "keyspace configs for the new keyspace\n"+ + "specify as comma separated key value pairs, e.g. --config k1=v1,k2=v2") + return r +} + +func createKeyspaceCommandFunc(cmd *cobra.Command, args []string) { + if len(args) != 1 { + cmd.Usage() + return + } + + configPairs, err := cmd.Flags().GetStringSlice(nmConfig) + if err != nil { + cmd.PrintErrln("Failed to parse flag: ", err) + return + } + config := map[string]string{} + for _, flag := range configPairs { + kvs := strings.Split(flag, ",") + for _, kv := range kvs { + pair := strings.Split(kv, "=") + if len(pair) != 2 { + cmd.PrintErrf("Failed to create keyspace: invalid kv pair %s\n", kv) + return + } + if _, exist := config[pair[0]]; exist { + cmd.PrintErrf("Failed to create keyspace: key %s is specified multiple times\n", pair[0]) + return + } + config[pair[0]] = pair[1] + } + } + params := handlers.CreateKeyspaceParams{ + Name: args[0], + Config: config, + } + body, err := json.Marshal(params) + if err != nil { + cmd.PrintErrln("Failed to encode the request body: ", err) + return + } + resp, err := doRequest(cmd, keyspacePrefix, http.MethodPost, http.Header{}, WithBody(bytes.NewBuffer(body))) + if err != nil { + cmd.PrintErrln("Failed to create the keyspace: ", err) + return + } + cmd.Println(resp) +} + +func newUpdateKeyspaceConfigCommand() *cobra.Command { + r := &cobra.Command{ + Use: "update-config ", + Short: "update keyspace config", + Run: updateKeyspaceConfigCommandFunc, + } + r.Flags().StringSlice(nmRemove, nil, "keys to remove from keyspace config\n"+ + "specify as comma separated keys, e.g. --remove k1,k2") + r.Flags().StringSlice(nmUpdate, nil, "kv pairs to upsert into keyspace config\n"+ + "specify as comma separated key value pairs, e.g. --update k1=v1,k2=v2") + return r +} + +func updateKeyspaceConfigCommandFunc(cmd *cobra.Command, args []string) { + if len(args) != 1 { + cmd.Usage() + return + } + configPatch := map[string]*string{} + removeFlags, err := cmd.Flags().GetStringSlice(nmRemove) + if err != nil { + cmd.PrintErrln("Failed to parse flag: ", err) + return + } + for _, flag := range removeFlags { + keys := strings.Split(flag, ",") + for _, key := range keys { + if _, exist := configPatch[key]; exist { + cmd.PrintErrf("Failed to update keyspace config: key %s is specified multiple times\n", key) + return + } + configPatch[key] = nil + } + } + updateFlags, err := cmd.Flags().GetStringSlice(nmUpdate) + if err != nil { + cmd.PrintErrln("Failed to parse flag: ", err) + return + } + for _, flag := range updateFlags { + kvs := strings.Split(flag, ",") + for _, kv := range kvs { + pair := strings.Split(kv, "=") + if len(pair) != 2 { + cmd.PrintErrf("Failed to update keyspace config: invalid kv pair %s\n", kv) + return + } + if _, exist := configPatch[pair[0]]; exist { + cmd.PrintErrf("Failed to update keyspace config: key %s is specified multiple times\n", pair[0]) + return + } + configPatch[pair[0]] = &pair[1] + } + } + params := handlers.UpdateConfigParams{Config: configPatch} + data, err := json.Marshal(params) + if err != nil { + cmd.PrintErrln("Failed to update keyspace config:", err) + return + } + url := fmt.Sprintf("%s/%s/config", keyspacePrefix, args[0]) + resp, err := doRequest(cmd, url, http.MethodPatch, http.Header{}, WithBody(bytes.NewBuffer(data))) + if err != nil { + cmd.PrintErrln("Failed to update the keyspace config: ", err) + return + } + cmd.Println(resp) +} + +func newUpdateKeyspaceStateCommand() *cobra.Command { + r := &cobra.Command{ + Use: "update-state ", + Long: "update keyspace state, state can be one of: ENABLED, DISABLED, ARCHIVED, TOMBSTONE", + Run: updateKeyspaceStateCommandFunc, + } + return r +} + +func updateKeyspaceStateCommandFunc(cmd *cobra.Command, args []string) { + if len(args) != 2 { + cmd.Usage() + return + } + params := handlers.UpdateStateParam{ + State: args[1], + } + data, err := json.Marshal(params) + if err != nil { + cmd.PrintErrln(err) + return + } + url := fmt.Sprintf("%s/%s/state", keyspacePrefix, args[0]) + resp, err := doRequest(cmd, url, http.MethodPut, http.Header{}, WithBody(bytes.NewBuffer(data))) + if err != nil { + cmd.PrintErrln("Failed to update the keyspace state: ", err) + return + } + cmd.Println(resp) +} + +func newListKeyspaceCommand() *cobra.Command { + r := &cobra.Command{ + Use: "list [flags]", + Short: "list keyspaces according to filters", + Run: listKeyspaceCommandFunc, + } + r.Flags().String(nmLimit, "", "The maximum number of keyspace metas to return. If not set, no limit is posed.") + r.Flags().String(nmPageToken, "", "The keyspace id of the scan start. If not set, scan from keyspace/keyspace group with id 0") + return r +} + +func listKeyspaceCommandFunc(cmd *cobra.Command, args []string) { + if len(args) != 0 { + cmd.Usage() + return + } + + url := keyspacePrefix + limit, err := cmd.Flags().GetString(nmLimit) + if err != nil { + cmd.PrintErrln("Failed to parse flag: ", err) + return + } + if limit != "" { + url += fmt.Sprintf("?limit=%s", limit) + } + pageToken, err := cmd.Flags().GetString(nmPageToken) + if err != nil { + cmd.PrintErrln("Failed to parse flag: ", err) + return + } + if pageToken != "" { + url += fmt.Sprintf("&page_token=%s", pageToken) + } + resp, err := doRequest(cmd, url, http.MethodGet, http.Header{}) + if err != nil { + cmd.PrintErrln("Failed to list keyspace: ", err) return } cmd.Println(resp)