enhance: [GoSDK] Add collection alias & rename API (#36990)

Related to #31293

---------

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
This commit is contained in:
congqixia 2024-10-21 18:45:33 +08:00 committed by GitHub
parent 0be81d95aa
commit e6cb7d6695
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 475 additions and 0 deletions

92
client/alias.go Normal file
View File

@ -0,0 +1,92 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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 client
import (
"context"
"google.golang.org/grpc"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus/client/v2/entity"
"github.com/milvus-io/milvus/pkg/util/merr"
)
func (c *Client) CreateAlias(ctx context.Context, option CreateAliasOption, callOptions ...grpc.CallOption) error {
req := option.Request()
return c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
resp, err := milvusService.CreateAlias(ctx, req, callOptions...)
return merr.CheckRPCCall(resp, err)
})
}
func (c *Client) DescribeAlias(ctx context.Context, option DescribeAliasOption, callOptions ...grpc.CallOption) (*entity.Alias, error) {
req := option.Request()
var alias *entity.Alias
err := c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
resp, err := milvusService.DescribeAlias(ctx, req, callOptions...)
if err := merr.CheckRPCCall(resp, err); err != nil {
return err
}
alias = &entity.Alias{
DbName: resp.GetDbName(),
Alias: resp.GetAlias(),
CollectionName: resp.GetCollection(),
}
return nil
})
return alias, err
}
func (c *Client) DropAlias(ctx context.Context, option DropAliasOption, callOptions ...grpc.CallOption) error {
req := option.Request()
return c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
resp, err := milvusService.DropAlias(ctx, req, callOptions...)
return merr.CheckRPCCall(resp, err)
})
}
func (c *Client) AlterAlias(ctx context.Context, option AlterAliasOption, callOptions ...grpc.CallOption) error {
req := option.Request()
return c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
resp, err := milvusService.AlterAlias(ctx, req, callOptions...)
return merr.CheckRPCCall(resp, err)
})
}
func (c *Client) ListAliases(ctx context.Context, option ListAliasesOption, callOptions ...grpc.CallOption) ([]string, error) {
req := option.Request()
var aliases []string
err := c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
resp, err := milvusService.ListAliases(ctx, req, callOptions...)
if err := merr.CheckRPCCall(resp, err); err != nil {
return err
}
aliases = resp.GetAliases()
return nil
})
return aliases, err
}

130
client/alias_options.go Normal file
View File

@ -0,0 +1,130 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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 client
import "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
// CreateCollectionOption is the interface builds CreateAliasRequest.
type CreateAliasOption interface {
Request() *milvuspb.CreateAliasRequest
}
type createAliasOption struct {
collectionName string
alias string
}
func (opt *createAliasOption) Request() *milvuspb.CreateAliasRequest {
return &milvuspb.CreateAliasRequest{
CollectionName: opt.collectionName,
Alias: opt.alias,
}
}
func NewCreateAliasOption(collectionName, alias string) *createAliasOption {
return &createAliasOption{
collectionName: collectionName,
alias: alias,
}
}
// DescribeAliasOption is the interface builds DescribeAliasOption.
type DescribeAliasOption interface {
Request() *milvuspb.DescribeAliasRequest
}
type describeAliasOption struct {
aliasName string
}
func (opt *describeAliasOption) Request() *milvuspb.DescribeAliasRequest {
return &milvuspb.DescribeAliasRequest{
Alias: opt.aliasName,
}
}
func NewDescribeAliasOption(alias string) *describeAliasOption {
return &describeAliasOption{
aliasName: alias,
}
}
// DropAliasOption is the interface builds DropAliasRequest.
type DropAliasOption interface {
Request() *milvuspb.DropAliasRequest
}
type dropAliasOption struct {
aliasName string
}
func (opt *dropAliasOption) Request() *milvuspb.DropAliasRequest {
return &milvuspb.DropAliasRequest{
Alias: opt.aliasName,
}
}
func NewDropAliasOption(alias string) *dropAliasOption {
return &dropAliasOption{
aliasName: alias,
}
}
// AlterAliasOption is the interface builds AlterAliasRequest.
type AlterAliasOption interface {
Request() *milvuspb.AlterAliasRequest
}
type alterAliasOption struct {
aliasName string
collectionName string
}
func (opt *alterAliasOption) Request() *milvuspb.AlterAliasRequest {
return &milvuspb.AlterAliasRequest{
Alias: opt.aliasName,
CollectionName: opt.collectionName,
}
}
func NewAlterAliasOption(alias, collectionName string) *alterAliasOption {
return &alterAliasOption{
aliasName: alias,
collectionName: collectionName,
}
}
// ListAliasesOption is the interface builds ListAliasesRequest.
type ListAliasesOption interface {
Request() *milvuspb.ListAliasesRequest
}
type listAliasesOption struct {
collectionName string
}
func (opt *listAliasesOption) Request() *milvuspb.ListAliasesRequest {
return &milvuspb.ListAliasesRequest{
CollectionName: opt.collectionName,
}
}
func NewListAliasesOption(collectionName string) *listAliasesOption {
return &listAliasesOption{
collectionName: collectionName,
}
}

