ResourceGroup interfaces for V2

Signed-off-by: yhmo <[email protected]>

sdk-core/src/main/java/io/milvus/common/resourcegroup/NodeInfo.java

@@ -0,0 +1,12 @@
+package io.milvus.common.resourcegroup;
+
+import lombok.Data;
+import lombok.experimental.SuperBuilder;
+
+@Data
+@SuperBuilder
+public class NodeInfo {
+    private Long nodeId;
+    private String address;
+    private String hostname;
+}

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

@@ -22,10 +22,7 @@
 import io.grpc.ManagedChannel;
 import io.grpc.Status;
 import io.grpc.StatusRuntimeException;
-import io.milvus.grpc.ClientInfo;
-import io.milvus.grpc.ConnectRequest;
-import io.milvus.grpc.ConnectResponse;
-import io.milvus.grpc.MilvusServiceGrpc;
+import io.milvus.grpc.*;
 import io.milvus.orm.iterator.QueryIterator;
 import io.milvus.orm.iterator.SearchIterator;
 
@@ -46,6 +43,9 @@
 import io.milvus.v2.service.rbac.RBACService;
 import io.milvus.v2.service.rbac.request.*;
 import io.milvus.v2.service.rbac.response.*;
+import io.milvus.v2.service.resourcegroup.ResourceGroupService;
+import io.milvus.v2.service.resourcegroup.request.*;
+import io.milvus.v2.service.resourcegroup.response.*;
 import io.milvus.v2.service.utility.UtilityService;
 import io.milvus.v2.service.utility.request.*;
 import io.milvus.v2.service.utility.response.*;
