enhance: [GoSDK] Add APIs for backup tool (#38791)

Related to #31291

This PR:

- Add Backup&Restore RBAC APIs
- Add GetPersistentSegmentInfo API
- Add GetServerVersion API
- Add flush task stats

---------

Signed-off-by: Congqi Xia <[email protected]>

client/entity/rbac.go

@@ -33,3 +33,31 @@ type GrantItem struct {
 	Grantor    string
 	Privilege  string
 }
+
+type UserInfo struct {
+	UserDescription
+	Password string
+}
+
+// UserDescription is the model for RBAC user description object.
+type UserDescription struct {
+	Name  string
+	Roles []string
+}
+
+type RBACMeta struct {
+	Users           []*UserInfo
+	Roles           []*Role
+	RoleGrants      []*RoleGrants
+	PrivilegeGroups []*PrivilegeGroup
+}
+
+// RoleGrants is the model for RBAC role description object.
+type RoleGrants struct {
+	Object        string
+	ObjectName    string
+	RoleName      string
+	GrantorName   string
+	PrivilegeName string
+	DbName        string
+}

client/entity/segment.go

@@ -0,0 +1,34 @@
+// 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
+
+import "github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
+
+// Segment represent segment in milvus
+type Segment struct {
+	ID           int64
+	CollectionID int64
+	ParititionID int64
+
+	NumRows int64
+	State   commonpb.SegmentState
+}
+
+// Flushed indicates segment is flushed
+func (s Segment) Flushed() bool {
+	return s.State == commonpb.SegmentState_Flushed
+}

client/milvusclient/admin.go

@@ -0,0 +1,230 @@
+// 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 milvusclient
+
+import (
+	"context"
+
+	"github.com/samber/lo"
+	"google.golang.org/grpc"
+
+	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
+	"github.com/milvus-io/milvus/client/v2/entity"
+	"github.com/milvus-io/milvus/pkg/util/merr"
+)
+
+// GetServerVersionOption is the interface for GetServerVersion request.
+type GetServerVersionOption interface {
+	Request() *milvuspb.GetVersionRequest
+}
+
+type getServerVersionOption struct{}
+
+func (opt *getServerVersionOption) Request() *milvuspb.GetVersionRequest {
+	return &milvuspb.GetVersionRequest{}
+}
+
+func NewGetServerVersionOption() *getServerVersionOption {
+	return &getServerVersionOption{}
+}
+
+// GetServerVersion returns connect Milvus instance version.
+func (c *Client) GetServerVersion(ctx context.Context, option GetServerVersionOption, callOptions ...grpc.CallOption) (string, error) {
+	req := option.Request()
+
+	var version string
+
+	err := c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
+		resp, err := milvusService.GetVersion(ctx, req, callOptions...)
+		version = resp.GetVersion()
+		return merr.CheckRPCCall(resp, err)
+	})
+	return version, err
+}
+
+type GetPersistentSegmentInfoOption interface {
+	Request() *milvuspb.GetPersistentSegmentInfoRequest
+}
+
+type getPersistentSegmentInfoOption struct {
+	collectionName string
+}
+
+func (opt *getPersistentSegmentInfoOption) Request() *milvuspb.GetPersistentSegmentInfoRequest {
+	return &milvuspb.GetPersistentSegmentInfoRequest{
+		CollectionName: opt.collectionName,
+	}
+}
+
+func NewGetPersistentSegmentInfoOption(collectionName string) GetPersistentSegmentInfoOption {
+	return &getPersistentSegmentInfoOption{
+		collectionName: collectionName,
+	}
+}
+
+func (c *Client) GetPersistentSegmentInfo(ctx context.Context, option GetPersistentSegmentInfoOption) ([]*entity.Segment, error) {
+	req := option.Request()
+
+	var segments []*entity.Segment
+
+	err := c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
+		resp, err := milvusService.GetPersistentSegmentInfo(ctx, req)
+		if err = merr.CheckRPCCall(resp, err); err != nil {
+			return err
+		}
+
+		segments = lo.Map(resp.GetInfos(), func(info *milvuspb.PersistentSegmentInfo, _ int) *entity.Segment {
+			return &entity.Segment{
+				ID:           info.GetSegmentID(),
+				CollectionID: info.GetCollectionID(),
+				ParititionID: info.GetPartitionID(),
+				NumRows:      info.GetNumRows(),
+				State:        info.GetState(),
+			}
+		})
+		return nil
+	})
+
+	return segments, err
+}
+
+type BackupRBACOption interface {
+	Request() *milvuspb.BackupRBACMetaRequest
+}
+
+type backupRBACOption struct{}
+
+func (opt *backupRBACOption) Request() *milvuspb.BackupRBACMetaRequest {
+	return &milvuspb.BackupRBACMetaRequest{}
+}
+
+func NewBackupRBACOption() BackupRBACOption {
+	return &backupRBACOption{}
+}
+
+func (c *Client) BackupRBAC(ctx context.Context, option BackupRBACOption, callOptions ...grpc.CallOption) (*entity.RBACMeta, error) {
+	req := option.Request()
+
+	var meta *entity.RBACMeta
+
+	err := c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
+		resp, err := milvusService.BackupRBAC(ctx, req, callOptions...)
+		if err = merr.CheckRPCCall(resp, err); err != nil {
+			return err
+		}
+
+		rbacMeta := resp.GetRBACMeta()
+
+		meta = &entity.RBACMeta{
+			Users: lo.Map(rbacMeta.GetUsers(), func(user *milvuspb.UserInfo, _ int) *entity.UserInfo {
+				return &entity.UserInfo{
+					UserDescription: entity.UserDescription{
+						Name:  user.GetUser(),
+						Roles: lo.Map(user.GetRoles(), func(role *milvuspb.RoleEntity, _ int) string { return role.GetName() }),
+					},
+					Password: user.GetPassword(),
+				}
+			}),
+			Roles: lo.Map(rbacMeta.GetRoles(), func(role *milvuspb.RoleEntity, _ int) *entity.Role {
+				return &entity.Role{
+					RoleName: role.GetName(),
+				}
+			}),
+			RoleGrants: lo.Map(rbacMeta.GetGrants(), func(grant *milvuspb.GrantEntity, _ int) *entity.RoleGrants {
+				return &entity.RoleGrants{
+					Object:        grant.GetObject().GetName(),
+					ObjectName:    grant.GetObjectName(),
+					RoleName:      grant.GetRole().GetName(),
+					GrantorName:   grant.GetGrantor().GetUser().GetName(),
+					PrivilegeName: grant.GetGrantor().GetPrivilege().GetName(),
+					DbName:        grant.GetDbName(),
+				}
+			}),
+			PrivilegeGroups: lo.Map(rbacMeta.GetPrivilegeGroups(), func(group *milvuspb.PrivilegeGroupInfo, _ int) *entity.PrivilegeGroup {
+				return &entity.PrivilegeGroup{
+					GroupName:  group.GetGroupName(),
+					Privileges: lo.Map(group.GetPrivileges(), func(privilege *milvuspb.PrivilegeEntity, _ int) string { return privilege.GetName() }),
+				}
+			}),
+		}
+
+		return nil
+	})
+	return meta, err
+}
+
+type RestoreRBACOption interface {
+	Request() *milvuspb.RestoreRBACMetaRequest
+}
+
+type restoreRBACOption struct {
+	meta *entity.RBACMeta
+}
+
+func (opt *restoreRBACOption) Request() *milvuspb.RestoreRBACMetaRequest {
+	return &milvuspb.RestoreRBACMetaRequest{
+		RBACMeta: &milvuspb.RBACMeta{
+			Users: lo.Map(opt.meta.Users, func(user *entity.UserInfo, _ int) *milvuspb.UserInfo {
+				return &milvuspb.UserInfo{
+					User:     user.Name,
+					Roles:    lo.Map(user.Roles, func(role string, _ int) *milvuspb.RoleEntity { return &milvuspb.RoleEntity{Name: role} }),
+					Password: user.Password,
+				}
+			}),
+			Roles: lo.Map(opt.meta.Roles, func(role *entity.Role, _ int) *milvuspb.RoleEntity {
+				return &milvuspb.RoleEntity{Name: role.RoleName}
+			}),
+			Grants: lo.Map(opt.meta.RoleGrants, func(grant *entity.RoleGrants, _ int) *milvuspb.GrantEntity {
+				return &milvuspb.GrantEntity{
+					Object:     &milvuspb.ObjectEntity{Name: grant.Object},
+					ObjectName: grant.ObjectName,
+					Role:       &milvuspb.RoleEntity{Name: grant.RoleName},
+					Grantor: &milvuspb.GrantorEntity{
+						User: &milvuspb.UserEntity{
+							Name: grant.GrantorName,
+						},
+						Privilege: &milvuspb.PrivilegeEntity{
+							Name: grant.PrivilegeName,
+						},
+					},
+					DbName: grant.DbName,
+				}
+			}),
+			PrivilegeGroups: lo.Map(opt.meta.PrivilegeGroups, func(group *entity.PrivilegeGroup, _ int) *milvuspb.PrivilegeGroupInfo {
+				return &milvuspb.PrivilegeGroupInfo{
+					GroupName: group.GroupName,
+					Privileges: lo.Map(group.Privileges, func(privilege string, _ int) *milvuspb.PrivilegeEntity {
+						return &milvuspb.PrivilegeEntity{Name: privilege}
+					}),
+				}
+			}),
+		},
+	}
+}
+
+func NewRestoreRBACOption(meta *entity.RBACMeta) RestoreRBACOption {
+	return &restoreRBACOption{meta: meta}
+}
+
+func (c *Client) RestoreRBAC(ctx context.Context, option RestoreRBACOption, callOptions ...grpc.CallOption) error {
+	req := option.Request()
+
+	return c.callService(func(milvusService milvuspb.MilvusServiceClient) error {
+		resp, err := milvusService.RestoreRBAC(ctx, req, callOptions...)
+		return merr.CheckRPCCall(resp, err)
+	})
+}