From 558a1dc7afc1a6f0804a5e8b337cf1e0fe49b764 Mon Sep 17 00:00:00 2001 From: Wei Liu Date: Fri, 5 Jul 2024 14:25:02 +0800 Subject: [PATCH] enhance: enable setting properties during create database Signed-off-by: Wei Liu --- client/database.go | 13 +++++++++---- client/options.go | 16 +++++++++++++++- client/options_msg_base.go | 4 ++-- examples/database/database.go | 2 +- go.mod | 4 ++-- go.sum | 5 +++++ test/base/milvus_client.go | 4 ++-- 7 files changed, 36 insertions(+), 12 deletions(-) diff --git a/client/database.go b/client/database.go index 8c78d651..d05136fd 100644 --- a/client/database.go +++ b/client/database.go @@ -49,12 +49,17 @@ func (c *GrpcClient) CreateDatabase(ctx context.Context, dbName string, opts ... if c.config.hasFlags(disableDatabase) { return ErrFeatureNotSupported } - req := &milvuspb.CreateDatabaseRequest{ - DbName: dbName, + + opt := &createDatabaseOpt{} + for _, o := range opts { + o(opt) } - for _, opt := range opts { - opt(req) + + req := &milvuspb.CreateDatabaseRequest{ + DbName: dbName, + Properties: entity.MapKvPairs(opt.Properties), } + resp, err := c.Service.CreateDatabase(ctx, req) if err != nil { return err diff --git a/client/options.go b/client/options.go index 0254dd7b..9490fef4 100644 --- a/client/options.go +++ b/client/options.go @@ -307,7 +307,21 @@ type ReleaseCollectionOption func(*milvuspb.ReleaseCollectionRequest) type FlushOption func(*milvuspb.FlushRequest) -type CreateDatabaseOption func(*milvuspb.CreateDatabaseRequest) +type createDatabaseOpt struct { + Base *commonpb.MsgBase + Properties map[string]string +} + +type CreateDatabaseOption func(*createDatabaseOpt) + +func WithDatabaseProperty(key, value string) CreateDatabaseOption { + return func(opt *createDatabaseOpt) { + if opt.Properties == nil { + opt.Properties = make(map[string]string, 0) + } + opt.Properties[key] = value + } +} type DropDatabaseOption func(*milvuspb.DropDatabaseRequest) diff --git a/client/options_msg_base.go b/client/options_msg_base.go index cafc4434..4f14c171 100644 --- a/client/options_msg_base.go +++ b/client/options_msg_base.go @@ -36,8 +36,8 @@ func WithFlushMsgBase(msgBase *commonpb.MsgBase) FlushOption { } func WithCreateDatabaseMsgBase(msgBase *commonpb.MsgBase) CreateDatabaseOption { - return func(req *milvuspb.CreateDatabaseRequest) { - req.Base = msgBase + return func(opt *createDatabaseOpt) { + opt.Base = msgBase } } diff --git a/examples/database/database.go b/examples/database/database.go index da84f505..0f04ac8f 100644 --- a/examples/database/database.go +++ b/examples/database/database.go @@ -20,7 +20,7 @@ func main() { clientDefault := mustConnect(ctx, cfg) defer clientDefault.Close() createCollection(ctx, clientDefault, "col1") - if err := clientDefault.CreateDatabase(ctx, "db1"); err != nil { + if err := clientDefault.CreateDatabase(ctx, "db1", client.WithDatabaseProperty("key1", "value1")); err != nil { log.Fatalf("create db1 failed, %+v", err) } dbs, err := clientDefault.ListDatabases(ctx) diff --git a/go.mod b/go.mod index 980c22af..01d97cc2 100644 --- a/go.mod +++ b/go.mod @@ -7,7 +7,7 @@ require ( github.com/go-faker/faker/v4 v4.1.0 github.com/golang/protobuf v1.5.2 github.com/grpc-ecosystem/go-grpc-middleware v1.3.0 - github.com/milvus-io/milvus-proto/go-api/v2 v2.4.3 + github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240710021812-d4a17c9e8a2f github.com/stretchr/testify v1.8.1 github.com/tidwall/gjson v1.14.4 github.com/x448/float16 v0.8.4 @@ -33,6 +33,6 @@ require ( golang.org/x/sys v0.13.0 // indirect golang.org/x/text v0.13.0 // indirect google.golang.org/genproto v0.0.0-20220503193339-ba3ae3f07e29 // indirect - google.golang.org/protobuf v1.30.0 // indirect + google.golang.org/protobuf v1.33.0 // indirect gopkg.in/yaml.v3 v3.0.1 // indirect ) diff --git a/go.sum b/go.sum index 988bf02e..f4fba65c 100644 --- a/go.sum +++ b/go.sum @@ -159,8 +159,12 @@ github.com/mediocregopher/radix/v3 v3.4.2/go.mod h1:8FL3F6UQRXHXIBSPUs5h0RybMF8i github.com/microcosm-cc/bluemonday v1.0.2/go.mod h1:iVP4YcDBq+n/5fb23BhYFvIMq/leAFZyRl6bYmGDlGc= github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240430025921-135167be0694 h1:iub0yx8peGNtnb9n11iuWNmhIhIXw3xfZooIDcrfeU8= github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240430025921-135167be0694/go.mod h1:1OIl0v5PQeNxIJhCvY+K55CBUOYDZevw9g9380u1Wek= +github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240710021812-d4a17c9e8a2f h1:2A2f9ds2PVKSWkz9LVaXA/oqgdXeRF6l/k+d4SwbLho= +github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240710021812-d4a17c9e8a2f/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs= github.com/milvus-io/milvus-proto/go-api/v2 v2.4.3 h1:KUSaWVePVlHMIluAXf2qmNffI1CMlGFLLiP+4iy9014= github.com/milvus-io/milvus-proto/go-api/v2 v2.4.3/go.mod h1:1OIl0v5PQeNxIJhCvY+K55CBUOYDZevw9g9380u1Wek= +github.com/milvus-io/milvus-proto/go-api/v2 v2.4.6-0.20240705061601-7a658711dc06 h1:b4ulIU5bNSxScgZ4eNb19Gdr6Xhv5pqaC2SOugFogao= +github.com/milvus-io/milvus-proto/go-api/v2 v2.4.6-0.20240705061601-7a658711dc06/go.mod h1:1OIl0v5PQeNxIJhCvY+K55CBUOYDZevw9g9380u1Wek= github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0= github.com/mitchellh/mapstructure v1.1.2/go.mod h1:FVVH3fgwuzCH5S8UJGiWEs2h04kUh9fWfEaFds41c1Y= github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q= @@ -419,6 +423,7 @@ google.golang.org/protobuf v1.27.1/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQ google.golang.org/protobuf v1.28.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I= google.golang.org/protobuf v1.30.0 h1:kPPoIgf3TsEvrm0PFe15JQ+570QVxYzEvvHqChK+cng= google.golang.org/protobuf v1.30.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I= +google.golang.org/protobuf v1.33.0/go.mod h1:c6P6GXX6sHbq/GpV6MGZEdwhWPcYBgnhAHhKbcUYpos= gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20180628173108-788fd7840127/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0= gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 h1:YR8cESwS4TdDjEe65xsg0ogRM/Nc3DYOhEAlW+xobZo= diff --git a/test/base/milvus_client.go b/test/base/milvus_client.go index e3a41e89..4dda3f34 100644 --- a/test/base/milvus_client.go +++ b/test/base/milvus_client.go @@ -104,9 +104,9 @@ func (mc *MilvusClient) ListDatabases(ctx context.Context) ([]entity.Database, e } // CreateDatabase create database with the given name. -func (mc *MilvusClient) CreateDatabase(ctx context.Context, dbName string) error { +func (mc *MilvusClient) CreateDatabase(ctx context.Context, dbName string, opts ...client.CreateDatabaseOption) error { preRequest("CreateDatabase", ctx, dbName) - err := mc.mClient.CreateDatabase(ctx, dbName) + err := mc.mClient.CreateDatabase(ctx, dbName, opts...) postResponse("CreateDatabase", err) return err }