171
client/alias_test.go Normal file
View File

@ -0,0 +1,171 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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 client
import (
"context"
"fmt"
"testing"
"github.com/stretchr/testify/mock"
"github.com/stretchr/testify/suite"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus/pkg/util/merr"
)
type AliasSuite struct {
MockSuiteBase
}
func (s *AliasSuite) TestCreateAlias() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
aliasName := fmt.Sprintf("test_alias_%s", s.randString(6))
collectionName := fmt.Sprintf("test_collection_%s", s.randString(6))
s.Run("success", func() {
s.mock.EXPECT().CreateAlias(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, car *milvuspb.CreateAliasRequest) (*commonpb.Status, error) {
s.Equal(aliasName, car.GetAlias())
s.Equal(collectionName, car.GetCollectionName())
return merr.Success(), nil
}).Once()
err := s.client.CreateAlias(ctx, NewCreateAliasOption(collectionName, aliasName))
s.NoError(err)
})
s.Run("failure", func() {
s.mock.EXPECT().CreateAlias(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once()
err := s.client.CreateAlias(ctx, NewCreateAliasOption(collectionName, aliasName))
s.Error(err)
})
}
func (s *AliasSuite) TestDropAlias() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
aliasName := fmt.Sprintf("test_alias_%s", s.randString(6))
s.Run("success", func() {
s.mock.EXPECT().DropAlias(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, dar *milvuspb.DropAliasRequest) (*commonpb.Status, error) {
s.Equal(aliasName, dar.GetAlias())
return merr.Success(), nil
}).Once()
err := s.client.DropAlias(ctx, NewDropAliasOption(aliasName))
s.NoError(err)
})
s.Run("failure", func() {
s.mock.EXPECT().DropAlias(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once()
err := s.client.DropAlias(ctx, NewDropAliasOption(aliasName))
s.Error(err)
})
}
func (s *AliasSuite) TestDescribeAlias() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
aliasName := fmt.Sprintf("test_alias_%s", s.randString(6))
collectionName := fmt.Sprintf("test_collection_%s", s.randString(6))
s.Run("success", func() {
s.mock.EXPECT().DescribeAlias(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, car *milvuspb.DescribeAliasRequest) (*milvuspb.DescribeAliasResponse, error) {
s.Equal(aliasName, car.GetAlias())
return &milvuspb.DescribeAliasResponse{
Alias: aliasName,
Collection: collectionName,
}, nil
}).Once()
alias, err := s.client.DescribeAlias(ctx, NewDescribeAliasOption(aliasName))
s.NoError(err)
s.Equal(aliasName, alias.Alias)
s.Equal(collectionName, alias.CollectionName)
})
s.Run("failure", func() {
s.mock.EXPECT().DescribeAlias(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once()
_, err := s.client.DescribeAlias(ctx, NewDescribeAliasOption(aliasName))
s.Error(err)
})
}
func (s *AliasSuite) TestAlterAlias() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
aliasName := fmt.Sprintf("test_alias_%s", s.randString(6))
collectionName := fmt.Sprintf("test_collection_%s", s.randString(6))
s.Run("success", func() {
s.mock.EXPECT().AlterAlias(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, dar *milvuspb.AlterAliasRequest) (*commonpb.Status, error) {
s.Equal(aliasName, dar.GetAlias())
s.Equal(collectionName, dar.GetCollectionName())
return merr.Success(), nil
}).Once()
err := s.client.AlterAlias(ctx, NewAlterAliasOption(aliasName, collectionName))
s.NoError(err)
})
s.Run("failure", func() {
s.mock.EXPECT().AlterAlias(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once()
err := s.client.AlterAlias(ctx, NewAlterAliasOption(aliasName, collectionName))
s.Error(err)
})
}
func (s *AliasSuite) TestListAliases() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
collectionName := fmt.Sprintf("test_collection_%s", s.randString(6))
s.Run("success", func() {
s.mock.EXPECT().ListAliases(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, lar *milvuspb.ListAliasesRequest) (*milvuspb.ListAliasesResponse, error) {
s.Equal(collectionName, lar.GetCollectionName())
return &milvuspb.ListAliasesResponse{
Aliases: []string{"test1", "test2", "test3"},
}, nil
}).Once()
names, err := s.client.ListAliases(ctx, NewListAliasesOption(collectionName))
s.NoError(err)
s.ElementsMatch([]string{"test1", "test2", "test3"}, names)
})
s.Run("failure", func() {
s.mock.EXPECT().ListAliases(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once()
_, err := s.client.ListAliases(ctx, NewListAliasesOption(collectionName))
s.Error(err)
})
}
func TestAlias(t *testing.T) {
suite.Run(t, new(AliasSuite))
}

