enhance: [GoSDK] Add alter collection API & expose options (#37365)

Related to #31293

This PR:

- Add `AlterCollection` API for collection property modification
- Expose hidden or missing option methods

---------

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>
This commit is contained in:
congqixia 2024-11-06 14:36:25 +08:00 committed by GitHub
parent b3de4b0edd
commit c83b93946e
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
6 changed files with 111 additions and 22 deletions

View File

@ -138,3 +138,12 @@ func (c *Client) RenameCollection(ctx context.Context, option RenameCollectionOp
return merr.CheckRPCCall(resp, err)
})
}
func (c *Client) AlterCollection(ctx context.Context, option AlterCollectionOption, 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)
})
}

View File

@ -17,6 +17,8 @@
package client
import (
"fmt"
"google.golang.org/protobuf/proto"
"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
@ -56,6 +58,8 @@ type createCollectionOption struct {
// partition key
numPartitions int64
indexOptions []CreateIndexOption
// is fast create collection
isFast bool
// fast creation with index
@ -83,6 +87,21 @@ func (opt *createCollectionOption) WithVarcharPK(varcharPK bool, maxLen int) *cr
return opt
}
func (opt *createCollectionOption) WithIndexOptions(indexOpts ...CreateIndexOption) *createCollectionOption {
opt.indexOptions = append(opt.indexOptions, indexOpts...)
return opt
}
func (opt *createCollectionOption) WithProperty(key string, value any) *createCollectionOption {
opt.properties[key] = fmt.Sprintf("%v", value)
return opt
}
func (opt *createCollectionOption) WithConsistencyLevel(cl entity.ConsistencyLevel) *createCollectionOption {
opt.consistencyLevel = cl
return opt
}
func (opt *createCollectionOption) Request() *milvuspb.CreateCollectionRequest {
// fast create collection
if opt.isFast {
@ -103,6 +122,7 @@ func (opt *createCollectionOption) Request() *milvuspb.CreateCollectionRequest {
var schemaBytes []byte
if opt.schema != nil {
opt.schema.WithName(opt.name)
schemaProto := opt.schema.ProtoMessage()
schemaBytes, _ = proto.Marshal(schemaProto)
}
@ -144,6 +164,7 @@ func SimpleCreateCollectionOptions(name string, dim int64) *createCollectionOpti
dim: dim,
enabledDynamicSchema: true,
consistencyLevel: entity.DefaultConsistencyLevel,
properties: make(map[string]string),
isFast: true,
metricType: entity.COSINE,
@ -157,6 +178,7 @@ func NewCreateCollectionOption(name string, collectionSchema *entity.Schema) *cr
shardNum: 1,
schema: collectionSchema,
consistencyLevel: entity.DefaultConsistencyLevel,
properties: make(map[string]string),
metricType: entity.COSINE,
}
@ -263,3 +285,28 @@ func NewRenameCollectionOption(oldName, newName string) *renameCollectionOption
newCollectionName: newName,
}
}
type AlterCollectionOption interface {
Request() *milvuspb.AlterCollectionRequest
}
type alterCollectionOption struct {
collectionName string
properties map[string]string
}
func (opt *alterCollectionOption) WithProperty(key string, value any) *alterCollectionOption {
opt.properties[key] = fmt.Sprintf("%v", value)
return opt
}
func (opt *alterCollectionOption) 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)}
}

View File

@ -30,6 +30,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/client/v2/entity"
"github.com/milvus-io/milvus/pkg/common"
"github.com/milvus-io/milvus/pkg/util/merr"
)
@ -117,11 +118,20 @@ func (s *CollectionSuite) TestCreateCollectionOptions() {
WithField(entity.NewField().WithName("int64").WithDataType(entity.FieldTypeInt64).WithIsPrimaryKey(true)).
WithField(entity.NewField().WithName("vector").WithDim(128).WithDataType(entity.FieldTypeFloatVector))
opt = NewCreateCollectionOption(collectionName, schema).WithShardNum(2)
opt = NewCreateCollectionOption(collectionName, schema).
WithShardNum(2).
WithConsistencyLevel(entity.ClEventually).
WithProperty(common.CollectionTTLConfigKey, 86400)
req = opt.Request()
s.Equal(collectionName, req.GetCollectionName())
s.EqualValues(2, req.GetShardsNum())
s.EqualValues(commonpb.ConsistencyLevel_Eventually, req.GetConsistencyLevel())
if s.Len(req.GetProperties(), 1) {
kv := req.GetProperties()[0]
s.Equal(common.CollectionTTLConfigKey, kv.GetKey())
s.Equal("86400", kv.GetValue())
}
collSchema = &schemapb.CollectionSchema{}
err = proto.Unmarshal(req.GetSchema(), collSchema)
@ -274,6 +284,37 @@ func (s *CollectionSuite) TestRenameCollection() {
})
}
func (s *CollectionSuite) TestAlterCollection() {
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
collName := fmt.Sprintf("test_collection_%s", s.randString(6))
key := s.randString(6)
value := s.randString(6)
s.Run("success", func() {
s.mock.EXPECT().AlterCollection(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, acr *milvuspb.AlterCollectionRequest) (*commonpb.Status, error) {
s.Equal(collName, acr.GetCollectionName())
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.AlterCollection(ctx, NewAlterCollectionOption(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))
s.Error(err)
})
}
func TestCollection(t *testing.T) {
suite.Run(t, new(CollectionSuite))
}

View File

@ -142,6 +142,11 @@ func (opt *searchOption) WithPartitions(partitionNames ...string) *searchOption
return opt
}
func (opt *searchOption) WithGroupByField(groupByField string) *searchOption {
opt.request.groupByField = groupByField
return opt
}
func NewSearchOption(collectionName string, limit int, vectors []entity.Vector) *searchOption {
return &searchOption{
collectionName: collectionName,

View File

@ -45,6 +45,7 @@ func (s *ReadSuite) TestSearch() {
s.mock.EXPECT().Search(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, sr *milvuspb.SearchRequest) (*milvuspb.SearchResults, error) {
s.Equal(collectionName, sr.GetCollectionName())
s.ElementsMatch([]string{partitionName}, sr.GetPartitionNames())
// s.Equal(s)
return &milvuspb.SearchResults{
Status: merr.Success(),
@ -71,7 +72,7 @@ func (s *ReadSuite) TestSearch() {
entity.FloatVector(lo.RepeatBy(128, func(_ int) float32 {
return rand.Float32()
})),
}).WithPartitions(partitionName))
}).WithPartitions(partitionName).WithGroupByField("group_by"))
s.NoError(err)
})

