pd-ctl: add keyspace commands

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)

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,19 +88,196 @@ 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))
 		return newGroupID == k.Config[keyspace.TSOKeyspaceGroupIDKey]
 	})
 
 	// 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")
 	re.NoError(failpoint.Disable("github.com/tikv/pd/pkg/keyspace/acceleratedAllocNodes"))
 	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.GetLeaderServer()
+	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)
+}