View File

@ -129,3 +129,12 @@ func (c *Client) DropCollection(ctx context.Context, option DropCollectionOption
}) })
return err return err
} }
func (c *Client) RenameCollection(ctx context.Context, option RenameCollectionOption, callOptions ...grpc.CallOption) error {
req := option.Request()
return c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
resp, err := milvusService.RenameCollection(ctx, req, callOptions...)
return merr.CheckRPCCall(resp, err)
})
}

View File

@ -240,3 +240,26 @@ func NewDropCollectionOption(name string) *dropCollectionOption {
name: name, name: name,
} }
} }
type RenameCollectionOption interface {
Request() *milvuspb.RenameCollectionRequest
}
type renameCollectionOption struct {
oldCollectionName string
newCollectionName string
}
func (opt *renameCollectionOption) Request() *milvuspb.RenameCollectionRequest {
return &milvuspb.RenameCollectionRequest{
OldName: opt.oldCollectionName,
NewName: opt.newCollectionName,
}
}
func NewRenameCollectionOption(oldName, newName string) *renameCollectionOption {
return &renameCollectionOption{
oldCollectionName: oldName,
newCollectionName: newName,
}
}

View File

@ -248,6 +248,32 @@ func (s *CollectionSuite) TestDropCollection() {
}) })
} }
func (s *CollectionSuite) TestRenameCollection() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
oldName := fmt.Sprintf("test_collection_%s", s.randString(6))
newName := fmt.Sprintf("%s_new", oldName)
s.Run("success", func() {
s.mock.EXPECT().RenameCollection(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, rcr *milvuspb.RenameCollectionRequest) (*commonpb.Status, error) {
s.Equal(oldName, rcr.GetOldName())
s.Equal(newName, rcr.GetNewName())
return merr.Success(), nil
}).Once()
err := s.client.RenameCollection(ctx, NewRenameCollectionOption(oldName, newName))
s.NoError(err)
})
s.Run("failure", func() {
s.mock.EXPECT().RenameCollection(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once()
err := s.client.RenameCollection(ctx, NewRenameCollectionOption(oldName, newName))
s.Error(err)
})
}
func TestCollection(t *testing.T) { func TestCollection(t *testing.T) {
suite.Run(t, new(CollectionSuite)) suite.Run(t, new(CollectionSuite))
} }

24
client/entity/alias.go Normal file
View File

@ -0,0 +1,24 @@
// Licensed to the LF AI & Data foundation under one
// or more contributor license agreements. See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership. The ASF licenses this file
// to you 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 entity
// Alias is the entity model for collection alias.
type Alias struct {
DbName string
Alias string
CollectionName string
}