Repository: milvus-io/milvus-sdk-go
Pull request #808 (2024-08-20 07:08:04)

enhance: Add BackupRBAC/RestoreRBAC API to enable rbac backup
issue: milvus-io/milvus#35443

sre-ci-robot: [APPROVALNOTIFIER] This PR is **NOT APPROVED**

This pull-request has been approved by: *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/808#" title="Author self-approved">weiliu1031</a>*
To complete the [pull request process](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process), please assign **congqixia** after the PR has been reviewed.
You can assign the PR to them by writing `/assign @congqixia` in a comment when ready.

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=milvus-io%2Fmilvus-sdk-go).

<details open>
Needs approval from an approver in each of these files:

- **[OWNERS](https://github.com/milvus-io/milvus-sdk-go/blob/master/OWNERS)**

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":["congqixia"]} -->
sre-ci-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/808#pullrequestreview-2249565384" title="LGTM">congqixia</a>*, *<a href="https://github.com/milvus-io/milvus-sdk-go/pull/808#" title="Author self-approved">weiliu1031</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=milvus-io%2Fmilvus-sdk-go).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/milvus-io/milvus-sdk-go/blob/master/OWNERS)~~ [congqixia]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->

Patch:
diff --git a/client/client.go b/client/client.go
index 123c74c..9330a0c 100644
--- a/client/client.go
+++ b/client/client.go
@@ -82,6 +82,10 @@ type Client interface {
 	// GetReplicas gets the replica groups as well as their querynodes and shards information
 	GetReplicas(ctx context.Context, collName string) ([]*entity.ReplicaGroup, error)
 
+	// BackupRBAC back up rbac meta info, include user/role/grants
+	BackupRBAC(ctx context.Context) (*entity.RBACMeta, error)
+	RestoreRBAC(ctx context.Context, meta *entity.RBACMeta) error
+
 	// -- authentication --
 
 	// CreateCredential create new user and password
diff --git a/client/rbac.go b/client/rbac.go
index 8b71eb3..a65dc2b 100644
--- a/client/rbac.go
+++ b/client/rbac.go
@@ -18,6 +18,7 @@ package client
 
 import (
 	"context"
+	"errors"
 
 	"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
 	"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
@@ -163,7 +164,6 @@ func (c *GrpcClient) DescribeUser(ctx context.Context, username string) (entity.
 	}
 
 	resp, err := c.Service.SelectUser(ctx, req)
-
 	if err != nil {
 		return entity.UserDescription{}, err
 	}
@@ -198,7 +198,6 @@ func (c *GrpcClient) DescribeUsers(ctx context.Context) ([]entity.UserDescriptio
 	}
 
 	resp, err := c.Service.SelectUser(ctx, req)
-
 	if err != nil {
 		return nil, err
 	}
@@ -240,7 +239,6 @@ func (c *GrpcClient) ListGrants(ctx context.Context, role string, dbName string)
 	}
 
 	resp, err := c.Service.SelectGrant(ctx, req)
-
 	if err != nil {
 		return RoleGrantsList, err
 	}
@@ -290,7 +288,6 @@ func (c *GrpcClient) ListGrant(ctx context.Context, role string, object string,
 	}
 
 	resp, err := c.Service.SelectGrant(ctx, req)
-
 	if err != nil {
 		return RoleGrantsList, err
 	}
@@ -395,3 +392,132 @@ func (c *GrpcClient) Revoke(ctx context.Context, role string, objectType entity.
 
 	return handleRespStatus(resp)
 }
+
+func (c *GrpcClient) BackupRBAC(ctx context.Context) (*entity.RBACMeta, error) {
+	if c.Service == nil {
+		return nil, ErrClientNotReady
+	}
+
+	req := &milvuspb.BackupRBACMetaRequest{}
+
+	resp, err := c.Service.BackupRBAC(ctx, req)
+	if err != nil {
+		return nil, err
+	}
+	if err = handleRespStatus(resp.GetStatus()); err != nil {
+		return nil, err
+	}
+
+	meta := resp.GetRBACMeta()
+	users := make([]*entity.UserInfo, 0, len(meta.GetUsers()))
+	for _, user := range meta.GetUsers() {
+		roles := make([]string, 0, len(user.GetRoles()))
+		for _, role := range user.GetRoles() {
+			roles = append(roles, role.GetName())
+		}
+
+		users = append(users, &entity.UserInfo{
+			UserDescription: entity.UserDescription{
+				Name:  user.GetUser(),
+				Roles: roles,
+			},
+			Password: user.GetPassword(),
+		})
+	}
+
+	roles := []*entity.Role{}
+	for _, role := range meta.GetRoles() {
+		roles = append(roles, &entity.Role{
+			Name: role.GetName(),
+		})
+	}
+
+	roleGrants := []*entity.RoleGrants{}
+	for _, grant := range meta.GetGrants() {
+		roleGrant := &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(),
+		}
+
+		roleGrants = append(roleGrants, roleGrant)
+	}
+
+	return &entity.RBACMeta{
+		Users:      users,
+		Roles:      roles,
+		RoleGrants: roleGrants,
+	}, nil
+}
+
+func (c *GrpcClient) RestoreRBAC(ctx context.Context, meta *entity.RBACMeta) error {
+	if c.Service == nil {
+		return ErrClientNotReady
+	}
+
+	if meta == nil {
+		return errors.New("failed to restore rbac meta, meta is nil")
+	}
+
+	users := make([]*milvuspb.UserInfo, 0, len(meta.Users))
+	for _, user := range meta.Users {
+		roles := make([]*milvuspb.RoleEntity, 0, len(user.Roles))
+		for _, role := range meta.Roles {
+			roles = append(roles, &milvuspb.RoleEntity{
+				Name: role.Name,
+			})
+		}
+		users = append(users, &milvuspb.UserInfo{
+			User:     user.Name,
+			Password: user.Password,
+			Roles:    roles,
+		})
+	}
+
+	roles := make([]*milvuspb.RoleEntity, 0, len(meta.Roles))
+	for _, role := range meta.Roles {
+		roles = append(roles, &milvuspb.RoleEntity{
+			Name: role.Name,
+		})
+	}
+
+	grants := make([]*milvuspb.GrantEntity, 0, len(meta.RoleGrants))
+	for _, grant := range meta.RoleGrants {
+		grants = append(grants, &milvuspb.GrantEntity{
+			Role: &milvuspb.RoleEntity{
+				Name: grant.RoleName,
+			},
+			Object: &milvuspb.ObjectEntity{
+				Name: grant.Object,
+			},
+			ObjectName: grant.ObjectName,
+			Grantor: &milvuspb.GrantorEntity{
+				User: &milvuspb.UserEntity{
+					Name: grant.GrantorName,
+				},
+				Privilege: &milvuspb.PrivilegeEntity{
+					Name: grant.PrivilegeName,
+				},
+			},
+			DbName: grant.DbName,
+		})
+	}
+
+	req := &milvuspb.RestoreRBACMetaRequest{
+		RBACMeta: &milvuspb.RBACMeta{
+			Users:  users,
+			Roles:  roles,
+			Grants: grants,
+		},
+	}
+
+	resp, err := c.Service.RestoreRBAC(ctx, req)
+	if err != nil {
+		return err
+	}
+
+	return handleRespStatus(resp)
+}
diff --git a/entity/rbac.go b/entity/rbac.go
index f215f28..22f473e 100644
--- a/entity/rbac.go
+++ b/entity/rbac.go
@@ -61,3 +61,14 @@ func WithOperatePrivilegeDatabase(database string) OperatePrivilegeOption {
 		o.Database = database
 	}
 }
