Repository: milvus-io/milvus-sdk-go (go)

Issue:
[Feature]: Support MultipleReplica in milvus go sdk
### Is there an existing issue for this?

- [X] I have searched the existing issues

### Is your feature request related to a problem? Please describe.

Milvus 2.1 allows users to load collection/partitions as multiple replicas to utilize extra querynodes' CPU and memory. With this ability, the overall qps and throughput can be scaled out with extra hardware.

Go SDK should support this feature. 

### Describe the solution you'd like.

* [x] Add ReplicaNumber option to specify replica number in collection loading #273
* [x] Add GetReplicas API to check information about replica groups, as well as their corresponding querynodes and shards #274

### Describe an alternate solution.

_No response_

### Anything else? (Additional Context)

_No response_

Pull request:
Add GetReplicas API for supporting memory replica
Add GetReplicas API for supporting memory replica
Close #272 

Signed-off-by: Letian Jiang <letian-jiang@zilliz.com>

Patch:
diff --git a/client/client.go b/client/client.go
index 01e2c8b..778bae9 100644
--- a/client/client.go
+++ b/client/client.go
@@ -57,6 +57,9 @@ type Client interface {
 	// AlterAlias changes collection alias to provided alias
 	AlterAlias(ctx context.Context, collName string, alias string) error
 
+	// GetReplicas gets the replica groups as well as their querynodes and shards information
+	GetReplicas(ctx context.Context, collName string) ([]*entity.ReplicaGroup, error)
+
 	// -- partition --
 
 	// CreatePartition create partition for collection
@@ -88,8 +91,6 @@ 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
@@ -106,7 +107,8 @@ type Client interface {
 	CalcDistance(ctx context.Context, collName string, partitions []string,
 		metricType entity.MetricType, opLeft, opRight entity.Column) (entity.Column, error)
 
-	// -- row basd apis --
+	// -- row based apis --
+
 	// CreateCollectionByRow create collection by row
 	CreateCollectionByRow(ctx context.Context, row entity.Row, shardNum int32) error
 	// InsertByRows insert by rows
diff --git a/client/client_grpc_collection.go b/client/client_grpc_collection.go
index 1043cf5..8ea331f 100644
--- a/client/client_grpc_collection.go
+++ b/client/client_grpc_collection.go
@@ -322,6 +322,7 @@ func (c *grpcClient) ShowCollection(ctx context.Context, collName string) (*enti
 		return nil, errors.New("response len not valid")
 	}
 	return &entity.Collection{
+		ID:     resp.CollectionIds[0],
 		Loaded: resp.InMemoryPercentages[0] == 100, // TODO silverxia, the percentage can be either 0 or 100
 	}, nil
 }
@@ -394,3 +395,46 @@ func (c *grpcClient) ReleaseCollection(ctx context.Context, collName string) err
 	}
 	return handleRespStatus(resp)
 }
+
+// GetReplicas gets the replica groups as well as their querynodes and shards information
+func (c *grpcClient) GetReplicas(ctx context.Context, collName string) ([]*entity.ReplicaGroup, error) {
+	if c.service == nil {
+		return nil, ErrClientNotReady
+	}
+	coll, err := c.ShowCollection(ctx, collName)
+	if err != nil {
+		return nil, err
+	}
+
+	req := &server.GetReplicasRequest{
+		CollectionID:   coll.ID,
+		WithShardNodes: true, // return nodes by default
+	}
+
+	resp, err := c.service.GetReplicas(ctx, req)
+	if err != nil {
+		return nil, err
+	}
+	if err = handleRespStatus(resp.GetStatus()); err != nil {
+		return nil, err
+	}
+
+	groups := make([]*entity.ReplicaGroup, 0, len(resp.GetReplicas()))
+	for _, rp := range resp.GetReplicas() {
+		group := &entity.ReplicaGroup{
+			ReplicaID:     rp.ReplicaID,
+			NodeIDs:       rp.NodeIds,
+			ShardReplicas: make([]*entity.ShardReplica, 0, len(rp.ShardReplicas)),
+		}
+		for _, s := range rp.ShardReplicas {
+			shard := &entity.ShardReplica{
+				LeaderID:      s.LeaderID,
+				NodesIDs:      s.NodeIds,
+				DmChannelName: s.DmChannelName,
+			}
+			group.ShardReplicas = append(group.ShardReplicas, shard)
+		}
+		groups = append(groups, group)
+	}
+	return groups, nil
+}
diff --git a/entity/collection.go b/entity/collection.go
index 5ff3fb0..d976082 100644
--- a/entity/collection.go
+++ b/entity/collection.go
@@ -28,3 +28,17 @@ type Partition struct {
 	Name   string // partition name
 	Loaded bool   // partition loaded
 }
