Repository: milvus-io/milvus-sdk-go
Pull request #787 (2024-07-16 03:29:48)

feat: add support for partition key isolation property
issue: https://github.com/milvus-io/milvus/issues/34332

sre-ci-robot: [APPROVALNOTIFIER] This PR is **NOT APPROVED**

This pull-request has been approved by: *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/787#" title="Author self-approved">PwzXxm</a>*
To complete the [pull request process](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process), please assign **congqixia** after the PR has been reviewed.
You can assign the PR to them by writing `/assign @congqixia` in a comment when ready.

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=milvus-io%2Fmilvus-sdk-go).

<details open>
Needs approval from an approver in each of these files:

- **[OWNERS](https://github.com/milvus-io/milvus-sdk-go/blob/master/OWNERS)**

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":["congqixia"]} -->

Patch:
diff --git a/client/database.go b/client/database.go
index 8c78d65..d05136f 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 0254dd7..9490fef 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 cafc443..4f14c17 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/entity/collection_attr.go b/entity/collection_attr.go
index 7ada438..0e6c27f 100644
--- a/entity/collection_attr.go
+++ b/entity/collection_attr.go
@@ -22,10 +22,11 @@ const (
 	// cakTTL const for collection attribute key TTL in seconds.
 	cakTTL = `collection.ttl.seconds`
 	// cakAutoCompaction const for collection attribute key autom compaction enabled.
-	cakAutoCompaction      = `collection.autocompaction.enabled`
-	akMmap                 = "mmap.enabled"
-	databaseReplica        = "database.replica"
-	databaseResourceGroups = "database.resource_groups"
+	cakAutoCompaction       = `collection.autocompaction.enabled`
+	akMmap                  = "mmap.enabled"
+	databaseReplica         = "database.replica"
+	databaseResourceGroups  = "database.resource_groups"
+	caPartitionKeyIsolation = "partitionkey.isolation"
 )
 
 // CollectionAttribute is the interface for altering collection attributes.
@@ -48,6 +49,10 @@ type ttlCollAttr struct {
 	collAttrBase
 }
 
+type partitionKeyIsolationCollAttr struct {
+	collAttrBase
+}
+
 // Valid implements CollectionAttribute.
 // checks ttl seconds is valid positive integer.
 func (ca collAttrBase) Valid() error {
@@ -114,6 +119,21 @@ func (ca mmapAttr) Valid() error {
 	return nil
 }
 
+func CollectionPartitionKeyIsolation(enabled bool) partitionKeyIsolationCollAttr {
+	ca := partitionKeyIsolationCollAttr{}
+	ca.key = caPartitionKeyIsolation
+	ca.value = strconv.FormatBool(enabled)
+	return ca
+}
+
+func (ca partitionKeyIsolationCollAttr) Valid() error {
+	_, err := strconv.ParseBool(ca.value)
+	if err != nil {
+		return errors.Wrap(err, "PartitionKeyIsolation setting is not a valid boolean")
+	}
+	return nil
+}
+
 type DatabaseAttribute = CollectionAttribute
 
 type databaseAttrBase = collAttrBase
diff --git a/examples/database/database.go b/examples/database/database.go
index da84f50..0f04ac8 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 980c22a..5164c9b 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.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
index 988bf02..71106d8 100644
--- 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=


Test patch:
diff --git a/client/collection_test.go b/client/collection_test.go
index d0021e9..c3a30ce 100644
--- a/client/collection_test.go
+++ b/client/collection_test.go
@@ -155,6 +155,43 @@ func (s *CollectionSuite) TestCreateCollection() {
 		s.NoError(err)
 	})
 
+	s.Run("create_with_partition_key_isolation", func() {
+		ds := defaultSchema()
+		partitionKeyFieldName := "partitionKeyField"
+		partitionKeyField := entity.NewField().WithName(partitionKeyFieldName).WithDataType(entity.FieldTypeInt64).WithIsPartitionKey(true)
+		ds.WithField(partitionKeyField)
+
+		shardsNum := int32(1)
+		expectedIsoStr := "true"
+		IsoKeystr := "partitionkey.isolation"
+		defer s.resetMock()
+		s.mock.EXPECT().CreateCollection(mock.Anything, mock.AnythingOfType("*milvuspb.CreateCollectionRequest")).
+			Run(func(ctx context.Context, req *milvuspb.CreateCollectionRequest) {
+				s.Equal(testCollectionName, req.GetCollectionName())
+				sschema := &schemapb.CollectionSchema{}
+				s.Require().NoError(proto.Unmarshal(req.GetSchema(), sschema))
+				s.Require().Equal(len(ds.Fields), len(sschema.Fields))
+				for idx, fieldSchema := range ds.Fields {
+					s.Equal(fieldSchema.Name, sschema.GetFields()[idx].GetName())
+					s.Equal(fieldSchema.PrimaryKey, sschema.GetFields()[idx].GetIsPrimaryKey())
+					s.Equal(fieldSchema.AutoID, sschema.GetFields()[idx].GetAutoID())
+					s.EqualValues(fieldSchema.DataType, sschema.GetFields()[idx].GetDataType())
+				}
+				s.Equal(shardsNum, req.GetShardsNum())
+				s.Equal(commonpb.ConsistencyLevel_Eventually, req.GetConsistencyLevel())
+				propsMp := make(map[string]string)
+				for _, prop := range req.GetProperties() {
+					propsMp[prop.Key] = prop.Value
+				}
+				s.Equal(expectedIsoStr, propsMp[IsoKeystr])
+			}).
+			Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil)
+		s.mock.EXPECT().HasCollection(mock.Anything, &milvuspb.HasCollectionRequest{CollectionName: testCollectionName}).Return(&milvuspb.BoolResponse{Status: &commonpb.Status{}, Value: false}, nil)
+
+		err := c.CreateCollection(ctx, ds, shardsNum, WithConsistencyLevel(entity.ClEventually), WithCollectionProperty(IsoKeystr, expectedIsoStr))
+		s.NoError(err)
+	})
+
 	s.Run("invalid_schemas", func() {
 		type testCase struct {
 			name   string
@@ -439,6 +476,9 @@ func (s *CollectionSuite) TestAlterCollection() {
 
 		err := c.AlterCollection(ctx, testCollectionName, entity.CollectionTTL(100000))
 		s.NoError(err)
+
+		err = c.AlterCollection(ctx, testCollectionName, entity.CollectionPartitionKeyIsolation(true))
+		s.NoError(err)
 	})
 
 	s.Run("collection_not_exist", func() {
diff --git a/entity/collection_attr_test.go b/entity/collection_attr_test.go
index b4f42e4..5ac4e7a 100644
--- a/entity/collection_attr_test.go
+++ b/entity/collection_attr_test.go
@@ -130,3 +130,52 @@ func (s *CollectionAutoCompactionSuite) TestCollectionAutoCompactionEnabled() {
 func TestCollectionAutoCompaction(t *testing.T) {
 	suite.Run(t, new(CollectionAutoCompactionSuite))
 }
+
+type CollectionPartitionKeyIsolationSuite struct {
+	suite.Suite
+}
+
+func (s *CollectionPartitionKeyIsolationSuite) TestValid() {
+	type testCase struct {
+		input     string
+		expectErr bool
+	}
+
+	cases := []testCase{
+		{input: "a", expectErr: true},
+		{input: "true", expectErr: false},
+		{input: "false", expectErr: false},
+		{input: "", expectErr: true},
+	}
+
+	for _, tc := range cases {
+		s.Run(tc.input, func() {
+			ca := partitionKeyIsolationCollAttr{}
+			ca.value = tc.input
+			err := ca.Valid()
+			if tc.expectErr {
+				s.Error(err)
+			} else {
+				s.NoError(err)
+			}
+		})
+	}
+}
+
+func (s *CollectionPartitionKeyIsolationSuite) TestCollectionAutoCompactionEnabled() {
+
+	cases := []bool{true, false}
+
+	for _, tc := range cases {
+		s.Run(fmt.Sprintf("%v", tc), func() {
+			ca := CollectionPartitionKeyIsolation(tc)
+			key, value := ca.KeyValue()
+			s.Equal(caPartitionKeyIsolation, key)
+			s.Equal(strconv.FormatBool(tc), value)
+		})
+	}
+}
+
+func TestCollectionPartitionKeyIsolationAttr(t *testing.T) {
+	suite.Run(t, new(CollectionPartitionKeyIsolationSuite))
+}
diff --git a/test/base/milvus_client.go b/test/base/milvus_client.go
index e3a41e8..4dda3f3 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
 }