Support describeReplicas() interface (#1317)

Signed-off-by: yhmo <yihua.mo@zilliz.com>

sdk-core/src/main/java/io/milvus/v2/client/MilvusClientV2.java

@@ -378,6 +378,15 @@ public class MilvusClientV2 {
         return rpcUtils.retry(()->collectionService.getLoadState(this.getRpcStub(), request));
     }
 
+    /**
+     * Get information of all replicas from a collection.
+     *
+     * @param request {@link DescribeReplicasReq}
+     */
+    public DescribeReplicasResp describeReplicas(DescribeReplicasReq request) {
+        return rpcUtils.retry(()->collectionService.describeReplicas(this.getRpcStub(), request));
+    }
+
     /////////////////////////////////////////////////////////////////////////////////////////////
     // Index Operations
     /////////////////////////////////////////////////////////////////////////////////////////////

sdk-core/src/main/java/io/milvus/v2/service/collection/CollectionService.java

@@ -358,6 +358,55 @@ public class CollectionService extends BaseService {
                 .build();
     }
 
+    public DescribeReplicasResp describeReplicas(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
+                                                 DescribeReplicasReq request) {
+        if (StringUtils.isEmpty(request.getCollectionName())) {
+            throw new MilvusClientException(ErrorCode.INVALID_PARAMS, "Invalid collection name");
+        }
+
+        String title = String.format("DescribeReplicas collectionName:%s", request.getCollectionName());
+
+        GetReplicasRequest.Builder requestBuilder = GetReplicasRequest.newBuilder()
+                .setCollectionName(request.getCollectionName())
+                .setWithShardNodes(true);
+
+        if (StringUtils.isNotEmpty(request.getDatabaseName())) {
+            requestBuilder.setDbName(request.getDatabaseName());
+        }
+
+        GetReplicasResponse response = blockingStub.getReplicas(requestBuilder.build());
+        rpcUtils.handleResponse(title, response.getStatus());
+
+        List<ReplicaInfo> replicas = new ArrayList<>();
+        List<io.milvus.grpc.ReplicaInfo> rpcReplicas = response.getReplicasList();
+        for (io.milvus.grpc.ReplicaInfo info : rpcReplicas) {
+            List<ShardReplica> shardReplicas = new ArrayList<>();
+            List<io.milvus.grpc.ShardReplica> rpcShardReplicas = info.getShardReplicasList();
+            for (io.milvus.grpc.ShardReplica shardReplica : rpcShardReplicas) {
+                shardReplicas.add(ShardReplica.builder()
+                        .leaderID(shardReplica.getLeaderID())
+                        .leaderAddress(shardReplica.getLeaderAddr())
+                        .channelName(shardReplica.getDmChannelName())
+                        .nodeIDs(shardReplica.getNodeIdsList())
+                        .build());
+            }
+
+            replicas.add(ReplicaInfo.builder()
+                    .replicaID(info.getReplicaID())
+                    .collectionID(info.getCollectionID())
+                    .partitionIDs(info.getPartitionIdsList())
+                    .nodeIDs(info.getNodeIdsList())
+                    .resourceGroupName(info.getResourceGroupName())
+                    .numOutboundNode(info.getNumOutboundNodeMap())
+                    .shardReplicas(shardReplicas)
+                    .build());
+        }
+
+        return DescribeReplicasResp.builder()
+                .replicas(replicas)
+                .build();
+    }
+
     public void waitForCollectionRelease(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub, ReleaseCollectionReq request) {
         boolean isLoaded = true;
         long startTime = System.currentTimeMillis(); // Capture start time/ Timeout in milliseconds (60 seconds)

sdk-core/src/main/java/io/milvus/v2/service/collection/ReplicaInfo.java

@@ -0,0 +1,27 @@
+package io.milvus.v2.service.collection;
+
+import lombok.Builder;
+import lombok.Data;
+import lombok.experimental.SuperBuilder;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+@Data
+@SuperBuilder
+public class ReplicaInfo {
+    private Long replicaID;
+    private Long collectionID;
+    @Builder.Default
+    private List<Long> partitionIDs = new ArrayList<>();
+    @Builder.Default
+    private List<ShardReplica> shardReplicas = new ArrayList<>();
+    @Builder.Default
+    private List<Long> nodeIDs = new ArrayList<>(); // include leaders
+    @Builder.Default
+    private String resourceGroupName = "";
+    @Builder.Default
+    private Map<String, Integer> numOutboundNode = new HashMap<>();
+}

sdk-core/src/main/java/io/milvus/v2/service/collection/ShardReplica.java

@@ -0,0 +1,19 @@
+package io.milvus.v2.service.collection;
+
+import lombok.Builder;
+import lombok.Data;
+import lombok.experimental.SuperBuilder;
+
+import java.util.ArrayList;
+import java.util.List;
+
+@Data
+@SuperBuilder
+public class ShardReplica {
+    private Long leaderID;
+    private String leaderAddress; // IP:port
+    @Builder.Default
+    private String channelName = "";
+    @Builder.Default
+    private List<Long> nodeIDs = new ArrayList<>();
+}

sdk-core/src/main/java/io/milvus/v2/service/collection/request/DescribeReplicasReq.java

@@ -0,0 +1,11 @@
+package io.milvus.v2.service.collection.request;
+
+import lombok.Data;
+import lombok.experimental.SuperBuilder;
+
+@Data
+@SuperBuilder
+public class DescribeReplicasReq {
+    private String collectionName;
+    private String databaseName;
+}

sdk-core/src/main/java/io/milvus/v2/service/collection/response/DescribeCollectionResp.java

@@ -33,6 +33,7 @@ import java.util.Map;
 @SuperBuilder
 public class DescribeCollectionResp {
     private String collectionName;
+    private Long collectionID;
     private String databaseName;
     private String description;
     private Long numOfPartitions;

sdk-core/src/main/java/io/milvus/v2/service/collection/response/DescribeReplicasResp.java

@@ -0,0 +1,16 @@
+package io.milvus.v2.service.collection.response;
+
+import io.milvus.v2.service.collection.ReplicaInfo;
+import lombok.Builder;
+import lombok.Data;
+import lombok.experimental.SuperBuilder;
+
+import java.util.ArrayList;
+import java.util.List;
+
+@Data
+@SuperBuilder
+public class DescribeReplicasResp {
+    @Builder.Default
+    private List<ReplicaInfo> replicas = new ArrayList<>();
+}

sdk-core/src/main/java/io/milvus/v2/utils/ConvertUtils.java

@@ -129,6 +129,7 @@ public class ConvertUtils {
 
         DescribeCollectionResp describeCollectionResp = DescribeCollectionResp.builder()
                 .collectionName(response.getCollectionName())
+                .collectionID(response.getCollectionID())
                 .databaseName(response.getDbName())
                 .description(response.getSchema().getDescription())
                 .numOfPartitions(response.getNumPartitions())

sdk-core/src/test/java/io/milvus/v2/client/MilvusClientV2DockerTest.java

@@ -2190,4 +2190,58 @@ class MilvusClientV2DockerTest {
                 .dimension(DIMENSION)
                 .build());
     }
+
+    @Test
+    void testReplica() {
+        String randomCollectionName = generator.generate(10);
+
+        String vectorFieldName = "float_vector";
+        CreateCollectionReq.CollectionSchema collectionSchema = baseSchema();
+        collectionSchema.addField(AddFieldReq.builder()
+                .fieldName(vectorFieldName)
+                .dataType(DataType.FloatVector)
+                .dimension(DIMENSION)
+                .build());
+
+        IndexParam indexParam = IndexParam.builder()
+                .fieldName(vectorFieldName)
+                .indexType(IndexParam.IndexType.AUTOINDEX)
+                .metricType(IndexParam.MetricType.COSINE)
+                .build();
+
+        CreateCollectionReq requestCreate = CreateCollectionReq.builder()
+                .collectionName(randomCollectionName)
+                .collectionSchema(collectionSchema)
+                .indexParams(Collections.singletonList(indexParam))
+                .build();
+        client.createCollection(requestCreate);
+
+        // insert rows
+        long count = 10000;
+        List<JsonObject> data = generateRandomData(collectionSchema, count);
+        InsertResp insertResp = client.insert(InsertReq.builder()
+                .collectionName(randomCollectionName)
+                .data(data)
+                .build());
+        Assertions.assertEquals(count, insertResp.getInsertCnt());
+
+        DescribeCollectionResp descCollResp = client.describeCollection(DescribeCollectionReq.builder()
+                .collectionName(randomCollectionName)
+                .build());
+
+        DescribeReplicasResp descReplicaResp = client.describeReplicas(DescribeReplicasReq.builder()
+                .collectionName(randomCollectionName)
+                .build());
+        Assertions.assertEquals(1, descReplicaResp.getReplicas().size());
+        io.milvus.v2.service.collection.ReplicaInfo info = descReplicaResp.getReplicas().get(0);
+        Assertions.assertEquals(descCollResp.getCollectionID(), info.getCollectionID());
+        Assertions.assertEquals(1, info.getNodeIDs().size());
+        Assertions.assertNotEquals(0L, info.getReplicaID());
+        Assertions.assertFalse(info.getResourceGroupName().isEmpty());
+        Assertions.assertEquals(1, info.getShardReplicas().size());
+        io.milvus.v2.service.collection.ShardReplica replica = info.getShardReplicas().get(0);
+        Assertions.assertFalse(replica.getChannelName().isEmpty());
+        Assertions.assertFalse(replica.getLeaderAddress().isEmpty());
+        Assertions.assertNotEquals(0L, replica.getLeaderID());
+    }
 }