Repository: milvus-io/milvus-sdk-go
Pull request #273 (2022-05-13 03:31:07)

Add ReplicaNumber option in LoadCollection
Add ReplicaNumber option in LoadCollection

Related to #272 

Signed-off-by: Letian Jiang <letian.jiang@outlook.com>

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/273#" title="Author self-approved">letian-jiang</a>*
To complete the [pull request process](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process), please assign **sunby** after the PR has been reviewed.
You can assign the PR to them by writing `/assign @sunby` 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":["sunby"]} -->
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/273#pullrequestreview-971679698" title="LGTM">congqixia</a>*, *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/273#" title="Author self-approved">letian-jiang</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":[]} -->
congqixia: /kind feature

Patch:
diff --git a/client/client.go b/client/client.go
index 1a059e9..01e2c8b 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, opts ...CreateCollectionOpt) error
+	CreateCollection(ctx context.Context, schema *entity.Schema, shardsNum int32, opts ...CreateCollectionOption) error
 	// DescribeCollection describe collection meta
 	DescribeCollection(ctx context.Context, collName string) (*entity.Collection, error)
 	// DropCollection drop the specified collection
@@ -44,7 +44,7 @@ type Client interface {
 	// GetCollectionStatistics get collection statistics
 	GetCollectionStatistics(ctx context.Context, collName string) (map[string]string, error)
 	// LoadCollection load collection into memory
-	LoadCollection(ctx context.Context, collName string, async bool) error
+	LoadCollection(ctx context.Context, collName string, async bool, opts ...LoadCollectionOption) error
 	// ReleaseCollection release loaded collection
 	ReleaseCollection(ctx context.Context, collName string) error
 	// HasCollection check whether collection exists
@@ -88,6 +88,8 @@ type Client interface {
 
 	// -- basic operation --
 
+	// TODO: add get_replicas API in Collection
+
 	// Insert column-based data into collection, returns id column values
 	Insert(ctx context.Context, collName string, partitionName string, columns ...entity.Column) (entity.Column, error)
 	// Flush flush collection, specified
diff --git a/client/client_grpc_collection.go b/client/client_grpc_collection.go
index 270709b..1043cf5 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, opts ...CreateCollectionOpt) error {
+func (c *grpcClient) CreateCollection(ctx context.Context, collSchema *entity.Schema, shardNum int32, opts ...CreateCollectionOption) error {
 	if c.service == nil {
 		return ErrClientNotReady
 	}
@@ -141,7 +141,7 @@ func (c *grpcClient) CreateCollection(ctx context.Context, collSchema *entity.Sc
 	}
 	// apply options on request
 	for _, opt := range opts {
-		opt.OptCreateCollection(req)
+		opt(req)
 	}
 	resp, err := c.service.CreateCollection(ctx, req)
 	if err != nil {
@@ -327,7 +327,7 @@ func (c *grpcClient) ShowCollection(ctx context.Context, collName string) (*enti
 }
 
 // LoadCollection load collection into memory
-func (c *grpcClient) LoadCollection(ctx context.Context, collName string, async bool) error {
+func (c *grpcClient) LoadCollection(ctx context.Context, collName string, async bool, opts ...LoadCollectionOption) error {
 	if c.service == nil {
 		return ErrClientNotReady
 	}
@@ -338,7 +338,13 @@ func (c *grpcClient) LoadCollection(ctx context.Context, collName string, async
 
 	req := &server.LoadCollectionRequest{
 		CollectionName: collName,
+		ReplicaNumber:  1, // default replica number
 	}
+
+	for _, opt := range opts {
+		opt(req)
+	}
+
 	resp, err := c.service.LoadCollection(ctx, req)
 	if err != nil {
 		return err
diff --git a/client/client_grpc_options.go b/client/client_grpc_options.go
index 71e364f..cbf1332 100644
--- a/client/client_grpc_options.go
+++ b/client/client_grpc_options.go
@@ -16,22 +16,22 @@ import (
 	"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)
+// CreateCollectionOption is an option that is used to modify CreateCollectionRequest
+type CreateCollectionOption func(*server.CreateCollectionRequest)
 
-// OptCreateCollection implements CreateCollectionOpt
-func (f createCollectionOptFunc) OptCreateCollection(req *server.CreateCollectionRequest) {
-	f(req)
+// WithConsistencyLevel specifies a specific ConsistencyLevel, rather than using the default ReaderProperties.
+func WithConsistencyLevel(cl entity.ConsistencyLevel) CreateCollectionOption {
+	return func(req *server.CreateCollectionRequest) {
+		req.ConsistencyLevel = cl.CommonConsisencyLevel()
+	}
 }
 
-// WithConsistencyLevel returns a CreateCollectionOpt with ConsistencyLevel
-func CollectionWithConsistencyLevel(cl entity.ConsistencyLevel) CreateCollectionOpt {
-	return createCollectionOptFunc(func(req *server.CreateCollectionRequest) {
-		req.ConsistencyLevel = cl.CommonConsisencyLevel()
-	})
+// LoadCollectionOption is an option that is used to modify LoadCollectionRequest
+type LoadCollectionOption func(*server.LoadCollectionRequest)
+
+// WithReplicaNumber specifies a specific ReplicaNumber, rather than using the default ReplicaNumber.
+func WithReplicaNumber(rn int32) LoadCollectionOption {
+	return func(req *server.LoadCollectionRequest) {
+		req.ReplicaNumber = rn
+	}
 }


Test patch:
diff --git a/client/client_grpc_collection_test.go b/client/client_grpc_collection_test.go
index 7255e7b..6376992 100644
--- a/client/client_grpc_collection_test.go
+++ b/client/client_grpc_collection_test.go
@@ -120,6 +120,23 @@ func TestGrpcClientCreateCollection(t *testing.T) {
 		mock.delInjection(mCreateCollection)
 	})
 
+	t.Run("Test create with consistency level", func(t *testing.T) {
+		ds := defaultSchema()
+		shardsNum := int32(1)
+		mock.setInjection(mCreateCollection, func(ctx context.Context, raw proto.Message) (proto.Message, error) {
+			req, ok := raw.(*server.CreateCollectionRequest)
+			if !ok {
+				return &common.Status{ErrorCode: common.ErrorCode_IllegalArgument}, errors.New("illegal request type")
+			}
+			assert.Equal(t, testCollectionName, req.GetCollectionName())
+			assert.Equal(t, common.ConsistencyLevel_Eventually, req.GetConsistencyLevel())
+
+			return &common.Status{ErrorCode: common.ErrorCode_Success}, nil
+		})
+		assert.Nil(t, c.CreateCollection(ctx, ds, shardsNum, WithConsistencyLevel(entity.CL_EVENTUALLY)))
+		mock.delInjection(mCreateCollection)
+	})
+
 	t.Run("Test invalid schemas", func(t *testing.T) {
 		cases := []*entity.Schema{
 			// empty fields
@@ -363,6 +380,34 @@ func TestGrpcClientLoadCollection(t *testing.T) {
 		// remove injection
 		mock.delInjection(mShowCollections)
 	})
+	t.Run("Load default replica", func(t *testing.T) {
+		mock.setInjection(mLoadCollection, func(ctx context.Context, raw proto.Message) (proto.Message, error) {
+			req, ok := raw.(*server.LoadCollectionRequest)
+			if !ok {
+				return badRequestStatus()
+			}
+			assert.Equal(t, testDefaultReplicaNumber, req.GetReplicaNumber())
+			assert.Equal(t, testCollectionName, req.GetCollectionName())
+			return successStatus()
+		})
+		defer mock.delInjection(mLoadCollection)
+		assert.Nil(t, c.LoadCollection(ctx, testCollectionName, true))
+	})
+	t.Run("Load multiple replica", func(t *testing.T) {
+		mock.delInjection(mLoadCollection)
+
+		mock.setInjection(mLoadCollection, func(ctx context.Context, raw proto.Message) (proto.Message, error) {
+			req, ok := raw.(*server.LoadCollectionRequest)
+			if !ok {
+				return badRequestStatus()
+			}
+			assert.Equal(t, testMultiReplicaNumber, req.GetReplicaNumber())
+			assert.Equal(t, testCollectionName, req.GetCollectionName())
+			return successStatus()
+		})
+		defer mock.delInjection(mLoadCollection)
+		assert.Nil(t, c.LoadCollection(ctx, testCollectionName, true, WithReplicaNumber(testMultiReplicaNumber)))
+	})
 }
 
 func TestReleaseCollection(t *testing.T) {
diff --git a/client/client_grpc_options_test.go b/client/client_grpc_options_test.go
index aee71b3..ba527ea 100644
--- a/client/client_grpc_options_test.go
+++ b/client/client_grpc_options_test.go
@@ -12,60 +12,33 @@
 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"
+	"testing"
 )
 
-func TestCollectionWithConsistencyLevel(t *testing.T) {
-	opt := CollectionWithConsistencyLevel(entity.ConsistencyLevel(entity.CL_BOUNDED))
+func TestCreateCollectionWithConsistencyLevel(t *testing.T) {
+	opt := WithConsistencyLevel(entity.ConsistencyLevel(entity.CL_BOUNDED))
 	assert.NotNil(t, opt)
 	req := &server.CreateCollectionRequest{}
 
 	assert.NotPanics(t, func() {
-		opt.OptCreateCollection(req)
+		opt(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)
+func TestLoadCollectionWithReplicaNumber(t *testing.T) {
+	opt := WithReplicaNumber(testMultiReplicaNumber)
+	assert.NotNil(t, opt)
+	req := &server.LoadCollectionRequest{}
 
-	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()
+	assert.NotPanics(t, func() {
+		opt(req)
+	})
 
-	err = c.CreateCollection(ctx, ds, shardsNum, opt)
-	assert.NoError(t, err)
-	opt.AssertExpectations(t)
+	assert.Equal(t, testMultiReplicaNumber, req.GetReplicaNumber())
 }
diff --git a/client/client_mock_test.go b/client/client_mock_test.go
index 92639c9..6015247 100644
--- a/client/client_mock_test.go
+++ b/client/client_mock_test.go
@@ -65,10 +65,12 @@ func testClient(ctx context.Context, t *testing.T) Client {
 }
 
 const (
-	testCollectionName = `test_go_sdk`
-	testPrimaryField   = `int64`
-	testVectorField    = `vector`
-	testVectorDim      = 128
+	testCollectionName       = `test_go_sdk`
+	testPrimaryField         = `int64`
+	testVectorField          = `vector`
+	testVectorDim            = 128
+	testDefaultReplicaNumber = int32(1)
+	testMultiReplicaNumber   = int32(2)
 )
 
 func defaultSchema() *entity.Schema {