@@ -74,6 +74,7 @@ public class MilvusClientV2 {
     private final VectorService vectorService = new VectorService();
     private final PartitionService partitionService = new PartitionService();
     private final RBACService rbacService = new RBACService();
+    private final ResourceGroupService rgroupService = new ResourceGroupService();
     private final UtilityService utilityService = new UtilityService();
     private ConnectConfig connectConfig;
     private RetryConfig retryConfig = RetryConfig.builder().build();
@@ -267,6 +268,9 @@ private <T> T retry(Callable<T> callable) {
         return null;
     }
 
+    /////////////////////////////////////////////////////////////////////////////////////////////
+    // Database Operations
+    /////////////////////////////////////////////////////////////////////////////////////////////
     /**
      * use Database
      * @param dbName databaseName
@@ -341,7 +345,9 @@ public DescribeDatabaseResp describeDatabase(DescribeDatabaseReq request) {
         return retry(()-> databaseService.describeDatabase(this.getRpcStub(), request));
     }
 
-    //Collection Operations
+    /////////////////////////////////////////////////////////////////////////////////////////////
+    // Collection Operations
+    /////////////////////////////////////////////////////////////////////////////////////////////
     /**
      * Creates a collection in Milvus.
      * @param request create collection request
@@ -481,7 +487,9 @@ public Boolean getLoadState(GetLoadStateReq request) {
         return retry(()->collectionService.getLoadState(this.getRpcStub(), request));
     }
 
-    //Index Operations
+    /////////////////////////////////////////////////////////////////////////////////////////////
+    // Index Operations
+    /////////////////////////////////////////////////////////////////////////////////////////////
     /**
      * Creates an index for a specified field in a collection in Milvus.
      *
@@ -547,7 +555,9 @@ public List<String> listIndexes(ListIndexesReq request) {
         return retry(()->indexService.listIndexes(this.getRpcStub(), request));
     }
 
+    /////////////////////////////////////////////////////////////////////////////////////////////
     // Vector Operations
+    /////////////////////////////////////////////////////////////////////////////////////////////
     /**
      * Inserts vectors into a collection in Milvus.
      *
@@ -634,7 +644,9 @@ public SearchIterator searchIterator(SearchIteratorReq request) {
         return retry(()->vectorService.searchIterator(this.getRpcStub(), request));
     }
 
+    /////////////////////////////////////////////////////////////////////////////////////////////
     // Partition Operations
+    /////////////////////////////////////////////////////////////////////////////////////////////
     /**
      * Creates a partition in a collection in Milvus.
      *
@@ -700,7 +712,9 @@ public void releasePartitions(ReleasePartitionsReq request) {
         retry(()->partitionService.releasePartitions(this.getRpcStub(), request));
     }
 
+    /////////////////////////////////////////////////////////////////////////////////////////////
     // RBAC operations
+    /////////////////////////////////////////////////////////////////////////////////////////////
     /**
      * list users
      *
@@ -837,8 +851,68 @@ public void revokePrivilegeV2(RevokePrivilegeReqV2 request) {
         retry(()->rbacService.revokePrivilegeV2(this.getRpcStub(), request));
     }
 
-    // Utility Operations
+    /////////////////////////////////////////////////////////////////////////////////////////////
+    // Resource group Operations
+    /////////////////////////////////////////////////////////////////////////////////////////////
+    /**
+     * Create a resource group.
+     *
+     * @param request {@link CreateResourceGroupReq}
+     */
+    public void createResourceGroup(CreateResourceGroupReq request){
+        retry(()->rgroupService.createResourceGroup(this.getRpcStub(), request));
+    }
+
+    /**
+     * Update resource groups.
+     *
+     * @param request {@link UpdateResourceGroupsReq}
+     */
+    public void updateResourceGroups(UpdateResourceGroupsReq request) {
+        retry(()->rgroupService.updateResourceGroups(this.getRpcStub(), request));
+    }
+
+    /**
+     * Drop a resource group.
+     *
+     * @param request {@link DropResourceGroupReq}
+     */
+    public void dropResourceGroup(DropResourceGroupReq request) {
+        retry(()->rgroupService.dropResourceGroup(this.getRpcStub(), request));
+    }
 
+    /**
+     * List resource groups.
+     *
+     * @param request {@link ListResourceGroupsReq}
+     * @return ListResourceGroupsResp
+     */
+    ListResourceGroupsResp listResourceGroups(ListResourceGroupsReq request) {
+        return retry(()->rgroupService.listResourceGroups(this.getRpcStub(), request));
+    }
+
+    /**
+     * Describe a resource group.
+     *
+     * @param request {@link DescribeResourceGroupReq}
+     * @return DescribeResourceGroupResp
+     */
+    DescribeResourceGroupResp describeResourceGroup(DescribeResourceGroupReq request) {
+        return retry(()->rgroupService.describeResourceGroup(this.getRpcStub(), request));
+    }
+
+    /**
+     * Transfer a replica from source resource group to target resource_group.
+     *
+     * @param request {@link TransferReplicaReq}
+     */
+    public void transferReplica(TransferReplicaReq request) {
+        retry(()->rgroupService.transferReplica(this.getRpcStub(), request));
+    }
+
+    /////////////////////////////////////////////////////////////////////////////////////////////
+    // Utility Operations
+    /////////////////////////////////////////////////////////////////////////////////////////////
     /**
      * create aliases
      *

sdk-core/src/main/java/io/milvus/v2/service/resourcegroup/ResourceGroupService.java

@@ -0,0 +1,182 @@
+package io.milvus.v2.service.resourcegroup;
+
+import io.milvus.grpc.*;
+import io.milvus.v2.exception.ErrorCode;
+import io.milvus.v2.exception.MilvusClientException;
+import io.milvus.v2.service.BaseService;
+import io.milvus.v2.service.resourcegroup.request.*;
+import io.milvus.v2.service.resourcegroup.response.*;
+import org.apache.commons.lang3.StringUtils;
+
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+public class ResourceGroupService extends BaseService {
+    private static ResourceGroupConfig convertResourceGroupConfig(io.milvus.common.resourcegroup.ResourceGroupConfig config) {
+        if (config == null) {
+            throw new MilvusClientException(ErrorCode.INVALID_PARAMS, "Invalid resource group config");
+        }
+
+        ResourceGroupConfig.Builder builder = ResourceGroupConfig.newBuilder();
+        builder.setRequests(ResourceGroupLimit.newBuilder()
+                .setNodeNum(config.getRequests().getNodeNum()))
+                .build();
+        builder.setLimits(ResourceGroupLimit.newBuilder()
+                .setNodeNum(config.getLimits().getNodeNum()))
+                .build();
+
+        for (io.milvus.common.resourcegroup.ResourceGroupTransfer groupFrom : config.getFrom()) {
+            builder.addTransferFrom(ResourceGroupTransfer.newBuilder()
+                    .setResourceGroup(groupFrom.getResourceGroupName()))
+                    .build();
+        }
+
+        for (io.milvus.common.resourcegroup.ResourceGroupTransfer groupTo : config.getTo()) {
+            builder.addTransferTo(ResourceGroupTransfer.newBuilder()
+                    .setResourceGroup(groupTo.getResourceGroupName()))
+                    .build();
+        }
+
+        return builder.build();
+    }
+
+    private static io.milvus.common.resourcegroup.ResourceGroupConfig convertResourceGroupConfig(ResourceGroupConfig config) {
+        List<io.milvus.common.resourcegroup.ResourceGroupTransfer> fromList = new ArrayList<>();
+        config.getTransferFromList().forEach((groupFrom)->{
+            fromList.add(new io.milvus.common.resourcegroup.ResourceGroupTransfer(groupFrom.getResourceGroup()));
+        });
+
+        List<io.milvus.common.resourcegroup.ResourceGroupTransfer> toList = new ArrayList<>();
+        config.getTransferToList().forEach((groupTo)->{
+            toList.add(new io.milvus.common.resourcegroup.ResourceGroupTransfer(groupTo.getResourceGroup()));
+        });
+
+        return io.milvus.common.resourcegroup.ResourceGroupConfig.newBuilder()
+                .withRequests(new io.milvus.common.resourcegroup.ResourceGroupLimit(config.getRequests().getNodeNum()))
+                .withLimits(new io.milvus.common.resourcegroup.ResourceGroupLimit(config.getLimits().getNodeNum()))
+                .withFrom(fromList)
+                .withTo(toList)
+                .build();
+    }
+
+    public Void createResourceGroup(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
+                                    CreateResourceGroupReq request) {
+        String title = String.format("CreateResourceGroupReq groupName:%s", request.getGroupName());
+
+        ResourceGroupConfig rpcConfig = convertResourceGroupConfig(request.getConfig());
+        CreateResourceGroupRequest rpcRequest = CreateResourceGroupRequest.newBuilder()
+                .setResourceGroup(request.getGroupName())
+                .setConfig(rpcConfig)
+                .build();
+
+        Status status = blockingStub.createResourceGroup(rpcRequest);
+        rpcUtils.handleResponse(title, status);
+        return null;
+    }
+
+    public Void updateResourceGroups(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
+                                     UpdateResourceGroupsReq request) {
+        String title = "UpdateResourceGroupsReq";
+
+        Map<String, io.milvus.common.resourcegroup.ResourceGroupConfig> resourceGroups = request.getResourceGroups();
+        if (resourceGroups.isEmpty()) {
+            throw new MilvusClientException(ErrorCode.INVALID_PARAMS, "Resource group configurations cannot be empty");
+        }
+
+        UpdateResourceGroupsRequest.Builder requestBuilder = UpdateResourceGroupsRequest.newBuilder();
+        resourceGroups.forEach((groupName, config) -> {
+            ResourceGroupConfig rpcConfig = convertResourceGroupConfig(config);
+            requestBuilder.putResourceGroups(groupName, rpcConfig);
+        });
+
+        Status status = blockingStub.updateResourceGroups(requestBuilder.build());
+        rpcUtils.handleResponse(title, status);
+        return null;
+    }
+
+    public Void dropResourceGroup(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
+                                  DropResourceGroupReq request) {
+        String title = String.format("DropResourceGroupReq groupName:%s", request.getGroupName());
+
+        DropResourceGroupRequest rpcRequest = DropResourceGroupRequest.newBuilder()
+                .setResourceGroup(request.getGroupName())
+                .build();
+
+        Status status = blockingStub.dropResourceGroup(rpcRequest);
+        rpcUtils.handleResponse(title, status);
+        return null;
+    }
+
+    public ListResourceGroupsResp listResourceGroups(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
+                                                     ListResourceGroupsReq request) {
+        String title = "ListResourceGroupsReq";
+        ListResourceGroupsResponse response = blockingStub.listResourceGroups(ListResourceGroupsRequest.newBuilder().build());
+        rpcUtils.handleResponse(title, response.getStatus());
+        return ListResourceGroupsResp.builder()
+                .groupNames(response.getResourceGroupsList())
+                .build();
+    }
+
+    public DescribeResourceGroupResp describeResourceGroup(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
+                                                           DescribeResourceGroupReq request) {
+        String title = String.format("DescribeResourceGroupReq groupName:%s", request.getGroupName());
+
+        DescribeResourceGroupRequest rpcRequest = DescribeResourceGroupRequest.newBuilder()
+                .setResourceGroup(request.getGroupName())
+                .build();
+
+        DescribeResourceGroupResponse response = blockingStub.describeResourceGroup(rpcRequest);
+        rpcUtils.handleResponse(title, response.getStatus());
+
+        ResourceGroup rgroup = response.getResourceGroup();
+        List<io.milvus.common.resourcegroup.NodeInfo> nodes = new ArrayList<>();
+        rgroup.getNodesList().forEach((node)->{
+            nodes.add(io.milvus.common.resourcegroup.NodeInfo.builder()
+                    .nodeId(node.getNodeId())
+                    .address(node.getAddress())
+                    .hostname(node.getHostname())
+                    .build());
+        });
+        return DescribeResourceGroupResp.builder()
+                .groupName(rgroup.getName())
+                .capacity(rgroup.getCapacity())
+                .numberOfAvailableNode(rgroup.getNumAvailableNode())
+                .numberOfLoadedReplica(rgroup.getNumLoadedReplicaMap())
+                .numberOfOutgoingNode(rgroup.getNumOutgoingNodeMap())
+                .numberOfIncomingNode(rgroup.getNumIncomingNodeMap())
+                .config(convertResourceGroupConfig(rgroup.getConfig()))
+                .nodes(nodes)
+                .build();
+    }
+
+    public Void transferReplica(MilvusServiceGrpc.MilvusServiceBlockingStub blockingStub,
+                                TransferReplicaReq request) {
+        if (StringUtils.isEmpty(request.getSourceGroupName())) {
+            throw new MilvusClientException(ErrorCode.INVALID_PARAMS, "Invalid source group name");
+        }
+        if (StringUtils.isEmpty(request.getTargetGroupName())) {
+            throw new MilvusClientException(ErrorCode.INVALID_PARAMS, "Invalid target group name");
+        }
+        if (StringUtils.isEmpty(request.getCollectionName())) {
+            throw new MilvusClientException(ErrorCode.INVALID_PARAMS, "Invalid collection name");
+        }
+
+        String title = String.format("TransferReplicaReq sourceGroupName:%s targetGroupName:%s collectionName:%s",
+                request.getSourceGroupName(), request.getTargetGroupName(), request.getCollectionName());
+
+        TransferReplicaRequest.Builder requestBuilder = TransferReplicaRequest.newBuilder()
+                .setSourceResourceGroup(request.getSourceGroupName())
+                .setTargetResourceGroup(request.getTargetGroupName())
+                .setCollectionName(request.getCollectionName())
+                .setNumReplica(request.getNumberOfReplicas());
+
+        if (StringUtils.isNotEmpty(request.getDatabaseName())) {
+            requestBuilder.setDbName(request.getDatabaseName());
+        }
+
+        Status status = blockingStub.transferReplica(requestBuilder.build());
+        rpcUtils.handleResponse(title, status);
+        return null;
+    }
+}

sdk-core/src/main/java/io/milvus/v2/service/resourcegroup/request/CreateResourceGroupReq.java

@@ -0,0 +1,12 @@
+package io.milvus.v2.service.resourcegroup.request;
+
+import io.milvus.common.resourcegroup.ResourceGroupConfig;
+import lombok.Data;
+import lombok.experimental.SuperBuilder;
+
+@Data
+@SuperBuilder
+public class CreateResourceGroupReq {
+    private String groupName;
+    private ResourceGroupConfig config;
+}

sdk-core/src/main/java/io/milvus/v2/service/resourcegroup/request/DescribeResourceGroupReq.java

@@ -0,0 +1,10 @@
+package io.milvus.v2.service.resourcegroup.request;
+
+import lombok.Data;
+import lombok.experimental.SuperBuilder;
+
+@Data
+@SuperBuilder
+public class DescribeResourceGroupReq {
+    private String groupName;
+}

sdk-core/src/main/java/io/milvus/v2/service/resourcegroup/request/DropResourceGroupReq.java

@@ -0,0 +1,10 @@
+package io.milvus.v2.service.resourcegroup.request;
+
+import lombok.Data;
+import lombok.experimental.SuperBuilder;
+
+@Data
+@SuperBuilder
+public class DropResourceGroupReq {
+    private String groupName;
+}