enhance: enable setting properties during create database

Signed-off-by: Wei Liu <[email protected]>
milvus-io · Jul 16, 2024 · 9d07bc8 · 9d07bc8
1 parent 869db52
commit 9d07bc8
Show file tree

Hide file tree

Showing 7 changed files with 35 additions and 18 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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.4.6
 	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
@@ -157,10 +157,8 @@ github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27k
 github.com/mattn/goveralls v0.0.2/go.mod h1:8d1ZMHsd7fW6IRPKQh46F2WRpyib5/X4FOpevwGNQEw=
 github.com/mediocregopher/radix/v3 v3.4.2/go.mod h1:8FL3F6UQRXHXIBSPUs5h0RybMF8i4n7wVopoX3x7Bv8=
 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.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 h1:rSkwp5Mg/7KBSUqXcrPBUgTQGZNdvYWEKB+rHo9YJtk=
+github.com/milvus-io/milvus-proto/go-api/v2 v2.4.6/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=
@@ -417,8 +415,8 @@ google.golang.org/protobuf v1.26.0-rc.1/go.mod h1:jlhhOSvTdKEhbULTjvd4ARK9grFBp0
 google.golang.org/protobuf v1.26.0/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
 google.golang.org/protobuf v1.27.1/go.mod h1:9q0QmTI4eRPtz6boOQmLYwt+qCgq0jsYwAQnmE0givc=
 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 h1:uNO2rsAINq/JlFpSdYEKIZ0uKD/R9cpdv0T+yoGwGmI=
+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
@@ -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
 }