+
+type UserInfo struct {
+	UserDescription
+	Password string
+}
+
+type RBACMeta struct {
+	Users      []*UserInfo
+	Roles      []*Role
+	RoleGrants []*RoleGrants
+}
diff --git a/examples/rbac/main.go b/examples/rbac/main.go
new file mode 100644
index 0000000..7e015e6
--- /dev/null
+++ b/examples/rbac/main.go
@@ -0,0 +1,69 @@
+package main
+
+import (
+	"context"
+	"log"
+
+	"github.com/milvus-io/milvus-sdk-go/v2/client"
+	"github.com/milvus-io/milvus-sdk-go/v2/entity"
+)
+
+const (
+	milvusAddr     = `localhost:19530`
+	nEntities, dim = 10000, 128
+	collectionName = "hello_partition_key"
+
+	idCol, keyCol, embeddingCol = "ID", "key", "embeddings"
+	topK                        = 3
+)
+
+func main() {
+	ctx := context.Background()
+
+	log.Println("start connecting to Milvus")
+	c, err := client.NewClient(ctx, client.Config{
+		Address: milvusAddr,
+	})
+	if err != nil {
+		log.Fatalf("failed to connect to milvus, err: %v", err)
+	}
+	defer c.Close()
+
+	// create user
+	err = c.CreateCredential(ctx, "user123", "passwd1")
+	if err != nil {
+		log.Fatalf("failed to create user, err: %v", err)
+	}
+
+	// create role
+	c.CreateRole(ctx, "role123")
+	if err != nil {
+		log.Fatalf("failed to create role, err: %v", err)
+	}
+	c.Grant(ctx, "role123", entity.PriviledegeObjectTypeGlobal, "*", "read")
+
+	// grant role to user
+	c.AddUserRole(ctx, "user123", "role123")
+
+	// backup rbac
+	meta, err := c.BackupRBAC(ctx)
+	if err != nil {
+		log.Fatalf("failed to backup rbac, err: %v", err)
+	}
+
+	// clean rbac to make restore works
+	c.DropRole(ctx, "role123")
+	c.DeleteCredential(ctx, "user123")
+	c.Revoke(ctx, "role123", entity.PriviledegeObjectTypeGlobal, "*", "read")
+
+	// restore rbac
+	err = c.RestoreRBAC(ctx, meta)
+	if err != nil {
+		log.Fatalf("failed to restore rbac, err: %v", err)
+	}
+
+	// clean rbac
+	c.DropRole(ctx, "role123")
+	c.DeleteCredential(ctx, "user123")
+	c.Revoke(ctx, "role123", entity.PriviledegeObjectTypeGlobal, "*", "read")
+}
diff --git a/go.mod b/go.mod
index 5164c9b..7108f99 100644
--- a/go.mod
+++ b/go.mod
@@ -7,7 +7,7 @@ require (
 	github.com/go-faker/faker/v4 v4.1.0
 	github.com/golang/protobuf v1.5.2
 	github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
-	github.com/milvus-io/milvus-proto/go-api/v2 v2.4.6
+	github.com/milvus-io/milvus-proto/go-api/v2 v2.4.10-0.20240819025435-512e3b98866a
 	github.com/stretchr/testify v1.8.1
 	github.com/tidwall/gjson v1.14.4
 	github.com/x448/float16 v0.8.4
diff --git a/go.sum b/go.sum
index 71106d8..9ca5360 100644
--- a/go.sum
+++ b/go.sum
@@ -157,8 +157,8 @@ github.com/mattn/go-isatty v0.0.14/go.mod h1:7GGIvUiUoEMVVmxf/4nioHXj79iQHKdU27k
 github.com/mattn/goveralls v0.0.2/go.mod h1:8d1ZMHsd7fW6IRPKQh46F2WRpyib5/X4FOpevwGNQEw=
 github.com/mediocregopher/radix/v3 v3.4.2/go.mod h1:8FL3F6UQRXHXIBSPUs5h0RybMF8i4n7wVopoX3x7Bv8=
 github.com/microcosm-cc/bluemonday v1.0.2/go.mod h1:iVP4YcDBq+n/5fb23BhYFvIMq/leAFZyRl6bYmGDlGc=
-github.com/milvus-io/milvus-proto/go-api/v2 v2.4.6 h1:rSkwp5Mg/7KBSUqXcrPBUgTQGZNdvYWEKB+rHo9YJtk=
-github.com/milvus-io/milvus-proto/go-api/v2 v2.4.6/go.mod h1:1OIl0v5PQeNxIJhCvY+K55CBUOYDZevw9g9380u1Wek=
+github.com/milvus-io/milvus-proto/go-api/v2 v2.4.10-0.20240819025435-512e3b98866a h1:0B/8Fo66D8Aa23Il0yrQvg1KKz92tE/BJ5BvkUxxAAk=
+github.com/milvus-io/milvus-proto/go-api/v2 v2.4.10-0.20240819025435-512e3b98866a/go.mod h1:1OIl0v5PQeNxIJhCvY+K55CBUOYDZevw9g9380u1Wek=
 github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=
 github.com/mitchellh/mapstructure v1.1.2/go.mod h1:FVVH3fgwuzCH5S8UJGiWEs2h04kUh9fWfEaFds41c1Y=
 github.com/modern-go/concurrent v0.0.0-20180228061459-e0a39a4cb421/go.mod h1:6dJC0mAP4ikYIbvyc7fijjWJddQyLn8Ig3JB5CqoB9Q=
diff --git a/mocks/MilvusServiceServer.go b/mocks/MilvusServiceServer.go
index 95f06fa..2ae5737 100644
--- a/mocks/MilvusServiceServer.go
+++ b/mocks/MilvusServiceServer.go
@@ -302,6 +302,61 @@ func (_c *MilvusServiceServer_AlterIndex_Call) RunAndReturn(run func(context.Con
 	return _c
 }
 
+// BackupRBAC provides a mock function with given fields: _a0, _a1
+func (_m *MilvusServiceServer) BackupRBAC(_a0 context.Context, _a1 *milvuspb.BackupRBACMetaRequest) (*milvuspb.BackupRBACMetaResponse, error) {
+	ret := _m.Called(_a0, _a1)
+
+	var r0 *milvuspb.BackupRBACMetaResponse
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.BackupRBACMetaRequest) (*milvuspb.BackupRBACMetaResponse, error)); ok {
+		return rf(_a0, _a1)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.BackupRBACMetaRequest) *milvuspb.BackupRBACMetaResponse); ok {
+		r0 = rf(_a0, _a1)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*milvuspb.BackupRBACMetaResponse)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, *milvuspb.BackupRBACMetaRequest) error); ok {
+		r1 = rf(_a0, _a1)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MilvusServiceServer_BackupRBAC_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'BackupRBAC'
+type MilvusServiceServer_BackupRBAC_Call struct {
+	*mock.Call
+}
+
+// BackupRBAC is a helper method to define mock.On call
+//   - _a0 context.Context
+//   - _a1 *milvuspb.BackupRBACMetaRequest
+func (_e *MilvusServiceServer_Expecter) BackupRBAC(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_BackupRBAC_Call {
+	return &MilvusServiceServer_BackupRBAC_Call{Call: _e.mock.On("BackupRBAC", _a0, _a1)}
+}
+
+func (_c *MilvusServiceServer_BackupRBAC_Call) Run(run func(_a0 context.Context, _a1 *milvuspb.BackupRBACMetaRequest)) *MilvusServiceServer_BackupRBAC_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(*milvuspb.BackupRBACMetaRequest))
+	})
+	return _c
+}
+
+func (_c *MilvusServiceServer_BackupRBAC_Call) Return(_a0 *milvuspb.BackupRBACMetaResponse, _a1 error) *MilvusServiceServer_BackupRBAC_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MilvusServiceServer_BackupRBAC_Call) RunAndReturn(run func(context.Context, *milvuspb.BackupRBACMetaRequest) (*milvuspb.BackupRBACMetaResponse, error)) *MilvusServiceServer_BackupRBAC_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // CalcDistance provides a mock function with given fields: _a0, _a1
 func (_m *MilvusServiceServer) CalcDistance(_a0 context.Context, _a1 *milvuspb.CalcDistanceRequest) (*milvuspb.CalcDistanceResults, error) {
 	ret := _m.Called(_a0, _a1)
@@ -4152,6 +4207,61 @@ func (_c *MilvusServiceServer_ReplicateMessage_Call) RunAndReturn(run func(conte
 	return _c
 }
 
+// RestoreRBAC provides a mock function with given fields: _a0, _a1
+func (_m *MilvusServiceServer) RestoreRBAC(_a0 context.Context, _a1 *milvuspb.RestoreRBACMetaRequest) (*commonpb.Status, error) {
+	ret := _m.Called(_a0, _a1)
+
+	var r0 *commonpb.Status
+	var r1 error
+	if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.RestoreRBACMetaRequest) (*commonpb.Status, error)); ok {
+		return rf(_a0, _a1)
+	}
+	if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.RestoreRBACMetaRequest) *commonpb.Status); ok {
+		r0 = rf(_a0, _a1)
+	} else {
+		if ret.Get(0) != nil {
+			r0 = ret.Get(0).(*commonpb.Status)
+		}
+	}
+
+	if rf, ok := ret.Get(1).(func(context.Context, *milvuspb.RestoreRBACMetaRequest) error); ok {
+		r1 = rf(_a0, _a1)
+	} else {
+		r1 = ret.Error(1)
+	}
+
+	return r0, r1
+}
+
+// MilvusServiceServer_RestoreRBAC_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RestoreRBAC'
+type MilvusServiceServer_RestoreRBAC_Call struct {
+	*mock.Call
+}
+
+// RestoreRBAC is a helper method to define mock.On call
+//   - _a0 context.Context
+//   - _a1 *milvuspb.RestoreRBACMetaRequest
+func (_e *MilvusServiceServer_Expecter) RestoreRBAC(_a0 interface{}, _a1 interface{}) *MilvusServiceServer_RestoreRBAC_Call {
+	return &MilvusServiceServer_RestoreRBAC_Call{Call: _e.mock.On("RestoreRBAC", _a0, _a1)}
+}
+
+func (_c *MilvusServiceServer_RestoreRBAC_Call) Run(run func(_a0 context.Context, _a1 *milvuspb.RestoreRBACMetaRequest)) *MilvusServiceServer_RestoreRBAC_Call {
+	_c.Call.Run(func(args mock.Arguments) {
+		run(args[0].(context.Context), args[1].(*milvuspb.RestoreRBACMetaRequest))
+	})
+	return _c
+}
+
+func (_c *MilvusServiceServer_RestoreRBAC_Call) Return(_a0 *commonpb.Status, _a1 error) *MilvusServiceServer_RestoreRBAC_Call {
+	_c.Call.Return(_a0, _a1)
+	return _c
+}
+
+func (_c *MilvusServiceServer_RestoreRBAC_Call) RunAndReturn(run func(context.Context, *milvuspb.RestoreRBACMetaRequest) (*commonpb.Status, error)) *MilvusServiceServer_RestoreRBAC_Call {
+	_c.Call.Return(run)
+	return _c
+}
+
 // Search provides a mock function with given fields: _a0, _a1
 func (_m *MilvusServiceServer) Search(_a0 context.Context, _a1 *milvuspb.SearchRequest) (*milvuspb.SearchResults, error) {
 	ret := _m.Called(_a0, _a1)


Test patch:
diff --git a/client/client_mock_test.go b/client/client_mock_test.go
index 5ef7dcb..38b7525 100644
--- a/client/client_mock_test.go
+++ b/client/client_mock_test.go
@@ -307,6 +307,8 @@ const (
 	MDescribeDatabase ServiceMethod = 1004
 
 	MReplicateMessage ServiceMethod = 1100
+	MBackupRBAC       ServiceMethod = 1101
+	MRestoreRBAC      ServiceMethod = 1102
 )
 
 // injection function definition
@@ -1055,3 +1057,22 @@ func (m *MockServer) HybridSearch(ctx context.Context, req *milvuspb.HybridSearc
 func (m *MockServer) UpdateResourceGroups(_ context.Context, _ *milvuspb.UpdateResourceGroupsRequest) (*commonpb.Status, error) {
 	return nil, errors.New("not implemented")
 }
+
+func (m *MockServer) BackupRBAC(ctx context.Context, req *milvuspb.BackupRBACMetaRequest) (*milvuspb.BackupRBACMetaResponse, error) {
+	f := m.GetInjection(MBackupRBAC)
+	if f != nil {
+		r, err := f(ctx, req)
+		return r.(*milvuspb.BackupRBACMetaResponse), err
+	}
+	s, err := SuccessStatus()
+	return &milvuspb.BackupRBACMetaResponse{Status: s}, err
+}
+
+func (m *MockServer) RestoreRBAC(ctx context.Context, req *milvuspb.RestoreRBACMetaRequest) (*commonpb.Status, error) {
+	f := m.GetInjection(MRestoreRBAC)
+	if f != nil {
+		r, err := f(ctx, req)
+		return r.(*commonpb.Status), err
+	}
+	return SuccessStatus()
+}
diff --git a/client/client_test.go b/client/client_test.go
index 86f3911..136013f 100644
--- a/client/client_test.go
+++ b/client/client_test.go
@@ -155,7 +155,7 @@ func TestGrpcClientNil(t *testing.T) {
 				m.Name == "HybridSearch" || // type alias MetricType treated as string
 				m.Name == "CalcDistance" ||
 				m.Name == "ManualCompaction" || // time.Duration hard to detect in reflect
-				m.Name == "Insert" || m.Name == "Upsert" { // complex methods with ...
+				m.Name == "Insert" || m.Name == "Upsert" || m.Name == "RestoreRBAC" { // complex methods with ...
 				t.Skip("method", m.Name, "skipped")
 			}
 			ins := make([]reflect.Value, 0, mt.NumIn())
diff --git a/client/rbac_test.go b/client/rbac_test.go
index 666406f..1306186 100644
--- a/client/rbac_test.go
+++ b/client/rbac_test.go
@@ -7,6 +7,7 @@ import (
 	"github.com/cockroachdb/errors"
 
 	"github.com/milvus-io/milvus-sdk-go/v2/entity"
+	"github.com/milvus-io/milvus-sdk-go/v2/merr"
 	"github.com/stretchr/testify/suite"
 
 	"github.com/stretchr/testify/mock"
@@ -740,6 +741,166 @@ func (s *RBACSuite) TestRevoke() {
 	})
 }
 
+func (s *RBACSuite) TestBackup() {
+	ctx := context.Background()
+
+	s.Run("normal run", func() {
+		ctx, cancel := context.WithCancel(ctx)
+		defer cancel()
+		defer s.resetMock()
+		s.mock.EXPECT().BackupRBAC(mock.Anything, mock.Anything).Run(func(_ context.Context, _ *milvuspb.BackupRBACMetaRequest) {
+		}).Return(&milvuspb.BackupRBACMetaResponse{
+			Status: merr.Success(),
+			RBACMeta: &milvuspb.RBACMeta{
+				Users: []*milvuspb.UserInfo{
+					{
+						User:     "user1",
+						Password: "passwd",
+						Roles: []*milvuspb.RoleEntity{
+							{Name: "role1"},
+						},
+					},
+				},
+				Roles: []*milvuspb.RoleEntity{
+					{Name: "role1"},
+				},
+				Grants: []*milvuspb.GrantEntity{
+					{
+						Object: &milvuspb.ObjectEntity{
+							Name: "testObject",
+						},
+						ObjectName: "testObjectName",
+						Role: &milvuspb.RoleEntity{
+							Name: "testRole",
+						},
+						Grantor: &milvuspb.GrantorEntity{
+							User: &milvuspb.UserEntity{
+								Name: "grantorUser",
+							},
+							Privilege: &milvuspb.PrivilegeEntity{
+								Name: "testPrivilege",
+							},
+						},
+						DbName: "testDB",
+					},
+				},
+			},
+		}, nil)
+
+		meta, err := s.client.BackupRBAC(ctx)
+		s.NoError(err)
+		s.Len(meta.Users, 1)
+		s.Len(meta.Roles, 1)
+		s.Len(meta.RoleGrants, 1)
+	})
+
+	s.Run("rpc error", func() {
+		ctx, cancel := context.WithCancel(ctx)
+		defer cancel()
+		defer s.resetMock()
+		s.mock.EXPECT().BackupRBAC(mock.Anything, mock.Anything).Return(nil, errors.New("mock error"))
+
+		_, err := s.client.BackupRBAC(ctx)
+		s.Error(err)
+	})
+
+	s.Run("status error", func() {
+		ctx, cancel := context.WithCancel(ctx)
+		defer cancel()
+		defer s.resetMock()
+		s.mock.EXPECT().BackupRBAC(mock.Anything, mock.Anything).Return(&milvuspb.BackupRBACMetaResponse{
+			Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_UnexpectedError},
+		}, nil)
+
+		_, err := s.client.BackupRBAC(ctx)
+		s.Error(err)
+	})
+
+	s.Run("service not ready", func() {
+		ctx, cancel := context.WithCancel(ctx)
+		defer cancel()
+		defer s.resetMock()
+
+		c := &GrpcClient{}
+		_, err := c.BackupRBAC(ctx)
+		s.Error(err)
+		s.ErrorIs(err, ErrClientNotReady)
+	})
+}
+
+func (s *RBACSuite) TestRestore() {
+	ctx := context.Background()
+
+	s.Run("normal run", func() {
+		ctx, cancel := context.WithCancel(ctx)
+		defer cancel()
+		defer s.resetMock()
+
+		s.mock.EXPECT().RestoreRBAC(mock.Anything, mock.Anything).Run(func(_ context.Context, req *milvuspb.RestoreRBACMetaRequest) {
+			s.Len(req.RBACMeta.Users, 1)
+			s.Len(req.RBACMeta.Roles, 1)
+			s.Len(req.RBACMeta.Grants, 1)
+		}).Return(merr.Success(), nil)
+
+		err := s.client.RestoreRBAC(ctx, &entity.RBACMeta{
+			Users: []*entity.UserInfo{
+				{
+					UserDescription: entity.UserDescription{
+						Name:  "user1",
+						Roles: []string{"role1"},
+					},
+					Password: "passwd",
+				},
+			},
+			Roles: []*entity.Role{
+				{Name: "role1"},
+			},
+			RoleGrants: []*entity.RoleGrants{
+				{
+					Object:        "testObject",
+					ObjectName:    "testObjectName",
+					RoleName:      "testRole",
+					GrantorName:   "grantorUser",
+					PrivilegeName: "testPrivilege",
+					DbName:        "testDB",
+				},
+			},
+		})
+		s.NoError(err)
+	})
+
+	s.Run("rpc error", func() {
+		ctx, cancel := context.WithCancel(ctx)
+		defer cancel()
+		defer s.resetMock()
+		s.mock.EXPECT().RestoreRBAC(mock.Anything, mock.Anything).Return(nil, errors.New("mock error"))
+
+		err := s.client.RestoreRBAC(ctx, &entity.RBACMeta{})
+		s.Error(err)
+	})
+
+	s.Run("status error", func() {
+		ctx, cancel := context.WithCancel(ctx)
+		defer cancel()
+		defer s.resetMock()
+		s.mock.EXPECT().RestoreRBAC(mock.Anything, mock.Anything).Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_UnexpectedError}, nil)
+
+		err := s.client.RestoreRBAC(ctx, &entity.RBACMeta{})
+		s.Error(err)
+	})
+
+	s.Run("service not ready", func() {
+		ctx, cancel := context.WithCancel(ctx)
+		defer cancel()
+		defer s.resetMock()
+
+		c := &GrpcClient{}
+		err := c.RestoreRBAC(ctx, &entity.RBACMeta{})
+		s.Error(err)
+		s.ErrorIs(err, ErrClientNotReady)
+	})
+}
+
 func TestRBACSuite(t *testing.T) {
 	suite.Run(t, new(RBACSuite))
 }