diff --git a/pkg/storage/endpoint/keyspace_group.go b/pkg/storage/endpoint/keyspace_group.go index 0681a2a4345b..d6d4d12ec84b 100644 --- a/pkg/storage/endpoint/keyspace_group.go +++ b/pkg/storage/endpoint/keyspace_group.go @@ -31,7 +31,7 @@ type KeyspaceGroup struct { // KeyspaceGroupStorage is the interface for keyspace group storage. type KeyspaceGroupStorage interface { - LoadKeyspaceGroups() ([]*KeyspaceGroup, error) + LoadKeyspaceGroups(startID uint32, limit int) ([]*KeyspaceGroup, error) LoadKeyspaceGroup(txn kv.Txn, id uint32) (*KeyspaceGroup, error) SaveKeyspaceGroup(txn kv.Txn, kg *KeyspaceGroup) error DeleteKeyspaceGroup(txn kv.Txn, id uint32) error @@ -70,10 +70,10 @@ func (se *StorageEndpoint) DeleteKeyspaceGroup(txn kv.Txn, id uint32) error { } // LoadKeyspaceGroups loads all keyspace groups. -func (se *StorageEndpoint) LoadKeyspaceGroups() ([]*KeyspaceGroup, error) { - prefix := KeyspaceGroupIDPrefix() - prefixEnd := clientv3.GetPrefixRangeEnd(prefix) - keys, values, err := se.LoadRange(prefix, prefixEnd, 0) +func (se *StorageEndpoint) LoadKeyspaceGroups(startID uint32, limit int) ([]*KeyspaceGroup, error) { + prefix := KeyspaceGroupIDPath(startID) + prefixEnd := clientv3.GetPrefixRangeEnd(KeyspaceGroupIDPrefix()) + keys, values, err := se.LoadRange(prefix, prefixEnd, limit) if err != nil { return nil, err } diff --git a/server/apiv2/handlers/keyspace_group.go b/server/apiv2/handlers/keyspace_group.go index 4578baf03fd2..2723c20cf4f7 100644 --- a/server/apiv2/handlers/keyspace_group.go +++ b/server/apiv2/handlers/keyspace_group.go @@ -27,7 +27,7 @@ import ( // RegisterTSOKeyspaceGroup registers keyspace group handlers to the server. func RegisterTSOKeyspaceGroup(r *gin.RouterGroup) { - router := r.Group("tso/keyspace-group") + router := r.Group("tso/keyspace-groups") router.Use(middlewares.BootstrapChecker()) router.POST("", CreateKeyspaceGroups) router.GET("", GetKeyspaceGroups) @@ -58,7 +58,7 @@ func CreateKeyspaceGroups(c *gin.Context) { } } - err = manager.CreateKeyspaces(createParams.KeyspaceGroups) + err = manager.CreateKeyspaceGroups(createParams.KeyspaceGroups) if err != nil { c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error()) return @@ -70,7 +70,12 @@ func CreateKeyspaceGroups(c *gin.Context) { func GetKeyspaceGroups(c *gin.Context) { svr := c.MustGet("server").(*server.Server) manager := svr.GetKeyspaceGroupManager() - keyspaceGroups, err := manager.GetKeyspaceGroups() + scanStart, scanLimit, err := parseLoadAllQuery(c) + if err != nil { + c.AbortWithStatusJSON(http.StatusBadRequest, err.Error()) + return + } + keyspaceGroups, err := manager.GetKeyspaceGroups(scanStart, scanLimit) if err != nil { c.AbortWithStatusJSON(http.StatusInternalServerError, err.Error()) return diff --git a/server/keyspace/keyspace_group.go b/server/keyspace/keyspace_group.go index f3d26235bdfb..d5e4eec06196 100644 --- a/server/keyspace/keyspace_group.go +++ b/server/keyspace/keyspace_group.go @@ -57,8 +57,8 @@ func (m *GroupManager) Bootstrap() error { return nil } -// CreateKeyspaces creates keyspace groups. -func (m *GroupManager) CreateKeyspaces(keyspaceGroups []*endpoint.KeyspaceGroup) error { +// CreateKeyspaceGroups creates keyspace groups. +func (m *GroupManager) CreateKeyspaceGroups(keyspaceGroups []*endpoint.KeyspaceGroup) error { for _, keyspaceGroup := range keyspaceGroups { // TODO: add replica count kg := &endpoint.KeyspaceGroup{ @@ -78,8 +78,8 @@ func (m *GroupManager) CreateKeyspaces(keyspaceGroups []*endpoint.KeyspaceGroup) } // GetKeyspaceGroups returns all keyspace groups. -func (m *GroupManager) GetKeyspaceGroups() ([]*endpoint.KeyspaceGroup, error) { - return m.store.LoadKeyspaceGroups() +func (m *GroupManager) GetKeyspaceGroups(startID uint32, limit int) ([]*endpoint.KeyspaceGroup, error) { + return m.store.LoadKeyspaceGroups(startID, limit) } // GetKeyspaceGroupByID returns the keyspace group by id. diff --git a/server/keyspace/keyspace_group_test.go b/server/keyspace/keyspace_group_test.go index 38ed18ed0e25..ef448e12b65c 100644 --- a/server/keyspace/keyspace_group_test.go +++ b/server/keyspace/keyspace_group_test.go @@ -55,12 +55,16 @@ func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupOperations() { UserKind: "business", }, } - err := suite.manager.CreateKeyspaces(keyspaceGroups) + err := suite.manager.CreateKeyspaceGroups(keyspaceGroups) re.NoError(err) // list all keyspace groups - kgs, err := suite.manager.GetKeyspaceGroups() + kgs, err := suite.manager.GetKeyspaceGroups(uint32(0), 0) re.NoError(err) re.Len(kgs, 4) + // list part of keyspace groups + kgs, err = suite.manager.GetKeyspaceGroups(uint32(1), 0) + re.NoError(err) + re.Len(kgs, 3) // get the default keyspace group kg, err := suite.manager.GetKeyspaceGroupByID(0) re.NoError(err) @@ -80,7 +84,7 @@ func (suite *keyspaceGroupTestSuite) TestKeyspaceGroupOperations() { // create an existing keyspace group keyspaceGroups = []*endpoint.KeyspaceGroup{{ID: uint32(1), UserKind: "business"}} - err = suite.manager.CreateKeyspaces(keyspaceGroups) + err = suite.manager.CreateKeyspaceGroups(keyspaceGroups) // only print the warning log re.NoError(err) } diff --git a/tests/server/apiv2/handlers/keyspace_group_test.go b/tests/server/apiv2/handlers/keyspace_group_test.go new file mode 100644 index 000000000000..e5cc9b8b582f --- /dev/null +++ b/tests/server/apiv2/handlers/keyspace_group_test.go @@ -0,0 +1,126 @@ +// 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 handlers_test + +import ( + "bytes" + "context" + "encoding/json" + "io" + "net/http" + "testing" + + "github.com/stretchr/testify/require" + "github.com/stretchr/testify/suite" + "github.com/tikv/pd/pkg/storage/endpoint" + "github.com/tikv/pd/server/apiv2/handlers" + "github.com/tikv/pd/tests" +) + +const keyspaceGroupsPrefix = "/pd/api/v2/tso/keyspace-groups" + +type keyspaceGroupTestSuite struct { + suite.Suite + cleanup func() + cluster *tests.TestCluster + server *tests.TestServer +} + +func TestKeyspaceGroupTestSuite(t *testing.T) { + suite.Run(t, new(keyspaceGroupTestSuite)) +} + +func (suite *keyspaceGroupTestSuite) SetupTest() { + ctx, cancel := context.WithCancel(context.Background()) + suite.cleanup = cancel + cluster, err := tests.NewTestCluster(ctx, 1) + suite.cluster = cluster + suite.NoError(err) + suite.NoError(cluster.RunInitialServers()) + suite.NotEmpty(cluster.WaitLeader()) + suite.server = cluster.GetServer(cluster.GetLeader()) + suite.NoError(suite.server.BootstrapCluster()) +} + +func (suite *keyspaceGroupTestSuite) TearDownTest() { + suite.cleanup() + suite.cluster.Destroy() +} + +func (suite *keyspaceGroupTestSuite) TestCreateKeyspaceGroup() { + re := suite.Require() + kgs := &handlers.CreateKeyspaceGroupParams{KeyspaceGroups: []*endpoint.KeyspaceGroup{ + { + ID: uint32(1), + UserKind: "business", + }, + { + ID: uint32(2), + UserKind: "business", + }, + }} + + mustCreateKeyspaceGroup(re, suite.server, kgs) +} + +func (suite *keyspaceGroupTestSuite) TestLoadKeyspaceGroup() { + re := suite.Require() + kgs := &handlers.CreateKeyspaceGroupParams{KeyspaceGroups: []*endpoint.KeyspaceGroup{ + { + ID: uint32(1), + UserKind: "business", + }, + { + ID: uint32(2), + UserKind: "business", + }, + }} + + mustCreateKeyspaceGroup(re, suite.server, kgs) + resp := sendLoadKeyspaceGroupRequest(re, suite.server, "0", "0") + re.Equal(3, len(resp)) +} + +func sendLoadKeyspaceGroupRequest(re *require.Assertions, server *tests.TestServer, token, limit string) []*endpoint.KeyspaceGroup { + // Construct load range request. + httpReq, err := http.NewRequest(http.MethodGet, server.GetAddr()+keyspaceGroupsPrefix, nil) + re.NoError(err) + query := httpReq.URL.Query() + query.Add("page_token", token) + query.Add("limit", limit) + httpReq.URL.RawQuery = query.Encode() + // Send request. + httpResp, err := dialClient.Do(httpReq) + re.NoError(err) + defer httpResp.Body.Close() + re.Equal(http.StatusOK, httpResp.StatusCode) + // Receive & decode response. + data, err := io.ReadAll(httpResp.Body) + re.NoError(err) + var resp []*endpoint.KeyspaceGroup + re.NoError(json.Unmarshal(data, &resp)) + return resp +} + +func mustCreateKeyspaceGroup(re *require.Assertions, server *tests.TestServer, request *handlers.CreateKeyspaceGroupParams) { + data, err := json.Marshal(request) + re.NoError(err) + httpReq, err := http.NewRequest(http.MethodPost, server.GetAddr()+keyspaceGroupsPrefix, bytes.NewBuffer(data)) + re.NoError(err) + resp, err := dialClient.Do(httpReq) + re.NoError(err) + defer resp.Body.Close() + re.Equal(http.StatusOK, resp.StatusCode) +}