diff --git a/client/alias.go b/client/alias.go new file mode 100644 index 0000000000000..c7cc76d2047a3 --- /dev/null +++ b/client/alias.go @@ -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 +} diff --git a/client/alias_options.go b/client/alias_options.go new file mode 100644 index 0000000000000..26123cafcc0c4 --- /dev/null +++ b/client/alias_options.go @@ -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, + } +} diff --git a/client/alias_test.go b/client/alias_test.go new file mode 100644 index 0000000000000..582ae0b89dea3 --- /dev/null +++ b/client/alias_test.go @@ -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)) +} diff --git a/client/collection.go b/client/collection.go index 4031c687d9993..44fd08588b604 100644 --- a/client/collection.go +++ b/client/collection.go @@ -129,3 +129,12 @@ func (c *Client) DropCollection(ctx context.Context, option DropCollectionOption }) 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) + }) +} diff --git a/client/collection_options.go b/client/collection_options.go index 046a2d21c5027..2263e60323b2b 100644 --- a/client/collection_options.go +++ b/client/collection_options.go @@ -240,3 +240,26 @@ func NewDropCollectionOption(name string) *dropCollectionOption { 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, + } +} diff --git a/client/collection_test.go b/client/collection_test.go index 0bab40f48335c..1c14c5267b776 100644 --- a/client/collection_test.go +++ b/client/collection_test.go @@ -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) { suite.Run(t, new(CollectionSuite)) } diff --git a/client/entity/alias.go b/client/entity/alias.go new file mode 100644 index 0000000000000..0e69c332d596d --- /dev/null +++ b/client/entity/alias.go @@ -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 +}