+
+// ReplicaGroup represents a replica group
+type ReplicaGroup struct {
+	ReplicaID     int64
+	NodeIDs       []int64
+	ShardReplicas []*ShardReplica
+}
+
+// ShardReplica represents a shard in the ReplicaGroup
+type ShardReplica struct {
+	LeaderID      int64
+	NodesIDs      []int64
+	DmChannelName string
+}


Test patch:
diff --git a/client/client_grpc_collection_test.go b/client/client_grpc_collection_test.go
index 6376992..110a156 100644
--- a/client/client_grpc_collection_test.go
+++ b/client/client_grpc_collection_test.go
@@ -550,3 +550,96 @@ func TestGrpcClientGetCollectionStatistics(t *testing.T) {
 		}
 	}
 }
+
+func TestGrpcClientGetReplicas(t *testing.T) {
+	ctx := context.Background()
+	c := testClient(ctx, t)
+
+	replicaID := rand.Int63()
+	nodeIds := []int64{1, 2, 3, 4}
+	mock.setInjection(mHasCollection, hasCollectionDefault)
+	defer mock.delInjection(mHasCollection)
+
+	mock.setInjection(mShowCollections, func(_ context.Context, raw proto.Message) (proto.Message, error) {
+		s, err := successStatus()
+		resp := &server.ShowCollectionsResponse{
+			Status:              s,
+			CollectionIds:       []int64{testCollectionID},
+			CollectionNames:     []string{testCollectionName},
+			InMemoryPercentages: []int64{100},
+		}
+		return resp, err
+	})
+	defer mock.delInjection(mShowCollections)
+
+	mock.setInjection(mGetReplicas, func(ctx context.Context, raw proto.Message) (proto.Message, error) {
+		req, ok := raw.(*server.GetReplicasRequest)
+		resp := &server.GetReplicasResponse{}
+		if !ok {
+			s, err := badRequestStatus()
+			resp.Status = s
+			return resp, err
+		}
+
+		assert.Equal(t, testCollectionID, req.CollectionID)
+
+		s, err := successStatus()
+		resp.Status = s
+		resp.Replicas = []*server.ReplicaInfo{{
+			ReplicaID: replicaID,
+			ShardReplicas: []*server.ShardReplica{
+				{
+					LeaderID:      1,
+					DmChannelName: "DML_channel_v1",
+				},
+				{
+					LeaderID:   2,
+					LeaderAddr: "DML_channel_v2",
+				},
+			},
+			NodeIds: nodeIds,
+		}}
+		return resp, err
+	})
+
+	t.Run("get replicas normal", func(t *testing.T) {
+		groups, err := c.GetReplicas(ctx, testCollectionName)
+		assert.Nil(t, err)
+		assert.NotNil(t, groups)
+		fmt.Println(groups)
+		assert.Equal(t, 1, len(groups))
+
+		assert.Equal(t, replicaID, groups[0].ReplicaID)
+		assert.Equal(t, nodeIds, groups[0].NodeIDs)
+		assert.Equal(t, 2, len(groups[0].ShardReplicas))
+	})
+
+	t.Run("get replicas invalid name", func(t *testing.T) {
+		_, err := c.GetReplicas(ctx, "invalid name")
+		assert.Error(t, err)
+	})
+
+	t.Run("get replicas grpc error", func(t *testing.T) {
+		mock.setInjection(mGetReplicas, func(ctx context.Context, raw proto.Message) (proto.Message, error) {
+			return &server.GetReplicasResponse{}, errors.New("mocked grpc error")
+		})
+		_, err := c.GetReplicas(ctx, testCollectionName)
+		assert.Error(t, err)
+	})
+
+	t.Run("get replicas server error", func(t *testing.T) {
+		mock.setInjection(mGetReplicas, func(ctx context.Context, raw proto.Message) (proto.Message, error) {
+			return &server.GetReplicasResponse{
+				Status: &common.Status{
+					ErrorCode: common.ErrorCode_UnexpectedError,
+					Reason:    "service is not healthy",
+				},
+				Replicas: nil,
+			}, nil
+		})
+		_, err := c.GetReplicas(ctx, testCollectionName)
+		assert.Error(t, err)
+	})
+
+	mock.delInjection(mGetReplicas)
+}
diff --git a/client/client_grpc_control_test.go b/client/client_grpc_control_test.go
index f08bd57..01dbcac 100644
--- a/client/client_grpc_control_test.go
+++ b/client/client_grpc_control_test.go
@@ -35,11 +35,13 @@ func TestGrpcManualCompaction(t *testing.T) {
 	c := testClient(ctx, t)
 	defer c.Close()
 
-	collID := int64(1)
+	mock.setInjection(mHasCollection, hasCollectionDefault)
+	defer mock.delInjection(mHasCollection)
+
 	compactionID := int64(1001)
 	t.Run("normal manual compaction", func(t *testing.T) {
 		now := time.Now()
-		mock.setInjection(mDescribeCollection, describeCollectionInjection(t, collID, testCollectionName, defaultSchema()))
+		mock.setInjection(mDescribeCollection, describeCollectionInjection(t, testCollectionID, testCollectionName, defaultSchema()))
 		defer mock.delInjection(mDescribeCollection)
 		mock.setInjection(mManualCompaction, func(ctx context.Context, raw proto.Message) (proto.Message, error) {
 			req, ok := raw.(*server.ManualCompactionRequest)
@@ -47,7 +49,7 @@ func TestGrpcManualCompaction(t *testing.T) {
 				t.FailNow()
 			}
 
-			assert.Equal(t, collID, req.GetCollectionID())
+			assert.Equal(t, testCollectionID, req.GetCollectionID())
 			ts, _ := tso.ParseTS(req.GetTimetravel())
 			assert.True(t, ts.Sub(now) < time.Second)
 
@@ -83,7 +85,7 @@ func TestGrpcManualCompaction(t *testing.T) {
 	})
 
 	t.Run("compaction service error", func(t *testing.T) {
-		mock.setInjection(mDescribeCollection, describeCollectionInjection(t, collID, testCollectionName, defaultSchema()))
+		mock.setInjection(mDescribeCollection, describeCollectionInjection(t, testCollectionID, testCollectionName, defaultSchema()))
 		defer mock.delInjection(mDescribeCollection)
 		mock.setInjection(mManualCompaction, func(ctx context.Context, raw proto.Message) (proto.Message, error) {
 			resp := &server.ManualCompactionResponse{
diff --git a/client/client_mock_test.go b/client/client_mock_test.go
index 6015247..8d897c6 100644
--- a/client/client_mock_test.go
+++ b/client/client_mock_test.go
@@ -66,6 +66,7 @@ func testClient(ctx context.Context, t *testing.T) Client {
 
 const (
 	testCollectionName       = `test_go_sdk`
+	testCollectionID         = int64(789)
 	testPrimaryField         = `int64`
 	testVectorField          = `vector`
 	testVectorDim            = 128
@@ -127,6 +128,7 @@ const (
 	mCreateAlias             serviceMethod = 16
 	mDropAlias               serviceMethod = 17
 	mAlterAlias              serviceMethod = 18
+	mGetReplicas             serviceMethod = 19
 
 	mCreatePartition   serviceMethod = 9
 	mDropPartition     serviceMethod = 10
@@ -565,8 +567,16 @@ func (m *mockServer) GetCompactionStateWithPlans(ctx context.Context, req *serve
 	return resp, err
 }
 
-func (m *mockServer) GetReplicas(_ context.Context, _ *server.GetReplicasRequest) (*server.GetReplicasResponse, error) {
-	panic("not implemented") // TODO: Implement
+func (m *mockServer) GetReplicas(ctx context.Context, req *server.GetReplicasRequest) (*server.GetReplicasResponse, error) {
+	f := m.getInjection(mGetReplicas)
+	if f != nil {
+		r, err := f(ctx, req)
+		return r.(*server.GetReplicasResponse), err
+	}
+	resp := &server.GetReplicasResponse{}
+	s, err := successStatus()
+	resp.Status = s
+	return resp, err
 }
 
 // https://wiki.lfaidata.foundation/display/MIL/MEP+24+--+Support+bulk+load