View File

@ -445,35 +445,22 @@ func TestCreateCollectionWithInvalidCollectionName(t *testing.T) {
// connect
ctx := hp.CreateContext(t, time.Second*common.DefaultTimeout)
mc := createDefaultMilvusClient(ctx, t)
collName := common.GenRandomString(prefix, 6)
// create collection and schema no name
schema := genDefaultSchema()
err2 := mc.CreateCollection(ctx, client.NewCreateCollectionOption(collName, schema))
err2 := mc.CreateCollection(ctx, client.NewCreateCollectionOption("", schema))
common.CheckErr(t, err2, false, "collection name should not be empty")
// create collection with invalid schema name
for _, invalidName := range common.GenInvalidNames() {
log.Debug("TestCreateCollectionWithInvalidCollectionName", zap.String("collectionName", invalidName))
// schema has invalid name
schema.WithName(invalidName)
err := mc.CreateCollection(ctx, client.NewCreateCollectionOption(collName, schema))
err := mc.CreateCollection(ctx, client.NewCreateCollectionOption(invalidName, schema))
common.CheckErr(t, err, false, "collection name should not be empty",
"the first character of a collection name must be an underscore or letter",
"collection name can only contain numbers, letters and underscores",
fmt.Sprintf("the length of a collection name must be less than %d characters", common.MaxCollectionNameLen))
// collection option has invalid name
schema.WithName(collName)
err2 := mc.CreateCollection(ctx, client.NewCreateCollectionOption(invalidName, schema))
common.CheckErr(t, err2, false, "collection name matches schema name")
}
// collection name not equal to schema name
schema.WithName(collName)
err3 := mc.CreateCollection(ctx, client.NewCreateCollectionOption(common.GenRandomString("pre", 4), schema))
common.CheckErr(t, err3, false, "collection name matches schema name")
}
// create collection missing pk field or vector field
@ -937,11 +924,10 @@ func TestCreateCollectionInvalid(t *testing.T) {
vecField := entity.NewField().WithName("vec").WithDataType(entity.FieldTypeFloatVector).WithDim(8)
mSchemaErrs := []mSchemaErr{
{schema: nil, errMsg: "schema does not contain vector field"},
{schema: entity.NewSchema().WithField(vecField), errMsg: "collection name should not be empty"}, // no collection name
{schema: entity.NewSchema().WithName("aaa").WithField(vecField), errMsg: "primary key is not specified"}, // no pk field
{schema: entity.NewSchema().WithName("aaa").WithField(vecField).WithField(entity.NewField()), errMsg: "primary key is not specified"},
{schema: entity.NewSchema().WithName("aaa").WithField(vecField).WithField(entity.NewField().WithIsPrimaryKey(true)), errMsg: "the data type of primary key should be Int64 or VarChar"},
{schema: entity.NewSchema().WithName("aaa").WithField(vecField).WithField(entity.NewField().WithIsPrimaryKey(true).WithDataType(entity.FieldTypeVarChar)), errMsg: "field name should not be empty"},
{schema: entity.NewSchema().WithField(vecField), errMsg: "primary key is not specified"}, // no pk field
{schema: entity.NewSchema().WithField(vecField).WithField(entity.NewField()), errMsg: "primary key is not specified"},
{schema: entity.NewSchema().WithField(vecField).WithField(entity.NewField().WithIsPrimaryKey(true)), errMsg: "the data type of primary key should be Int64 or VarChar"},
{schema: entity.NewSchema().WithField(vecField).WithField(entity.NewField().WithIsPrimaryKey(true).WithDataType(entity.FieldTypeVarChar)), errMsg: "field name should not be empty"},
}
for _, mSchema := range mSchemaErrs {
err := mc.CreateCollection(ctx, client.NewCreateCollectionOption(collName, mSchema.schema))