From c39452f384dcde3e69f9d2be3b05ea99acfcfef1 Mon Sep 17 00:00:00 2001 From: congqixia Date: Fri, 27 Dec 2024 18:26:50 +0800 Subject: [PATCH] enhance: [GoSDK] Support alter properties APIs (#38812) Related to #31293 This PR: - Add AlterDatabaseProperties API - Add DropDatabaseProperties API - Add DescribeDatabase API - Rename AlterCollection to AlterCollectionProperties - Add DropCollectionProperties API - Add AlterCollectionFieldProperties API Signed-off-by: Congqi Xia --- client/entity/database.go | 22 +++++ client/milvusclient/collection.go | 20 ++++- .../milvusclient/collection_example_test.go | 2 +- client/milvusclient/collection_options.go | 66 +++++++++++++-- client/milvusclient/collection_test.go | 64 +++++++++++++- client/milvusclient/database.go | 40 +++++++++ client/milvusclient/database_options.go | 78 ++++++++++++++++- client/milvusclient/database_test.go | 83 +++++++++++++++++++ 8 files changed, 363 insertions(+), 12 deletions(-) create mode 100644 client/entity/database.go diff --git a/client/entity/database.go b/client/entity/database.go new file mode 100644 index 0000000000..87cfdbd580 --- /dev/null +++ b/client/entity/database.go @@ -0,0 +1,22 @@ +// 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 + +type Database struct { + Name string + Properties map[string]string +} diff --git a/client/milvusclient/collection.go b/client/milvusclient/collection.go index 025942bcfa..f723916c2b 100644 --- a/client/milvusclient/collection.go +++ b/client/milvusclient/collection.go @@ -139,7 +139,7 @@ func (c *Client) RenameCollection(ctx context.Context, option RenameCollectionOp }) } -func (c *Client) AlterCollection(ctx context.Context, option AlterCollectionOption, callOptions ...grpc.CallOption) error { +func (c *Client) AlterCollectionProperties(ctx context.Context, option AlterCollectionPropertiesOption, callOptions ...grpc.CallOption) error { req := option.Request() return c.callService(func(milvusService milvuspb.MilvusServiceClient) error { @@ -148,6 +148,24 @@ func (c *Client) AlterCollection(ctx context.Context, option AlterCollectionOpti }) } +func (c *Client) DropCollectionProperties(ctx context.Context, option DropCollectionPropertiesOption, callOptions ...grpc.CallOption) error { + req := option.Request() + + return c.callService(func(milvusService milvuspb.MilvusServiceClient) error { + resp, err := milvusService.AlterCollection(ctx, req, callOptions...) + return merr.CheckRPCCall(resp, err) + }) +} + +func (c *Client) AlterCollectionFieldProperty(ctx context.Context, option AlterCollectionFieldPropertiesOption, callOptions ...grpc.CallOption) error { + req := option.Request() + + return c.callService(func(milvusService milvuspb.MilvusServiceClient) error { + resp, err := milvusService.AlterCollectionField(ctx, req, callOptions...) + return merr.CheckRPCCall(resp, err) + }) +} + type GetCollectionOption interface { Request() *milvuspb.GetCollectionStatisticsRequest } diff --git a/client/milvusclient/collection_example_test.go b/client/milvusclient/collection_example_test.go index 612335cff8..1d573e27e4 100644 --- a/client/milvusclient/collection_example_test.go +++ b/client/milvusclient/collection_example_test.go @@ -260,7 +260,7 @@ func ExampleClient_AlterCollection_setTTL() { defer cli.Close(ctx) - err = cli.AlterCollection(ctx, milvusclient.NewAlterCollectionOption("my_collection").WithProperty(common.CollectionTTLConfigKey, 60)) + err = cli.AlterCollectionProperties(ctx, milvusclient.NewAlterCollectionPropertiesOption("my_collection").WithProperty(common.CollectionTTLConfigKey, 60)) if err != nil { // handle error } diff --git a/client/milvusclient/collection_options.go b/client/milvusclient/collection_options.go index 0907cb6c8f..1bd8029d63 100644 --- a/client/milvusclient/collection_options.go +++ b/client/milvusclient/collection_options.go @@ -286,29 +286,83 @@ func NewRenameCollectionOption(oldName, newName string) *renameCollectionOption } } -type AlterCollectionOption interface { +type AlterCollectionPropertiesOption interface { Request() *milvuspb.AlterCollectionRequest } -type alterCollectionOption struct { +type alterCollectionPropertiesOption struct { collectionName string properties map[string]string } -func (opt *alterCollectionOption) WithProperty(key string, value any) *alterCollectionOption { +func (opt *alterCollectionPropertiesOption) WithProperty(key string, value any) *alterCollectionPropertiesOption { opt.properties[key] = fmt.Sprintf("%v", value) return opt } -func (opt *alterCollectionOption) Request() *milvuspb.AlterCollectionRequest { +func (opt *alterCollectionPropertiesOption) Request() *milvuspb.AlterCollectionRequest { return &milvuspb.AlterCollectionRequest{ CollectionName: opt.collectionName, Properties: entity.MapKvPairs(opt.properties), } } -func NewAlterCollectionOption(collection string) *alterCollectionOption { - return &alterCollectionOption{collectionName: collection, properties: make(map[string]string)} +func NewAlterCollectionPropertiesOption(collection string) *alterCollectionPropertiesOption { + return &alterCollectionPropertiesOption{collectionName: collection, properties: make(map[string]string)} +} + +type DropCollectionPropertiesOption interface { + Request() *milvuspb.AlterCollectionRequest +} + +type dropCollectionPropertiesOption struct { + collectionName string + keys []string +} + +func (opt *dropCollectionPropertiesOption) Request() *milvuspb.AlterCollectionRequest { + return &milvuspb.AlterCollectionRequest{ + CollectionName: opt.collectionName, + DeleteKeys: opt.keys, + } +} + +func NewDropCollectionPropertiesOption(collection string, propertyKeys ...string) *dropCollectionPropertiesOption { + return &dropCollectionPropertiesOption{ + collectionName: collection, + keys: propertyKeys, + } +} + +type AlterCollectionFieldPropertiesOption interface { + Request() *milvuspb.AlterCollectionFieldRequest +} + +type alterCollectionFieldPropertiesOption struct { + collectionName string + fieldName string + properties map[string]string +} + +func (opt *alterCollectionFieldPropertiesOption) WithProperty(key string, value any) *alterCollectionFieldPropertiesOption { + opt.properties[key] = fmt.Sprintf("%v", value) + return opt +} + +func (opt *alterCollectionFieldPropertiesOption) Request() *milvuspb.AlterCollectionFieldRequest { + return &milvuspb.AlterCollectionFieldRequest{ + CollectionName: opt.collectionName, + FieldName: opt.fieldName, + Properties: entity.MapKvPairs(opt.properties), + } +} + +func NewAlterCollectionFieldPropertiesOption(collectionName string, fieldName string) *alterCollectionFieldPropertiesOption { + return &alterCollectionFieldPropertiesOption{ + collectionName: collectionName, + fieldName: fieldName, + properties: make(map[string]string), + } } type getCollectionStatsOption struct { diff --git a/client/milvusclient/collection_test.go b/client/milvusclient/collection_test.go index 4dc8e62d87..b905b24d58 100644 --- a/client/milvusclient/collection_test.go +++ b/client/milvusclient/collection_test.go @@ -285,7 +285,7 @@ func (s *CollectionSuite) TestRenameCollection() { }) } -func (s *CollectionSuite) TestAlterCollection() { +func (s *CollectionSuite) TestAlterCollectionProperties() { ctx, cancel := context.WithCancel(context.Background()) defer cancel() @@ -304,14 +304,72 @@ func (s *CollectionSuite) TestAlterCollection() { return merr.Success(), nil }).Once() - err := s.client.AlterCollection(ctx, NewAlterCollectionOption(collName).WithProperty(key, value)) + err := s.client.AlterCollectionProperties(ctx, NewAlterCollectionPropertiesOption(collName).WithProperty(key, value)) s.NoError(err) }) s.Run("failure", func() { s.mock.EXPECT().AlterCollection(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once() - err := s.client.AlterCollection(ctx, NewAlterCollectionOption(collName).WithProperty(key, value)) + err := s.client.AlterCollectionProperties(ctx, NewAlterCollectionPropertiesOption(collName).WithProperty(key, value)) + s.Error(err) + }) +} + +func (s *CollectionSuite) TestDropCollectionProperties() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + s.Run("success", func() { + dbName := fmt.Sprintf("dt_%s", s.randString(6)) + key := fmt.Sprintf("key_%s", s.randString(4)) + s.mock.EXPECT().AlterCollection(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, adr *milvuspb.AlterCollectionRequest) (*commonpb.Status, error) { + s.Equal([]string{key}, adr.GetDeleteKeys()) + return merr.Success(), nil + }).Once() + + err := s.client.DropCollectionProperties(ctx, NewDropCollectionPropertiesOption(dbName, key)) + s.NoError(err) + }) + + s.Run("failure", func() { + dbName := fmt.Sprintf("dt_%s", s.randString(6)) + s.mock.EXPECT().AlterCollection(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once() + + err := s.client.DropCollectionProperties(ctx, NewDropCollectionPropertiesOption(dbName, "key")) + s.Error(err) + }) +} + +func (s *CollectionSuite) TestAlterCollectionFieldProperties() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + collName := fmt.Sprintf("test_collection_%s", s.randString(6)) + fieldName := fmt.Sprintf("field_%s", s.randString(4)) + key := s.randString(6) + value := s.randString(6) + + s.Run("success", func() { + s.mock.EXPECT().AlterCollectionField(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, acr *milvuspb.AlterCollectionFieldRequest) (*commonpb.Status, error) { + s.Equal(collName, acr.GetCollectionName()) + s.Equal(fieldName, acr.GetFieldName()) + if s.Len(acr.GetProperties(), 1) { + item := acr.GetProperties()[0] + s.Equal(key, item.GetKey()) + s.Equal(value, item.GetValue()) + } + return merr.Success(), nil + }).Once() + + err := s.client.AlterCollectionFieldProperty(ctx, NewAlterCollectionFieldPropertiesOption(collName, fieldName).WithProperty(key, value)) + s.NoError(err) + }) + + s.Run("failure", func() { + s.mock.EXPECT().AlterCollectionField(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once() + + err := s.client.AlterCollectionFieldProperty(ctx, NewAlterCollectionFieldPropertiesOption("coll", "field").WithProperty(key, value)) s.Error(err) }) } diff --git a/client/milvusclient/database.go b/client/milvusclient/database.go index eb5b352963..681dca15a1 100644 --- a/client/milvusclient/database.go +++ b/client/milvusclient/database.go @@ -22,6 +22,7 @@ import ( "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" ) @@ -64,3 +65,42 @@ func (c *Client) DropDatabase(ctx context.Context, option DropDatabaseOption, ca return merr.CheckRPCCall(resp, err) }) } + +func (c *Client) DescribeDatabase(ctx context.Context, option DescribeDatabaseOption, callOptions ...grpc.CallOption) (*entity.Database, error) { + req := option.Request() + + var db *entity.Database + err := c.callService(func(milvusService milvuspb.MilvusServiceClient) error { + resp, err := milvusService.DescribeDatabase(ctx, req, callOptions...) + err = merr.CheckRPCCall(resp, err) + if err != nil { + return err + } + // databaseInfo = resp + db = &entity.Database{ + Name: resp.GetDbName(), + Properties: entity.KvPairsMap(resp.GetProperties()), + } + return nil + }) + + return db, err +} + +func (c *Client) AlterDatabaseProperies(ctx context.Context, option AlterDatabasePropertiesOption, callOptions ...grpc.CallOption) error { + req := option.Request() + + return c.callService(func(milvusService milvuspb.MilvusServiceClient) error { + resp, err := milvusService.AlterDatabase(ctx, req, callOptions...) + return merr.CheckRPCCall(resp, err) + }) +} + +func (c *Client) DropDatabaseProperties(ctx context.Context, option DropDatabasePropertiesOption, callOptions ...grpc.CallOption) error { + req := option.Request() + + return c.callService(func(milvusService milvuspb.MilvusServiceClient) error { + resp, err := milvusService.AlterDatabase(ctx, req, callOptions...) + return merr.CheckRPCCall(resp, err) + }) +} diff --git a/client/milvusclient/database_options.go b/client/milvusclient/database_options.go index 48542f9e58..4d644467ee 100644 --- a/client/milvusclient/database_options.go +++ b/client/milvusclient/database_options.go @@ -16,7 +16,12 @@ package milvusclient -import "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" +import ( + "fmt" + + "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" + "github.com/milvus-io/milvus/client/v2/entity" +) type UseDatabaseOption interface { DbName() string @@ -90,3 +95,74 @@ func NewDropDatabaseOption(dbName string) *dropDatabaseOption { dbName: dbName, } } + +type DescribeDatabaseOption interface { + Request() *milvuspb.DescribeDatabaseRequest +} + +type describeDatabaseOption struct { + dbName string +} + +func (opt *describeDatabaseOption) Request() *milvuspb.DescribeDatabaseRequest { + return &milvuspb.DescribeDatabaseRequest{ + DbName: opt.dbName, + } +} + +func NewDescribeDatabaseOption(dbName string) *describeDatabaseOption { + return &describeDatabaseOption{ + dbName: dbName, + } +} + +type AlterDatabasePropertiesOption interface { + Request() *milvuspb.AlterDatabaseRequest +} + +type alterDatabasePropertiesOption struct { + dbName string + properties map[string]string +} + +func (opt *alterDatabasePropertiesOption) Request() *milvuspb.AlterDatabaseRequest { + return &milvuspb.AlterDatabaseRequest{ + DbName: opt.dbName, + Properties: entity.MapKvPairs(opt.properties), + } +} + +func (opt *alterDatabasePropertiesOption) WithProperty(key string, value any) *alterDatabasePropertiesOption { + opt.properties[key] = fmt.Sprintf("%v", value) + return opt +} + +func NewAlterDatabasePropertiesOption(dbName string) *alterDatabasePropertiesOption { + return &alterDatabasePropertiesOption{ + dbName: dbName, + properties: make(map[string]string), + } +} + +type DropDatabasePropertiesOption interface { + Request() *milvuspb.AlterDatabaseRequest +} + +type dropDatabasePropertiesOption struct { + dbName string + keys []string +} + +func (opt *dropDatabasePropertiesOption) Request() *milvuspb.AlterDatabaseRequest { + return &milvuspb.AlterDatabaseRequest{ + DbName: opt.dbName, + DeleteKeys: opt.keys, + } +} + +func NewDropDatabasePropertiesOption(dbName string, propertyKeys ...string) *dropDatabasePropertiesOption { + return &dropDatabasePropertiesOption{ + dbName: dbName, + keys: propertyKeys, + } +} diff --git a/client/milvusclient/database_test.go b/client/milvusclient/database_test.go index 3cb3b7017f..541f095d22 100644 --- a/client/milvusclient/database_test.go +++ b/client/milvusclient/database_test.go @@ -108,6 +108,89 @@ func (s *DatabaseSuite) TestUseDatabase() { }) } +func (s *DatabaseSuite) TestDescribeDatabase() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + s.Run("success", func() { + dbName := fmt.Sprintf("dt_%s", s.randString(6)) + key := fmt.Sprintf("key_%s", s.randString(4)) + value := s.randString(6) + s.mock.EXPECT().DescribeDatabase(mock.Anything, mock.Anything).Return(&milvuspb.DescribeDatabaseResponse{ + Status: merr.Success(), + DbName: dbName, + Properties: []*commonpb.KeyValuePair{ + {Key: key, Value: value}, + }, + }, nil).Once() + + db, err := s.client.DescribeDatabase(ctx, NewDescribeDatabaseOption(dbName)) + s.NoError(err) + s.Equal(dbName, db.Name) + s.Equal(value, db.Properties[key]) + }) + + s.Run("failure", func() { + dbName := fmt.Sprintf("dt_%s", s.randString(6)) + s.mock.EXPECT().DescribeDatabase(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once() + + _, err := s.client.DescribeDatabase(ctx, NewDescribeDatabaseOption(dbName)) + s.Error(err) + }) +} + +func (s *DatabaseSuite) TestAlterDatabaseProperties() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + s.Run("success", func() { + dbName := fmt.Sprintf("dt_%s", s.randString(6)) + key := fmt.Sprintf("key_%s", s.randString(4)) + value := s.randString(6) + s.mock.EXPECT().AlterDatabase(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, adr *milvuspb.AlterDatabaseRequest) (*commonpb.Status, error) { + s.Equal(dbName, adr.GetDbName()) + s.Len(adr.GetProperties(), 1) + return merr.Success(), nil + }).Once() + + err := s.client.AlterDatabaseProperies(ctx, NewAlterDatabasePropertiesOption(dbName).WithProperty(key, value)) + s.NoError(err) + }) + + s.Run("failure", func() { + dbName := fmt.Sprintf("dt_%s", s.randString(6)) + s.mock.EXPECT().AlterDatabase(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once() + + err := s.client.AlterDatabaseProperies(ctx, NewAlterDatabasePropertiesOption(dbName).WithProperty("key", "value")) + s.Error(err) + }) +} + +func (s *DatabaseSuite) TestDropDatabaseProperties() { + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + s.Run("success", func() { + dbName := fmt.Sprintf("dt_%s", s.randString(6)) + key := fmt.Sprintf("key_%s", s.randString(4)) + s.mock.EXPECT().AlterDatabase(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, adr *milvuspb.AlterDatabaseRequest) (*commonpb.Status, error) { + s.Equal([]string{key}, adr.GetDeleteKeys()) + return merr.Success(), nil + }).Once() + + err := s.client.DropDatabaseProperties(ctx, NewDropDatabasePropertiesOption(dbName, key)) + s.NoError(err) + }) + + s.Run("failure", func() { + dbName := fmt.Sprintf("dt_%s", s.randString(6)) + s.mock.EXPECT().AlterDatabase(mock.Anything, mock.Anything).Return(nil, merr.WrapErrServiceInternal("mocked")).Once() + + err := s.client.DropDatabaseProperties(ctx, NewDropDatabasePropertiesOption(dbName, "key")) + s.Error(err) + }) +} + func TestDatabase(t *testing.T) { suite.Run(t, new(DatabaseSuite)) }