enhance: [GoSDK] Support Resource Group Controlling APIs

Related to milvus-io#30647

This PR add APIs:

- DescribeResourceGroup
- UpdateResourceGroup
- TransferReplica

Signed-off-by: Congqi Xia <congqi.xia@zilliz.com>

client/entity/resource_group.go

@@ -0,0 +1,54 @@
+// Licensed to the LF AI & Data foundation 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 entity
+
+type ResourceGroup struct {
+	Name             string
+	Capacity         int32
+	NumAvailableNode int32
+	NumLoadedReplica map[string]int32
+	NumOutgoingNode  map[string]int32
+	NumIncomingNode  map[string]int32
+	Config           *ResourceGroupConfig
+	Nodes            []NodeInfo
+}
+
+type NodeInfo struct {
+	NodeID   int64
+	Address  string
+	HostName string
+}
+
+type ResourceGroupLimit struct {
+	NodeNum int32
+}
+
+type ResourceGroupTransfer struct {
+	ResourceGroup string
+}
+
+type ResourceGroupNodeFilter struct {
+	NodeLabels map[string]string
+}
+
+type ResourceGroupConfig struct {
+	Requests     ResourceGroupLimit
+	Limits       ResourceGroupLimit
+	TransferFrom []*ResourceGroupTransfer
+	TransferTo   []*ResourceGroupTransfer
+	NodeFilter   ResourceGroupNodeFilter
+}

client/milvusclient/resource_group.go

@@ -19,9 +19,13 @@ package milvusclient
 import (
 	"context"
 
+	"github.com/samber/lo"
 	"google.golang.org/grpc"
 
+	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
+	"github.com/milvus-io/milvus-proto/go-api/v2/rgpb"
+	"github.com/milvus-io/milvus/client/v2/entity"
 	"github.com/milvus-io/milvus/pkg/util/merr"
 )
 
@@ -63,3 +67,78 @@ func (c *Client) DropResourceGroup(ctx context.Context, opt DropResourceGroupOpt
 
 	return err
 }
+
+func (c *Client) DescribeResourceGroup(ctx context.Context, opt DescribeResourceGroupOption, callOptions ...grpc.CallOption) (*entity.ResourceGroup, error) {
+	req := opt.Request()
+
+	var rg *entity.ResourceGroup
+	err := c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
+		resp, err := milvusService.DescribeResourceGroup(ctx, req, callOptions...)
+		if err = merr.CheckRPCCall(resp, err); err != nil {
+			return err
+		}
+
+		resultRg := resp.GetResourceGroup()
+		rg = &entity.ResourceGroup{
+			Name:             resultRg.GetName(),
+			Capacity:         resultRg.GetCapacity(),
+			NumAvailableNode: resultRg.GetNumAvailableNode(),
+			NumLoadedReplica: resultRg.GetNumLoadedReplica(),
+			NumOutgoingNode:  resultRg.GetNumOutgoingNode(),
+			NumIncomingNode:  resultRg.GetNumIncomingNode(),
+			Config: &entity.ResourceGroupConfig{
+				Requests: entity.ResourceGroupLimit{
+					NodeNum: resultRg.GetConfig().GetRequests().GetNodeNum(),
+				},
+				Limits: entity.ResourceGroupLimit{
+					NodeNum: resultRg.GetConfig().GetLimits().GetNodeNum(),
+				},
+				TransferFrom: lo.Map(resultRg.GetConfig().GetTransferFrom(), func(transfer *rgpb.ResourceGroupTransfer, i int) *entity.ResourceGroupTransfer {
+					return &entity.ResourceGroupTransfer{
+						ResourceGroup: transfer.GetResourceGroup(),
+					}
+				}),
+				TransferTo: lo.Map(resultRg.GetConfig().GetTransferTo(), func(transfer *rgpb.ResourceGroupTransfer, i int) *entity.ResourceGroupTransfer {
+					return &entity.ResourceGroupTransfer{
+						ResourceGroup: transfer.GetResourceGroup(),
+					}
+				}),
+				NodeFilter: entity.ResourceGroupNodeFilter{
+					NodeLabels: entity.KvPairsMap(resultRg.GetConfig().GetNodeFilter().GetNodeLabels()),
+				},
+			},
+			Nodes: lo.Map(resultRg.GetNodes(), func(node *commonpb.NodeInfo, i int) entity.NodeInfo {
+				return entity.NodeInfo{
+					NodeID:   node.GetNodeId(),
+					Address:  node.GetAddress(),
+					HostName: node.GetHostname(),
+				}
+			}),
+		}
+
+		return nil
+	})
+	return rg, err
+}
+
+func (c *Client) UpdateResourceGroup(ctx context.Context, opt UpdateResourceGroupOption, callOptions ...grpc.CallOption) error {
+	req := opt.Request()
+
+	err := c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
+		resp, err := milvusService.UpdateResourceGroups(ctx, req, callOptions...)
+		return merr.CheckRPCCall(resp, err)
+	})
+
+	return err
+}
+
+func (c *Client) TransferReplica(ctx context.Context, opt TransferReplicaOption, callOptions ...grpc.CallOption) error {
+	req := opt.Request()
+
+	err := c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
+		resp, err := milvusService.TransferReplica(ctx, req, callOptions...)
+		return merr.CheckRPCCall(resp, err)
+	})
+
+	return err
+}

