ResourceGroup interfaces for V2

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

src/main/java/io/milvus/common/resourcegroup/NodeInfo.java

@@ -0,0 +1,31 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 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;
+}

src/main/java/io/milvus/common/resourcegroup/ResourceGroupConfig.java

@@ -1,3 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 io.milvus.common.resourcegroup;
 
 import java.util.stream.Collectors;
@@ -12,10 +31,12 @@ public class ResourceGroupConfig {
     private final ResourceGroupLimit limits;
     private final List<ResourceGroupTransfer> from;
     private final List<ResourceGroupTransfer> to;
+    private final ResourceGroupNodeFilter nodeFilter;
 
     private ResourceGroupConfig(Builder builder) {
         this.requests = builder.requests;
         this.limits = builder.limits;
+        this.nodeFilter = builder.nodeFilter;
 
         if (null == builder.from) {
             this.from = new ArrayList<>();
@@ -39,6 +60,7 @@ public static final class Builder {
         private ResourceGroupLimit limits;
         private List<ResourceGroupTransfer> from;
         private List<ResourceGroupTransfer> to;
+        private ResourceGroupNodeFilter nodeFilter;
 
         private Builder() {
         }
@@ -87,6 +109,17 @@ public Builder withTo(@NonNull List<ResourceGroupTransfer> to) {
             return this;
         }
 
+        /**
+         * Set the node filter.
+         * @param nodeFilter if node filter set, resource group will prefer to accept node which match node filter.
+         * @return <code>Builder</code>
+         */
+
+        public Builder withNodeFilter(@NonNull ResourceGroupNodeFilter nodeFilter) {
+            this.nodeFilter = nodeFilter;
+            return this;
+        }
+
         public ResourceGroupConfig build() {
             return new ResourceGroupConfig(this);
         }
@@ -101,12 +134,14 @@ public ResourceGroupConfig(@NonNull io.milvus.grpc.ResourceGroupConfig grpcConfi
         this.to = grpcConfig.getTransferToList().stream()
                 .map(transfer -> new ResourceGroupTransfer(transfer))
                 .collect(Collectors.toList());
+        this.nodeFilter = new ResourceGroupNodeFilter(grpcConfig.getNodeFilter());
     }
 
     public @NonNull io.milvus.grpc.ResourceGroupConfig toGRPC() {
         io.milvus.grpc.ResourceGroupConfig.Builder builder = io.milvus.grpc.ResourceGroupConfig.newBuilder()
                 .setRequests(io.milvus.grpc.ResourceGroupLimit.newBuilder().setNodeNum(requests.getNodeNum()))
-                .setLimits(io.milvus.grpc.ResourceGroupLimit.newBuilder().setNodeNum(limits.getNodeNum()));
+                .setLimits(io.milvus.grpc.ResourceGroupLimit.newBuilder().setNodeNum(limits.getNodeNum()))
+                .setNodeFilter(nodeFilter.toGRPC());
         for (ResourceGroupTransfer transfer : from) {
             builder.addTransferFrom(transfer.toGRPC());
         }

src/main/java/io/milvus/common/resourcegroup/ResourceGroupLimit.java

@@ -1,3 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 io.milvus.common.resourcegroup;
 
 import org.jetbrains.annotations.NotNull;

src/main/java/io/milvus/common/resourcegroup/ResourceGroupNodeFilter.java

@@ -0,0 +1,83 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 io.milvus.common.resourcegroup;
+
+import java.util.List;
+import java.util.Map;
+import java.util.stream.Collectors;
+
+import io.milvus.grpc.KeyValuePair;
+import io.milvus.param.ParamUtils;
+import lombok.Getter;
+import lombok.NonNull;
+
+@Getter
+public class ResourceGroupNodeFilter {
+    private final Map<String, String> nodeLabels;
+
+    private ResourceGroupNodeFilter(Builder builder) {
+        this.nodeLabels = builder.nodeLabels;
+    }
+
+    public static Builder newBuilder() {
+        return new Builder();
+    }
+
+    public static final class Builder {
+        private Map<String, String> nodeLabels;
+        private Builder() {
+        }
+
+        /**
+         * Set the node label filter
+         * @param key label name
+         * @param value label value
+         * @return <code>Builder</code>
+         */
+        public Builder withNodeLabel(@NonNull String key, @NonNull String value) {
+            this.nodeLabels.put(key, value);
+            return this;
+        }
+
+        public ResourceGroupNodeFilter build() {
+            return new ResourceGroupNodeFilter(this);
+        }
+    }
+
+    /**
+     * Transfer to grpc
+     * @return io.milvus.grpc.ResourceGroupNodeFilter
+     */
+    public @NonNull io.milvus.grpc.ResourceGroupNodeFilter toGRPC() {
+        List<KeyValuePair> pair = ParamUtils.AssembleKvPair(nodeLabels);
+        return io.milvus.grpc.ResourceGroupNodeFilter.newBuilder()
+                .addAllNodeLabels(pair)
+               .build();
+    }
+
+    /**
+     * Constructor from grpc
+     * @param filter grpc filter object
+     */
+    public ResourceGroupNodeFilter(io.milvus.grpc.ResourceGroupNodeFilter filter) {
+        this.nodeLabels = filter.getNodeLabelsList().stream().collect(Collectors.toMap(KeyValuePair::getKey, KeyValuePair::getValue));
+    }
+
+}

src/main/java/io/milvus/common/resourcegroup/ResourceGroupTransfer.java

@@ -1,3 +1,22 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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 io.milvus.common.resourcegroup;
 
 import lombok.Getter;

src/main/java/io/milvus/v2/client/MilvusClientV2.java

@@ -46,6 +46,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 +77,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 +271,9 @@ private <T> T retry(Callable<T> callable) {
         return null;
     }
 
+    /////////////////////////////////////////////////////////////////////////////////////////////
+    // Database Operations
+    /////////////////////////////////////////////////////////////////////////////////////////////
     /**
      * use Database
      * @param dbName databaseName
@@ -341,7 +348,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
@@ -480,7 +489,9 @@ public Boolean getLoadState(GetLoadStateReq request) {
         return retry(()->collectionService.getLoadState(this.getRpcStub(), request));
     }
 
+    /////////////////////////////////////////////////////////////////////////////////////////////
     //Index Operations
+    /////////////////////////////////////////////////////////////////////////////////////////////
     /**
      * Creates an index for a specified field in a collection in Milvus.
      *
@@ -546,7 +557,9 @@ public List<String> listIndexes(ListIndexesReq request) {
         return retry(()->indexService.listIndexes(this.getRpcStub(), request));
     }
 
+    /////////////////////////////////////////////////////////////////////////////////////////////
     // Vector Operations
+    /////////////////////////////////////////////////////////////////////////////////////////////
     /**
      * Inserts vectors into a collection in Milvus.
      *
@@ -633,7 +646,9 @@ public SearchIterator searchIterator(SearchIteratorReq request) {
         return retry(()->vectorService.searchIterator(this.getRpcStub(), request));
     }
 
+    /////////////////////////////////////////////////////////////////////////////////////////////
     // Partition Operations
+    /////////////////////////////////////////////////////////////////////////////////////////////
     /**
      * Creates a partition in a collection in Milvus.
      *
@@ -699,7 +714,9 @@ public void releasePartitions(ReleasePartitionsReq request) {
         retry(()->partitionService.releasePartitions(this.getRpcStub(), request));
     }
 
-    // RBAC operations
+    /////////////////////////////////////////////////////////////////////////////////////////////
+    // RBAC Operations
+    /////////////////////////////////////////////////////////////////////////////////////////////
     /**
      * list users
      *
@@ -836,8 +853,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
      *