Repository: milvus-io/milvus-sdk-go
Pull request #264 (2022-04-15 07:54:16)

Add ConsistencyLevel support in CreateCollection
Related issue #263
Add ConsistencyLevel in entity
Add CreateCollection options and ConsistencyLevelOption

/kind feature

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>

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

This pull-request has been approved by: *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/264#" title="Author self-approved">congqixia</a>*

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).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

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

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

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":[]} -->
codecov[bot]: # [Codecov](https://codecov.io/gh/milvus-io/milvus-sdk-go/pull/264?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) Report
> Merging [#264](https://codecov.io/gh/milvus-io/milvus-sdk-go/pull/264?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) (fbd556d) into [master](https://codecov.io/gh/milvus-io/milvus-sdk-go/commit/8ae70016002b47a2c342870d0a06e86c589664fb?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) (8ae7001) will **decrease** coverage by `0.09%`.
> The diff coverage is `85.71%`.

[![Impacted file tree graph](https://codecov.io/gh/milvus-io/milvus-sdk-go/pull/264/graphs/tree.svg?width=650&height=150&src=pr&token=WJK43PV390&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io)](https://codecov.io/gh/milvus-io/milvus-sdk-go/pull/264?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io)

```diff
@@            Coverage Diff             @@
##           master     #264      +/-   ##
==========================================
- Coverage   90.49%   90.39%   -0.10%     
==========================================
  Files          19       20       +1     
  Lines        2461     2447      -14     
==========================================
- Hits         2227     2212      -15     
- Misses        159      160       +1     
  Partials       75       75              
```


| [Impacted Files](https://codecov.io/gh/milvus-io/milvus-sdk-go/pull/264?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io) | Coverage Δ | |
|---|---|---|
| [client/client.go](https://codecov.io/gh/milvus-io/milvus-sdk-go/pull/264/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L2NsaWVudC5nbw==) | `100.00% <ø> (ø)` | |
| [entity/schema.go](https://codecov.io/gh/milvus-io/milvus-sdk-go/pull/264/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-ZW50aXR5L3NjaGVtYS5nbw==) | `91.81% <0.00%> (-0.85%)` | :arrow_down: |
| [client/client\_grpc\_collection.go](https://codecov.io/gh/milvus-io/milvus-sdk-go/pull/264/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L2NsaWVudF9ncnBjX2NvbGxlY3Rpb24uZ28=) | `84.00% <100.00%> (+0.16%)` | :arrow_up: |
| [client/client\_grpc\_options.go](https://codecov.io/gh/milvus-io/milvus-sdk-go/pull/264/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L2NsaWVudF9ncnBjX29wdGlvbnMuZ28=) | `100.00% <100.00%> (ø)` | |
| [client/client\_grpc\_alias.go](https://codecov.io/gh/milvus-io/milvus-sdk-go/pull/264/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=milvus-io#diff-Y2xpZW50L2NsaWVudF9ncnBjX2FsaWFzLmdv) | `100.00% <0.00%> (ø)` | |

sunby: /lgtm

Patch:
diff --git a/client/client.go b/client/client.go
index 629d096..1a059e9 100644
--- a/client/client.go
+++ b/client/client.go
@@ -36,7 +36,7 @@ type Client interface {
 	// ListCollections list collections from connection
 	ListCollections(ctx context.Context) ([]*entity.Collection, error)
 	// CreateCollection create collection using provided schema
-	CreateCollection(ctx context.Context, schema *entity.Schema, shardsNum int32) error
+	CreateCollection(ctx context.Context, schema *entity.Schema, shardsNum int32, opts ...CreateCollectionOpt) error
 	// DescribeCollection describe collection meta
 	DescribeCollection(ctx context.Context, collName string) (*entity.Collection, error)
 	// DropCollection drop the specified collection
diff --git a/client/client_grpc_collection.go b/client/client_grpc_collection.go
index 96f3a93..467548b 100644
--- a/client/client_grpc_collection.go
+++ b/client/client_grpc_collection.go
@@ -110,7 +110,7 @@ func (c *grpcClient) ListCollections(ctx context.Context) ([]*entity.Collection,
 }
 
 // CreateCollection create collection with specified schema
-func (c *grpcClient) CreateCollection(ctx context.Context, collSchema *entity.Schema, shardNum int32) error {
+func (c *grpcClient) CreateCollection(ctx context.Context, collSchema *entity.Schema, shardNum int32, opts ...CreateCollectionOpt) error {
 	if c.service == nil {
 		return ErrClientNotReady
 	}
@@ -135,6 +135,13 @@ func (c *grpcClient) CreateCollection(ctx context.Context, collSchema *entity.Sc
 		CollectionName: collSchema.CollectionName,
 		Schema:         bs,
 		ShardsNum:      shardNum,
+		// default consistency level is strong
+		// to be consistent with previous version
+		ConsistencyLevel: common.ConsistencyLevel_Strong,
+	}
+	// apply options on request
+	for _, opt := range opts {
+		opt.OptCreateCollection(req)
 	}
 	resp, err := c.service.CreateCollection(ctx, req)
 	if err != nil {
diff --git a/client/client_grpc_options.go b/client/client_grpc_options.go
new file mode 100644
index 0000000..71e364f
--- /dev/null
+++ b/client/client_grpc_options.go
@@ -0,0 +1,37 @@
+// Copyright (C) 2019-2021 Zilliz. All rights reserved.
+//
+// Licensed 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 client
+
+import (
+	"github.com/milvus-io/milvus-sdk-go/v2/entity"
+	"github.com/milvus-io/milvus-sdk-go/v2/internal/proto/server"
+)
+
+// CreateCollectionOpt is option interface for CreateCollection
+type CreateCollectionOpt interface {
+	OptCreateCollection(req *server.CreateCollectionRequest)
+}
+
+// createCollectionOptFunc shortcut function type
+type createCollectionOptFunc func(*server.CreateCollectionRequest)
+
+// OptCreateCollection implements CreateCollectionOpt
+func (f createCollectionOptFunc) OptCreateCollection(req *server.CreateCollectionRequest) {
+	f(req)
+}
+
+// WithConsistencyLevel returns a CreateCollectionOpt with ConsistencyLevel
+func CollectionWithConsistencyLevel(cl entity.ConsistencyLevel) CreateCollectionOpt {
+	return createCollectionOptFunc(func(req *server.CreateCollectionRequest) {
+		req.ConsistencyLevel = cl.CommonConsisencyLevel()
+	})
+}
diff --git a/entity/schema.go b/entity/schema.go
index 5036daf..d8b2f96 100644
--- a/entity/schema.go
+++ b/entity/schema.go
@@ -19,8 +19,25 @@ import (
 const (
 	// TYPE_PARAM_DIM is the const for field type param dimension
 	TYPE_PARAM_DIM = "dim"
+
+	// CL_STRONG strong consistency level
+	CL_STRONG ConsistencyLevel = ConsistencyLevel(common.ConsistencyLevel_Strong)
+	// CL_BOUNDED bounded consistency level with default tolerance of 5 seconds
+	CL_BOUNDED ConsistencyLevel = ConsistencyLevel(common.ConsistencyLevel_Bounded)
+	// CL_SESSION session consistency level
+	CL_SESSION ConsistencyLevel = ConsistencyLevel(common.ConsistencyLevel_Session)
+	// CL_EVENTUALLY eventually consistencty level
+	CL_EVENTUALLY ConsistencyLevel = ConsistencyLevel(common.ConsistencyLevel_Eventually)
 )
 
+// ConsistencyLevel enum type for collection Consistency Level
+type ConsistencyLevel common.ConsistencyLevel
+
+// CommonConsisencyLevel returns corresponding common.ConsisencyLevel
+func (cl ConsistencyLevel) CommonConsisencyLevel() common.ConsistencyLevel {
+	return common.ConsistencyLevel(cl)
+}
+
 // Schema represents schema info of collection in milvus
 type Schema struct {
 	CollectionName string


Test patch:
diff --git a/client/client_grpc_collection_test.go b/client/client_grpc_collection_test.go
index b045ec8..7255e7b 100644
--- a/client/client_grpc_collection_test.go
+++ b/client/client_grpc_collection_test.go
@@ -408,6 +408,32 @@ func TestGrpcClientHasCollection(t *testing.T) {
 	assert.True(t, has)
 }
 
+// return injection asserts collection name matchs
+// partition name request in partitionNames if flag is true
+func hasCollectionInjection(t *testing.T, mustIn bool, collNames ...string) func(context.Context, proto.Message) (proto.Message, error) {
+	return func(_ context.Context, raw proto.Message) (proto.Message, error) {
+		req, ok := raw.(*server.HasCollectionRequest)
+		resp := &server.BoolResponse{}
+		if !ok {
+			s, err := badRequestStatus()
+			resp.Status = s
+			return resp, err
+		}
+		if mustIn {
+			resp.Value = assert.Contains(t, collNames, req.GetCollectionName())
+		} else {
+			for _, pn := range collNames {
+				if pn == req.GetCollectionName() {
+					resp.Value = true
+				}
+			}
+		}
+		s, err := successStatus()
+		resp.Status = s
+		return resp, err
+	}
+}
+
 func describeCollectionInjection(t *testing.T, collID int64, collName string, sch *entity.Schema) func(_ context.Context, raw proto.Message) (proto.Message, error) {
 	return func(_ context.Context, raw proto.Message) (proto.Message, error) {
 		req, ok := raw.(*server.DescribeCollectionRequest)
diff --git a/client/client_grpc_options_test.go b/client/client_grpc_options_test.go
new file mode 100644
index 0000000..aee71b3
--- /dev/null
+++ b/client/client_grpc_options_test.go
@@ -0,0 +1,71 @@
+// Copyright (C) 2019-2021 Zilliz. All rights reserved.
+//
+// Licensed 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 client
+
+import (
+	"context"
+	"testing"
+
+	"github.com/golang/protobuf/proto"
+	"github.com/milvus-io/milvus-sdk-go/v2/entity"
+	"github.com/milvus-io/milvus-sdk-go/v2/internal/proto/common"
+	"github.com/milvus-io/milvus-sdk-go/v2/internal/proto/server"
+	"github.com/stretchr/testify/assert"
+	tmock "github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
+)
+
+func TestCollectionWithConsistencyLevel(t *testing.T) {
+	opt := CollectionWithConsistencyLevel(entity.ConsistencyLevel(entity.CL_BOUNDED))
+	assert.NotNil(t, opt)
+	req := &server.CreateCollectionRequest{}
+
+	assert.NotPanics(t, func() {
+		opt.OptCreateCollection(req)
+	})
+
+	assert.Equal(t, common.ConsistencyLevel_Bounded, req.GetConsistencyLevel())
+}
+
+type mockCreateCollectionOpt struct {
+	tmock.Mock
+}
+
+func (m *mockCreateCollectionOpt) OptCreateCollection(req *server.CreateCollectionRequest) {
+	m.Called(req)
+}
+
+func TestCreateCollection_WithConsistencyLevel(t *testing.T) {
+	ctx := context.Background()
+	c := testClient(ctx, t)
+	// default, all collection name returns false
+	mock.delInjection(mHasCollection)
+
+	ds := defaultSchema()
+	shardsNum := int32(1)
+	opt := &mockCreateCollectionOpt{}
+	sch := ds.ProtoMessage()
+	bs, err := proto.Marshal(sch)
+	require.NoError(t, err)
+	req := &server.CreateCollectionRequest{
+		DbName:           "", // reserved fields, not used for now
+		CollectionName:   ds.CollectionName,
+		Schema:           bs,
+		ShardsNum:        shardsNum,
+		ConsistencyLevel: common.ConsistencyLevel_Strong,
+	}
+	opt.On("OptCreateCollection", req).Return()
+
+	err = c.CreateCollection(ctx, ds, shardsNum, opt)
+	assert.NoError(t, err)
+	opt.AssertExpectations(t)
+}
diff --git a/client/client_test.go b/client/client_test.go
index e6a70ac..4992f80 100644
--- a/client/client_test.go
+++ b/client/client_test.go
@@ -3,6 +3,7 @@ package client
 import (
 	"context"
 	"errors"
+	"fmt"
 	"reflect"
 	"testing"
 
@@ -47,60 +48,70 @@ func TestGrpcClientNil(t *testing.T) {
 
 	for i := 0; i < tp.NumMethod(); i++ {
 		m := tp.Method(i)
-		mt := m.Type                                   // type of function
-		if m.Name == "Close" || m.Name == "Connect" || // skip connect & close
-			m.Name == "Search" || // type alias MetricType treated as string
-			m.Name == "CalcDistance" ||
-			m.Name == "ManualCompaction" || // time.Duration hard to detect in reflect
-			m.Name == "Insert" { // complex methods with ...
-			continue
-		}
-		ins := make([]reflect.Value, 0, mt.NumIn())
-		for j := 1; j < mt.NumIn(); j++ { // idx == 0, is the receiver v
-			if j == 1 {
-				//non-general solution, hard code context!
-				ins = append(ins, reflect.ValueOf(ctx))
-				continue
+		t.Run(fmt.Sprintf("TestGrpcClientNil_%s", m.Name), func(t *testing.T) {
+			mt := m.Type                                   // type of function
+			if m.Name == "Close" || m.Name == "Connect" || // skip connect & close
+				m.Name == "Search" || // type alias MetricType treated as string
+				m.Name == "CalcDistance" ||
+				m.Name == "ManualCompaction" || // time.Duration hard to detect in reflect
+				m.Name == "Insert" { // complex methods with ...
+				t.Skip("method", m.Name, "skipped")
 			}
-			inT := mt.In(j)
-
-			switch inT.Kind() {
-			case reflect.String: // pass empty
-				ins = append(ins, reflect.ValueOf(""))
-			case reflect.Int:
-				ins = append(ins, reflect.ValueOf(0))
-			case reflect.Int64:
-				ins = append(ins, reflect.ValueOf(int64(0)))
-			case reflect.Bool:
-				ins = append(ins, reflect.ValueOf(false))
-			case reflect.Interface:
-				idxType := reflect.TypeOf((*entity.Index)(nil)).Elem()
-				rowType := reflect.TypeOf((*entity.Row)(nil)).Elem()
-				colType := reflect.TypeOf((*entity.Column)(nil)).Elem()
-				switch {
-				case inT.Implements(idxType):
-					ins = append(ins, reflect.ValueOf(entity.NewFlatIndex("flat_index", entity.L2)))
-				case inT.Implements(rowType):
-					ins = append(ins, reflect.ValueOf(&ValidStruct{}))
-				case inT.Implements(colType):
-					ins = append(ins, reflect.ValueOf(entity.NewColumnInt64("id", []int64{})))
+			ins := make([]reflect.Value, 0, mt.NumIn())
+			for j := 1; j < mt.NumIn(); j++ { // idx == 0, is the receiver v
+				if j == 1 {
+					//non-general solution, hard code context!
+					ins = append(ins, reflect.ValueOf(ctx))
+					continue
+				}
+				if mt.IsVariadic() {
+					// Variadic function, skip last parameter
+					// func m (arg1 interface, opts ... options)
+					if j == mt.NumIn()-1 {
+						continue
+					}
+				}
+				inT := mt.In(j)
+
+				switch inT.Kind() {
+				case reflect.String: // pass empty
+					ins = append(ins, reflect.ValueOf(""))
+				case reflect.Int:
+					ins = append(ins, reflect.ValueOf(0))
+				case reflect.Int64:
+					ins = append(ins, reflect.ValueOf(int64(0)))
+				case reflect.Bool:
+					ins = append(ins, reflect.ValueOf(false))
+				case reflect.Interface:
+					idxType := reflect.TypeOf((*entity.Index)(nil)).Elem()
+					rowType := reflect.TypeOf((*entity.Row)(nil)).Elem()
+					colType := reflect.TypeOf((*entity.Column)(nil)).Elem()
+					switch {
+					case inT.Implements(idxType):
+						ins = append(ins, reflect.ValueOf(entity.NewFlatIndex("flat_index", entity.L2)))
+					case inT.Implements(rowType):
+						ins = append(ins, reflect.ValueOf(&ValidStruct{}))
+					case inT.Implements(colType):
+						ins = append(ins, reflect.ValueOf(entity.NewColumnInt64("id", []int64{})))
+					}
+				default:
+					ins = append(ins, reflect.Zero(inT))
 				}
-			default:
-				ins = append(ins, reflect.Zero(inT))
 			}
-		}
-		outs := v.MethodByName(m.Name).Call(ins)
-		assert.True(t, len(outs) > 0)
-		assert.EqualValues(t, ErrClientNotReady, outs[len(outs)-1].Interface())
+			outs := v.MethodByName(m.Name).Call(ins)
+			assert.True(t, len(outs) > 0)
+			assert.EqualValues(t, ErrClientNotReady, outs[len(outs)-1].Interface())
 
-		// ctx done
+			// ctx done
 
-		if len(ins) > 0 { // with context param
-			ins[0] = reflect.ValueOf(ctxDone)
-			outs := v2.MethodByName(m.Name).Call(ins)
-			assert.True(t, len(outs) > 0)
-			assert.False(t, outs[len(outs)-1].IsNil())
-		}
+			if len(ins) > 0 { // with context param
+				ins[0] = reflect.ValueOf(ctxDone)
+				outs := v2.MethodByName(m.Name).Call(ins)
+				assert.True(t, len(outs) > 0)
+				assert.False(t, outs[len(outs)-1].IsNil())
+			}
+
+		})
 	}
 }
 
diff --git a/entity/schema_test.go b/entity/schema_test.go
index 9c838a7..48e5a6a 100644
--- a/entity/schema_test.go
+++ b/entity/schema_test.go
@@ -3,9 +3,22 @@ package entity
 import (
 	"testing"
 
+	"github.com/milvus-io/milvus-sdk-go/v2/internal/proto/common"
 	"github.com/stretchr/testify/assert"
 )
 
+func TestCL_CommonCL(t *testing.T) {
+	cls := []ConsistencyLevel{
+		CL_STRONG,
+		CL_BOUNDED,
+		CL_SESSION,
+		CL_EVENTUALLY,
+	}
+	for _, cl := range cls {
+		assert.EqualValues(t, common.ConsistencyLevel(cl), cl.CommonConsisencyLevel())
+	}
+}
+
 func TestFieldSchema(t *testing.T) {
 	fields := []*Field{
 		{