client/milvusclient/resource_group_option.go

@@ -17,8 +17,11 @@
 package milvusclient
 
 import (
+	"github.com/samber/lo"
+
 	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/rgpb"
+	"github.com/milvus-io/milvus/client/v2/entity"
 )
 
 type ListResourceGroupsOption interface {
@@ -90,3 +93,101 @@ func (opt *dropResourceGroupOption) Request() *milvuspb.DropResourceGroupRequest
 func NewDropResourceGroupOption(name string) *dropResourceGroupOption {
 	return &dropResourceGroupOption{name: name}
 }
+
+type DescribeResourceGroupOption interface {
+	Request() *milvuspb.DescribeResourceGroupRequest
+}
+
+type describeResourceGroupOption struct {
+	name string
+}
+
+func (opt *describeResourceGroupOption) Request() *milvuspb.DescribeResourceGroupRequest {
+	return &milvuspb.DescribeResourceGroupRequest{
+		ResourceGroup: opt.name,
+	}
+}
+
+func NewDescribeResourceGroupOption(name string) *describeResourceGroupOption {
+	return &describeResourceGroupOption{name: name}
+}
+
+type UpdateResourceGroupOption interface {
+	Request() *milvuspb.UpdateResourceGroupsRequest
+}
+
+type updateResourceGroupOption struct {
+	name     string
+	rgConfig *entity.ResourceGroupConfig
+}
+
+func (opt *updateResourceGroupOption) Request() *milvuspb.UpdateResourceGroupsRequest {
+	return &milvuspb.UpdateResourceGroupsRequest{
+		ResourceGroups: map[string]*rgpb.ResourceGroupConfig{
+			opt.name: {
+				Requests: &rgpb.ResourceGroupLimit{
+					NodeNum: opt.rgConfig.Requests.NodeNum,
+				},
+				Limits: &rgpb.ResourceGroupLimit{
+					NodeNum: opt.rgConfig.Limits.NodeNum,
+				},
+				TransferFrom: lo.Map(opt.rgConfig.TransferFrom, func(transfer *entity.ResourceGroupTransfer, i int) *rgpb.ResourceGroupTransfer {
+					return &rgpb.ResourceGroupTransfer{
+						ResourceGroup: transfer.ResourceGroup,
+					}
+				}),
+				TransferTo: lo.Map(opt.rgConfig.TransferTo, func(transfer *entity.ResourceGroupTransfer, i int) *rgpb.ResourceGroupTransfer {
+					return &rgpb.ResourceGroupTransfer{
+						ResourceGroup: transfer.ResourceGroup,
+					}
+				}),
+				NodeFilter: &rgpb.ResourceGroupNodeFilter{
+					NodeLabels: entity.MapKvPairs(opt.rgConfig.NodeFilter.NodeLabels),
+				},
+			},
+		},
+	}
+}
+
+func NewUpdateResourceGroupOption(name string, resourceGroupConfig *entity.ResourceGroupConfig) *updateResourceGroupOption {
+	return &updateResourceGroupOption{
+		name:     name,
+		rgConfig: resourceGroupConfig,
+	}
+}
+
+type TransferReplicaOption interface {
+	Request() *milvuspb.TransferReplicaRequest
+}
+
+type transferReplicaOption struct {
+	collectionName string
+	sourceRG       string
+	targetRG       string
+	replicaNum     int64
+	dbName         string
+}
+
+func (opt *transferReplicaOption) WithDBName(dbName string) *transferReplicaOption {
+	opt.dbName = dbName
+	return opt
+}
+
+func (opt *transferReplicaOption) Request() *milvuspb.TransferReplicaRequest {
+	return &milvuspb.TransferReplicaRequest{
+		CollectionName:      opt.collectionName,
+		SourceResourceGroup: opt.sourceRG,
+		TargetResourceGroup: opt.targetRG,
+		NumReplica:          opt.replicaNum,
+		DbName:              opt.dbName,
+	}
+}
+
+func NewTransferReplicaOption(collectionName, sourceGroup, targetGroup string, replicaNum int64) *transferReplicaOption {
+	return &transferReplicaOption{
+		collectionName: collectionName,
+		sourceRG:       sourceGroup,
+		targetRG:       targetGroup,
+		replicaNum:     replicaNum,
+	}
+}