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 <[email protected]>
milvus-io · Dec 27, 2024 · c39452f · c39452f
1 parent 2557e3f
commit c39452f
Show file tree

Hide file tree

Showing 8 changed files with 363 additions and 12 deletions.
diff --git a/client/entity/database.go 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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,
+	}
+}