milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

enhance: Add BackupRBAC/RestoreRBAC API to enable rbac backup (#35444) 

issue: #35443

---------

Signed-off-by: Wei Liu <wei.liu@zilliz.com>
pull/35507/head
wei liu 2024-08-16 10:10:53 +08:00 committed by GitHub
parent 1275005fc1
commit 1d49358f82
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
26 changed files with 1420 additions and 18 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
go.mod
View File

@ -22,7 +22,7 @@ require (
github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
github.com/klauspost/compress v1.17.7
github.com/mgutz/ansi v0.0.0-20200706080929-d51e80ef957d
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240717062137-3ffb1db01632
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240815113856-e2789dca8b59
github.com/minio/minio-go/v7 v7.0.61
github.com/pingcap/log v1.1.1-0.20221015072633-39906604fb81
github.com/prometheus/client_golang v1.14.0
@ -57,6 +57,7 @@ require (
require (
github.com/bits-and-blooms/bitset v1.10.0
github.com/bytedance/sonic v1.9.1
github.com/cenkalti/backoff/v4 v4.2.1
github.com/cockroachdb/redact v1.1.3
github.com/greatroar/blobloom v0.0.0-00010101000000-000000000000
@ -91,7 +92,6 @@ require (
github.com/benbjohnson/clock v1.1.0 // indirect
github.com/benesch/cgosymbolizer v0.0.0-20190515212042-bec6fe6e597b // indirect
github.com/beorn7/perks v1.0.1 // indirect
github.com/bytedance/sonic v1.9.1 // indirect
github.com/campoy/embedmd v1.0.0 // indirect
github.com/cespare/xxhash/v2 v2.3.0 // indirect
github.com/chenzhuoyu/base64x v0.0.0-20221115062448-fe3a3abad311 // indirect

4
go.sum
View File

@ -598,8 +598,8 @@ github.com/milvus-io/cgosymbolizer v0.0.0-20240722103217-b7dee0e50119 h1:9VXijWu
github.com/milvus-io/cgosymbolizer v0.0.0-20240722103217-b7dee0e50119/go.mod h1:DvXTE/K/RtHehxU8/GtDs4vFtfw64jJ3PaCnFri8CRg=
github.com/milvus-io/gorocksdb v0.0.0-20220624081344-8c5f4212846b h1:TfeY0NxYxZzUfIfYe5qYDBzt4ZYRqzUjTR6CvUzjat8=
github.com/milvus-io/gorocksdb v0.0.0-20220624081344-8c5f4212846b/go.mod h1:iwW+9cWfIzzDseEBCCeDSN5SD16Tidvy8cwQ7ZY8Qj4=
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240717062137-3ffb1db01632 h1:CXig0DNtUsCLzchCFe3PR2KgOdobbz9gK2nSV7195PM=
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240717062137-3ffb1db01632/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs=
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240815113856-e2789dca8b59 h1:mKekr0GmCKMpIQh9OJ67TlKVKxDt08600ltARc/JUXY=
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240815113856-e2789dca8b59/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs=
github.com/milvus-io/pulsar-client-go v0.6.10 h1:eqpJjU+/QX0iIhEo3nhOqMNXL+TyInAs1IAHZCrCM/A=
github.com/milvus-io/pulsar-client-go v0.6.10/go.mod h1:lQqCkgwDF8YFYjKA+zOheTk1tev2B+bKj5j7+nm8M1w=
github.com/minio/asm2plan9s v0.0.0-20200509001527-cdd76441f9d8 h1:AMFGa4R4MiIpspGNG7Z948v4n35fFGB3RR3G/ry4FWs=

8
internal/datacoord/mock_test.go
View File

@ -624,6 +624,14 @@ func (m *mockRootCoordClient) GetMetrics(ctx context.Context, req *milvuspb.GetM
}, nil
}
func (m *mockRootCoordClient) BackupRBAC(ctx context.Context, req *milvuspb.BackupRBACMetaRequest, opts ...grpc.CallOption) (*milvuspb.BackupRBACMetaResponse, error) {
panic("not implemented") // TODO: Implement
}
func (m *mockRootCoordClient) RestoreRBAC(ctx context.Context, req *milvuspb.RestoreRBACMetaRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
panic("not implemented") // TODO: Implement
}
type mockCompactionTrigger struct {
methods map[string]interface{}
}

8
internal/distributed/proxy/service.go
View File

@ -1145,6 +1145,14 @@ func (s *Server) SelectGrant(ctx context.Context, req *milvuspb.SelectGrantReque
return s.proxy.SelectGrant(ctx, req)
}
func (s *Server) BackupRBAC(ctx context.Context, req *milvuspb.BackupRBACMetaRequest) (*milvuspb.BackupRBACMetaResponse, error) {
return s.proxy.BackupRBAC(ctx, req)
}
func (s *Server) RestoreRBAC(ctx context.Context, req *milvuspb.RestoreRBACMetaRequest) (*commonpb.Status, error) {
return s.proxy.RestoreRBAC(ctx, req)
}
func (s *Server) RefreshPolicyInfoCache(ctx context.Context, req *proxypb.RefreshPolicyInfoCacheRequest) (*commonpb.Status, error) {
return s.proxy.RefreshPolicyInfoCache(ctx, req)
}

24
internal/distributed/rootcoord/client/client.go
View File

@ -682,3 +682,27 @@ func (c *Client) AlterDatabase(ctx context.Context, request *rootcoordpb.AlterDa
return client.AlterDatabase(ctx, request)
})
}
func (c *Client) BackupRBAC(ctx context.Context, in *milvuspb.BackupRBACMetaRequest, opts ...grpc.CallOption) (*milvuspb.BackupRBACMetaResponse, error) {
in = typeutil.Clone(in)
commonpbutil.UpdateMsgBase(
in.GetBase(),
commonpbutil.FillMsgBaseFromClient(paramtable.GetNodeID(), commonpbutil.WithTargetID(c.sess.ServerID)),
)
return wrapGrpcCall(ctx, c, func(client rootcoordpb.RootCoordClient) (*milvuspb.BackupRBACMetaResponse, error) {
return client.BackupRBAC(ctx, in)
})
}
func (c *Client) RestoreRBAC(ctx context.Context, in *milvuspb.RestoreRBACMetaRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
in = typeutil.Clone(in)
commonpbutil.UpdateMsgBase(
in.GetBase(),
commonpbutil.FillMsgBaseFromClient(paramtable.GetNodeID(), commonpbutil.WithTargetID(c.sess.ServerID)),
)
return wrapGrpcCall(ctx, c, func(client rootcoordpb.RootCoordClient) (*commonpb.Status, error) {
return client.RestoreRBAC(ctx, in)
})
}

8
internal/distributed/rootcoord/service.go
View File

@ -533,3 +533,11 @@ func (s *Server) AlterCollection(ctx context.Context, request *milvuspb.AlterCol
func (s *Server) RenameCollection(ctx context.Context, request *milvuspb.RenameCollectionRequest) (*commonpb.Status, error) {
return s.rootCoord.RenameCollection(ctx, request)
}
func (s *Server) BackupRBAC(ctx context.Context, request *milvuspb.BackupRBACMetaRequest) (*milvuspb.BackupRBACMetaResponse, error) {
return s.rootCoord.BackupRBAC(ctx, request)
}
func (s *Server) RestoreRBAC(ctx context.Context, request *milvuspb.RestoreRBACMetaRequest) (*commonpb.Status, error) {
return s.rootCoord.RestoreRBAC(ctx, request)
}

4
internal/metastore/catalog.go
View File

@ -82,6 +82,10 @@ type RootCoordCatalog interface {
// For example []string{"user1/role1"}
ListUserRole(ctx context.Context, tenant string) ([]string, error)
ListCredentialsWithPasswd(ctx context.Context) (map[string]string, error)
BackupRBAC(ctx context.Context, tenant string) (*milvuspb.RBACMeta, error)
RestoreRBAC(ctx context.Context, tenant string, meta *milvuspb.RBACMeta) error
Close()
}

184
internal/metastore/kv/rootcoord/kv_catalog.go
View File

@ -6,6 +6,7 @@ import (
"fmt"
"github.com/cockroachdb/errors"
"github.com/samber/lo"
"go.uber.org/zap"
"google.golang.org/protobuf/proto"
@ -740,23 +741,37 @@ func (kc *Catalog) ListAliases(ctx context.Context, dbID int64, ts typeutil.Time
}
func (kc *Catalog) ListCredentials(ctx context.Context) ([]string, error) {
keys, _, err := kc.Txn.LoadWithPrefix(CredentialPrefix)
users, err := kc.ListCredentialsWithPasswd(ctx)
if err != nil {
return nil, err
}
return lo.Keys(users), nil
}
func (kc *Catalog) ListCredentialsWithPasswd(ctx context.Context) (map[string]string, error) {
keys, values, err := kc.Txn.LoadWithPrefix(CredentialPrefix)
if err != nil {
log.Error("list all credential usernames fail", zap.String("prefix", CredentialPrefix), zap.Error(err))
return nil, err
}
var usernames []string
for _, path := range keys {
username := typeutil.After(path, UserSubPrefix+"/")
users := make(map[string]string)
for i := range keys {
username := typeutil.After(keys[i], UserSubPrefix+"/")
if len(username) == 0 {
log.Warn("no username extract from path:", zap.String("path", path))
log.Warn("no username extract from path:", zap.String("path", keys[i]))
continue
}
usernames = append(usernames, username)
credential := &internalpb.CredentialInfo{}
err := json.Unmarshal([]byte(values[i]), credential)
if err != nil {
log.Error("credential unmarshal fail", zap.String("key", keys[i]), zap.Error(err))
return nil, err
}
users[username] = credential.EncryptedPassword
}
return usernames, nil
return users, nil
}
func (kc *Catalog) save(k string) error {
@ -1210,6 +1225,161 @@ func (kc *Catalog) ListUserRole(ctx context.Context, tenant string) ([]string, e
return userRoles, nil
}
func (kc *Catalog) BackupRBAC(ctx context.Context, tenant string) (*milvuspb.RBACMeta, error) {
users, err := kc.ListUser(ctx, tenant, nil, true)
if err != nil {
return nil, err
}
credentials, err := kc.ListCredentialsWithPasswd(ctx)
if err != nil {
return nil, err
}
userInfos := lo.FilterMap(users, func(entity *milvuspb.UserResult, _ int) (*milvuspb.UserInfo, bool) {
userName := entity.GetUser().GetName()
if userName == util.UserRoot {
return nil, false
}
return &milvuspb.UserInfo{
User: userName,
Password: credentials[userName],
Roles: entity.GetRoles(),
}, true
})
roles, err := kc.ListRole(ctx, tenant, nil, false)
if err != nil {
return nil, err
}
roleEntity := lo.FilterMap(roles, func(entity *milvuspb.RoleResult, _ int) (*milvuspb.RoleEntity, bool) {
roleName := entity.GetRole().GetName()
if roleName == util.RoleAdmin || roleName == util.RolePublic {
return nil, false
}
return entity.GetRole(), true
})
grantsEntity := make([]*milvuspb.GrantEntity, 0)
for _, role := range roleEntity {
grants, err := kc.ListGrant(ctx, tenant, &milvuspb.GrantEntity{
Role: role,
DbName: util.AnyWord,
})
if err != nil {
return nil, err
}
grantsEntity = append(grantsEntity, grants...)
}
return &milvuspb.RBACMeta{
Users: userInfos,
Roles: roleEntity,
Grants: grantsEntity,
}, nil
}
func (kc *Catalog) RestoreRBAC(ctx context.Context, tenant string, meta *milvuspb.RBACMeta) error {
var err error
needRollbackUser := make([]*milvuspb.UserInfo, 0)
needRollbackRole := make([]*milvuspb.RoleEntity, 0)
needRollbackGrants := make([]*milvuspb.GrantEntity, 0)
defer func() {
if err != nil {
log.Warn("failed to restore rbac, try to rollback", zap.Error(err))
// roll back role
for _, role := range needRollbackRole {
err = kc.DropRole(ctx, tenant, role.Name)
if err != nil {
log.Warn("failed to rollback roles after restore failed", zap.Error(err))
}
}
// roll back grant
for _, grant := range needRollbackGrants {
err = kc.AlterGrant(ctx, tenant, grant, milvuspb.OperatePrivilegeType_Revoke)
if err != nil {
log.Warn("failed to rollback grants after restore failed", zap.Error(err))
}
}
for _, user := range needRollbackUser {
// roll back user
err = kc.DropCredential(ctx, user.User)
if err != nil {
log.Warn("failed to rollback users after restore failed", zap.Error(err))
}
}
}
}()
// restore role
existRoles, err := kc.ListRole(ctx, tenant, nil, false)
if err != nil {
return err
}
existRoleMap := lo.SliceToMap(existRoles, func(entity *milvuspb.RoleResult) (string, struct{}) { return entity.GetRole().GetName(), struct{}{} })
for _, role := range meta.Roles {
if _, ok := existRoleMap[role.GetName()]; ok {
log.Warn("failed to restore, role already exists", zap.String("role", role.GetName()))
err = errors.Newf("role [%s] already exists", role.GetName())
return err
}
err = kc.CreateRole(ctx, tenant, role)
if err != nil {
return err
}
needRollbackRole = append(needRollbackRole, role)
}
// restore grant
for _, grant := range meta.Grants {
err = kc.AlterGrant(ctx, tenant, grant, milvuspb.OperatePrivilegeType_Grant)
if err != nil {
return err
}
needRollbackGrants = append(needRollbackGrants, grant)
}
// need rollback user
existUser, err := kc.ListUser(ctx, tenant, nil, false)
if err != nil {
return err
}
existUserMap := lo.SliceToMap(existUser, func(entity *milvuspb.UserResult) (string, struct{}) { return entity.GetUser().GetName(), struct{}{} })
for _, user := range meta.Users {
if _, ok := existUserMap[user.GetUser()]; ok {
log.Info("failed to restore, user already exists", zap.String("user", user.GetUser()))
err = errors.Newf("user [%s] already exists", user.GetUser())
return err
}
// restore user
err = kc.CreateCredential(ctx, &model.Credential{
Username: user.User,
EncryptedPassword: user.Password,
})
if err != nil {
return err
}
needRollbackUser = append(needRollbackUser, user)
// restore user role mapping
entity := &milvuspb.UserEntity{
Name: user.User,
}
for _, role := range user.Roles {
err = kc.AlterUserRole(ctx, tenant, entity, role, milvuspb.OperateUserRoleType_AddUserToRole)
if err != nil {
return err
}
}
}
return err
}
func (kc *Catalog) Close() {
// do nothing
}

206
internal/metastore/kv/rootcoord/kv_catalog_test.go
View File

@ -21,6 +21,7 @@ import (
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus-proto/go-api/v2/schemapb"
"github.com/milvus-io/milvus/internal/kv"
etcdkv "github.com/milvus-io/milvus/internal/kv/etcd"
"github.com/milvus-io/milvus/internal/kv/mocks"
"github.com/milvus-io/milvus/internal/metastore"
"github.com/milvus-io/milvus/internal/metastore/model"
@ -30,6 +31,7 @@ import (
"github.com/milvus-io/milvus/pkg/log"
"github.com/milvus-io/milvus/pkg/util"
"github.com/milvus-io/milvus/pkg/util/crypto"
"github.com/milvus-io/milvus/pkg/util/etcd"
"github.com/milvus-io/milvus/pkg/util/funcutil"
"github.com/milvus-io/milvus/pkg/util/merr"
"github.com/milvus-io/milvus/pkg/util/typeutil"
@ -1473,7 +1475,20 @@ func TestRBAC_Credential(t *testing.T) {
}
return nil
},
nil,
func(key string) []string {
cmu.RLock()
defer cmu.RUnlock()
passwd, _ := json.Marshal(&model.Credential{EncryptedPassword: crypto.Base64Encode("passwd")})
if count == 0 {
return []string{
string(passwd),
string(passwd),
string(passwd),
string(passwd),
}
}
return nil
},
func(key string) error {
cmu.RLock()
defer cmu.RUnlock()
@ -1931,7 +1946,14 @@ func TestRBAC_Role(t *testing.T) {
}
}
return nil
}, nil,
},
func(key string) []string {
if loadCredentialPrefixReturn.Load() {
passwd, _ := json.Marshal(&model.Credential{EncryptedPassword: crypto.Base64Encode("passwd")})
return []string{string(passwd)}
}
return nil
},
func(key string) error {
if loadCredentialPrefixReturn.Load() {
return nil
@ -2544,6 +2566,186 @@ func TestRBAC_Grant(t *testing.T) {
})
}
func TestRBAC_Backup(t *testing.T) {
etcdCli, _ := etcd.GetEtcdClient(
Params.EtcdCfg.UseEmbedEtcd.GetAsBool(),
Params.EtcdCfg.EtcdUseSSL.GetAsBool(),
Params.EtcdCfg.Endpoints.GetAsStrings(),
Params.EtcdCfg.EtcdTLSCert.GetValue(),
Params.EtcdCfg.EtcdTLSKey.GetValue(),
Params.EtcdCfg.EtcdTLSCACert.GetValue(),
Params.EtcdCfg.EtcdTLSMinVersion.GetValue())
rootPath := "/test/rbac"
metaKV := etcdkv.NewEtcdKV(etcdCli, rootPath)
defer metaKV.RemoveWithPrefix("")
defer metaKV.Close()
c := &Catalog{Txn: metaKV}
ctx := context.Background()
c.CreateRole(ctx, util.DefaultTenant, &milvuspb.RoleEntity{Name: "role1"})
c.AlterGrant(ctx, util.DefaultTenant, &milvuspb.GrantEntity{
Role: &milvuspb.RoleEntity{Name: "role1"},
Object: &milvuspb.ObjectEntity{Name: "obj1"},
ObjectName: "obj_name1",
DbName: util.DefaultDBName,
Grantor: &milvuspb.GrantorEntity{
User: &milvuspb.UserEntity{Name: "user1"},
Privilege: &milvuspb.PrivilegeEntity{Name: "PrivilegeLoad"},
},
}, milvuspb.OperatePrivilegeType_Grant)
c.CreateCredential(ctx, &model.Credential{
Username: "user1",
EncryptedPassword: "passwd",
})
c.AlterUserRole(ctx, util.DefaultTenant, &milvuspb.UserEntity{Name: "user1"}, &milvuspb.RoleEntity{Name: "role1"}, milvuspb.OperateUserRoleType_AddUserToRole)
// test backup success
backup, err := c.BackupRBAC(ctx, util.DefaultTenant)
assert.NoError(t, err)
assert.Equal(t, 1, len(backup.Grants))
assert.Equal(t, "obj_name1", backup.Grants[0].ObjectName)
assert.Equal(t, "role1", backup.Grants[0].Role.Name)
assert.Equal(t, 1, len(backup.Users))
assert.Equal(t, "user1", backup.Users[0].User)
assert.Equal(t, 1, len(backup.Users[0].Roles))
assert.Equal(t, 1, len(backup.Roles))
}
func TestRBAC_Restore(t *testing.T) {
etcdCli, _ := etcd.GetEtcdClient(
Params.EtcdCfg.UseEmbedEtcd.GetAsBool(),
Params.EtcdCfg.EtcdUseSSL.GetAsBool(),
Params.EtcdCfg.Endpoints.GetAsStrings(),
Params.EtcdCfg.EtcdTLSCert.GetValue(),
Params.EtcdCfg.EtcdTLSKey.GetValue(),
Params.EtcdCfg.EtcdTLSCACert.GetValue(),
Params.EtcdCfg.EtcdTLSMinVersion.GetValue())
rootPath := "/test/rbac"
metaKV := etcdkv.NewEtcdKV(etcdCli, rootPath)
defer metaKV.RemoveWithPrefix("")
defer metaKV.Close()
c := &Catalog{Txn: metaKV}
ctx := context.Background()
rbacMeta := &milvuspb.RBACMeta{
Users: []*milvuspb.UserInfo{
{
User: "user1",
Password: "passwd",
Roles: []*milvuspb.RoleEntity{
{
Name: "role1",
},
},
},
},
Roles: []*milvuspb.RoleEntity{
{
Name: "role1",
},
},
Grants: []*milvuspb.GrantEntity{
{
Role: &milvuspb.RoleEntity{Name: "role1"},
Object: &milvuspb.ObjectEntity{Name: "obj1"},
ObjectName: "obj_name1",
DbName: util.DefaultDBName,
Grantor: &milvuspb.GrantorEntity{
User: &milvuspb.UserEntity{Name: "user1"},
Privilege: &milvuspb.PrivilegeEntity{Name: "PrivilegeLoad"},
},
},
},
}
// test restore success
err := c.RestoreRBAC(ctx, util.DefaultTenant, rbacMeta)
assert.NoError(t, err)
// check user
users, err := c.ListCredentialsWithPasswd(ctx)
assert.NoError(t, err)
assert.Len(t, users, 1)
assert.Equal(t, users["user1"], "passwd")
// check role
roles, err := c.ListRole(ctx, util.DefaultTenant, nil, false)
assert.NoError(t, err)
assert.Len(t, roles, 1)
assert.Equal(t, "role1", roles[0].Role.Name)
// check grant
grants, err := c.ListGrant(ctx, util.DefaultTenant, &milvuspb.GrantEntity{
Role: roles[0].Role,
DbName: util.AnyWord,
})
assert.NoError(t, err)
assert.Len(t, grants, 1)
assert.Equal(t, "obj_name1", grants[0].ObjectName)
assert.Equal(t, "role1", grants[0].Role.Name)
assert.Equal(t, "user1", grants[0].Grantor.User.Name)
rbacMeta2 := &milvuspb.RBACMeta{
Users: []*milvuspb.UserInfo{
{
User: "user2",
Password: "passwd",
Roles: []*milvuspb.RoleEntity{
{
Name: "role2",
},
},
},
{
User: "user1",
Password: "passwd",
Roles: []*milvuspb.RoleEntity{
{
Name: "role2",
},
},
},
},
Roles: []*milvuspb.RoleEntity{
{
Name: "role2",
},
},
Grants: []*milvuspb.GrantEntity{
{
Role: &milvuspb.RoleEntity{Name: "role2"},
Object: &milvuspb.ObjectEntity{Name: "obj2"},
ObjectName: "obj_name2",
DbName: util.DefaultDBName,
Grantor: &milvuspb.GrantorEntity{
User: &milvuspb.UserEntity{Name: "user2"},
Privilege: &milvuspb.PrivilegeEntity{Name: "PrivilegeLoad"},
},
},
},
}
// test restore failed and roll back
err = c.RestoreRBAC(ctx, util.DefaultTenant, rbacMeta2)
assert.Error(t, err)
// check user
users, err = c.ListCredentialsWithPasswd(ctx)
assert.NoError(t, err)
assert.Len(t, users, 1)
// check role
roles, err = c.ListRole(ctx, util.DefaultTenant, nil, false)
assert.NoError(t, err)
assert.Len(t, roles, 1)
// check grant
grants, err = c.ListGrant(ctx, util.DefaultTenant, &milvuspb.GrantEntity{
Role: roles[0].Role,
DbName: util.AnyWord,
})
assert.NoError(t, err)
assert.Len(t, grants, 1)
}
func TestCatalog_AlterDatabase(t *testing.T) {
kvmock := mocks.NewSnapShotKV(t)
c := &Catalog{Snapshot: kvmock}

153
internal/metastore/mocks/mock_rootcoord_catalog.go
View File

@ -341,6 +341,61 @@ func (_c *RootCoordCatalog_AlterUserRole_Call) RunAndReturn(run func(context.Con
return _c
}
// BackupRBAC provides a mock function with given fields: ctx, tenant
func (_m *RootCoordCatalog) BackupRBAC(ctx context.Context, tenant string) (*milvuspb.RBACMeta, error) {
ret := _m.Called(ctx, tenant)
var r0 *milvuspb.RBACMeta
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, string) (*milvuspb.RBACMeta, error)); ok {
return rf(ctx, tenant)
}
if rf, ok := ret.Get(0).(func(context.Context, string) *milvuspb.RBACMeta); ok {
r0 = rf(ctx, tenant)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*milvuspb.RBACMeta)
}
}
if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
r1 = rf(ctx, tenant)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// RootCoordCatalog_BackupRBAC_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'BackupRBAC'
type RootCoordCatalog_BackupRBAC_Call struct {
*mock.Call
}
// BackupRBAC is a helper method to define mock.On call
// - ctx context.Context
// - tenant string
func (_e *RootCoordCatalog_Expecter) BackupRBAC(ctx interface{}, tenant interface{}) *RootCoordCatalog_BackupRBAC_Call {
return &RootCoordCatalog_BackupRBAC_Call{Call: _e.mock.On("BackupRBAC", ctx, tenant)}
}
func (_c *RootCoordCatalog_BackupRBAC_Call) Run(run func(ctx context.Context, tenant string)) *RootCoordCatalog_BackupRBAC_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(string))
})
return _c
}
func (_c *RootCoordCatalog_BackupRBAC_Call) Return(_a0 *milvuspb.RBACMeta, _a1 error) *RootCoordCatalog_BackupRBAC_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *RootCoordCatalog_BackupRBAC_Call) RunAndReturn(run func(context.Context, string) (*milvuspb.RBACMeta, error)) *RootCoordCatalog_BackupRBAC_Call {
_c.Call.Return(run)
return _c
}
// Close provides a mock function with given fields:
func (_m *RootCoordCatalog) Close() {
_m.Called()
@ -1327,6 +1382,60 @@ func (_c *RootCoordCatalog_ListCredentials_Call) RunAndReturn(run func(context.C
return _c
}
// ListCredentialsWithPasswd provides a mock function with given fields: ctx
func (_m *RootCoordCatalog) ListCredentialsWithPasswd(ctx context.Context) (map[string]string, error) {
ret := _m.Called(ctx)
var r0 map[string]string
var r1 error
if rf, ok := ret.Get(0).(func(context.Context) (map[string]string, error)); ok {
return rf(ctx)
}
if rf, ok := ret.Get(0).(func(context.Context) map[string]string); ok {
r0 = rf(ctx)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(map[string]string)
}
}
if rf, ok := ret.Get(1).(func(context.Context) error); ok {
r1 = rf(ctx)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// RootCoordCatalog_ListCredentialsWithPasswd_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'ListCredentialsWithPasswd'
type RootCoordCatalog_ListCredentialsWithPasswd_Call struct {
*mock.Call
}
// ListCredentialsWithPasswd is a helper method to define mock.On call
// - ctx context.Context
func (_e *RootCoordCatalog_Expecter) ListCredentialsWithPasswd(ctx interface{}) *RootCoordCatalog_ListCredentialsWithPasswd_Call {
return &RootCoordCatalog_ListCredentialsWithPasswd_Call{Call: _e.mock.On("ListCredentialsWithPasswd", ctx)}
}
func (_c *RootCoordCatalog_ListCredentialsWithPasswd_Call) Run(run func(ctx context.Context)) *RootCoordCatalog_ListCredentialsWithPasswd_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context))
})
return _c
}
func (_c *RootCoordCatalog_ListCredentialsWithPasswd_Call) Return(_a0 map[string]string, _a1 error) *RootCoordCatalog_ListCredentialsWithPasswd_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *RootCoordCatalog_ListCredentialsWithPasswd_Call) RunAndReturn(run func(context.Context) (map[string]string, error)) *RootCoordCatalog_ListCredentialsWithPasswd_Call {
_c.Call.Return(run)
return _c
}
// ListDatabases provides a mock function with given fields: ctx, ts
func (_m *RootCoordCatalog) ListDatabases(ctx context.Context, ts uint64) ([]*model.Database, error) {
ret := _m.Called(ctx, ts)
@ -1662,6 +1771,50 @@ func (_c *RootCoordCatalog_ListUserRole_Call) RunAndReturn(run func(context.Cont
return _c
}
// RestoreRBAC provides a mock function with given fields: ctx, tenant, meta
func (_m *RootCoordCatalog) RestoreRBAC(ctx context.Context, tenant string, meta *milvuspb.RBACMeta) error {
ret := _m.Called(ctx, tenant, meta)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, string, *milvuspb.RBACMeta) error); ok {
r0 = rf(ctx, tenant, meta)
} else {
r0 = ret.Error(0)
}
return r0
}
// RootCoordCatalog_RestoreRBAC_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RestoreRBAC'
type RootCoordCatalog_RestoreRBAC_Call struct {
*mock.Call
}
// RestoreRBAC is a helper method to define mock.On call
// - ctx context.Context
// - tenant string
// - meta *milvuspb.RBACMeta
func (_e *RootCoordCatalog_Expecter) RestoreRBAC(ctx interface{}, tenant interface{}, meta interface{}) *RootCoordCatalog_RestoreRBAC_Call {
return &RootCoordCatalog_RestoreRBAC_Call{Call: _e.mock.On("RestoreRBAC", ctx, tenant, meta)}
}
func (_c *RootCoordCatalog_RestoreRBAC_Call) Run(run func(ctx context.Context, tenant string, meta *milvuspb.RBACMeta)) *RootCoordCatalog_RestoreRBAC_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(string), args[2].(*milvuspb.RBACMeta))
})
return _c
}
func (_c *RootCoordCatalog_RestoreRBAC_Call) Return(_a0 error) *RootCoordCatalog_RestoreRBAC_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *RootCoordCatalog_RestoreRBAC_Call) RunAndReturn(run func(context.Context, string, *milvuspb.RBACMeta) error) *RootCoordCatalog_RestoreRBAC_Call {
_c.Call.Return(run)
return _c
}
// NewRootCoordCatalog creates a new instance of RootCoordCatalog. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations.
// The first argument is typically a *testing.T value.
func NewRootCoordCatalog(t interface {

110
internal/mocks/mock_proxy.go
View File

@ -309,6 +309,61 @@ func (_c *MockProxy_AlterIndex_Call) RunAndReturn(run func(context.Context, *mil
return _c
}
// BackupRBAC provides a mock function with given fields: _a0, _a1
func (_m *MockProxy) 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
}
// MockProxy_BackupRBAC_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'BackupRBAC'
type MockProxy_BackupRBAC_Call struct {
*mock.Call
}
// BackupRBAC is a helper method to define mock.On call
// - _a0 context.Context
// - _a1 *milvuspb.BackupRBACMetaRequest
func (_e *MockProxy_Expecter) BackupRBAC(_a0 interface{}, _a1 interface{}) *MockProxy_BackupRBAC_Call {
return &MockProxy_BackupRBAC_Call{Call: _e.mock.On("BackupRBAC", _a0, _a1)}
}
func (_c *MockProxy_BackupRBAC_Call) Run(run func(_a0 context.Context, _a1 *milvuspb.BackupRBACMetaRequest)) *MockProxy_BackupRBAC_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(*milvuspb.BackupRBACMetaRequest))
})
return _c
}
func (_c *MockProxy_BackupRBAC_Call) Return(_a0 *milvuspb.BackupRBACMetaResponse, _a1 error) *MockProxy_BackupRBAC_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockProxy_BackupRBAC_Call) RunAndReturn(run func(context.Context, *milvuspb.BackupRBACMetaRequest) (*milvuspb.BackupRBACMetaResponse, error)) *MockProxy_BackupRBAC_Call {
_c.Call.Return(run)
return _c
}
// CalcDistance provides a mock function with given fields: _a0, _a1
func (_m *MockProxy) CalcDistance(_a0 context.Context, _a1 *milvuspb.CalcDistanceRequest) (*milvuspb.CalcDistanceResults, error) {
ret := _m.Called(_a0, _a1)
@ -4940,6 +4995,61 @@ func (_c *MockProxy_ReplicateMessage_Call) RunAndReturn(run func(context.Context
return _c
}
// RestoreRBAC provides a mock function with given fields: _a0, _a1
func (_m *MockProxy) 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
}
// MockProxy_RestoreRBAC_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RestoreRBAC'
type MockProxy_RestoreRBAC_Call struct {
*mock.Call
}
// RestoreRBAC is a helper method to define mock.On call
// - _a0 context.Context
// - _a1 *milvuspb.RestoreRBACMetaRequest
func (_e *MockProxy_Expecter) RestoreRBAC(_a0 interface{}, _a1 interface{}) *MockProxy_RestoreRBAC_Call {
return &MockProxy_RestoreRBAC_Call{Call: _e.mock.On("RestoreRBAC", _a0, _a1)}
}
func (_c *MockProxy_RestoreRBAC_Call) Run(run func(_a0 context.Context, _a1 *milvuspb.RestoreRBACMetaRequest)) *MockProxy_RestoreRBAC_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(*milvuspb.RestoreRBACMetaRequest))
})
return _c
}
func (_c *MockProxy_RestoreRBAC_Call) Return(_a0 *commonpb.Status, _a1 error) *MockProxy_RestoreRBAC_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockProxy_RestoreRBAC_Call) RunAndReturn(run func(context.Context, *milvuspb.RestoreRBACMetaRequest) (*commonpb.Status, error)) *MockProxy_RestoreRBAC_Call {
_c.Call.Return(run)
return _c
}
// Search provides a mock function with given fields: _a0, _a1
func (_m *MockProxy) Search(_a0 context.Context, _a1 *milvuspb.SearchRequest) (*milvuspb.SearchResults, error) {
ret := _m.Called(_a0, _a1)

110
internal/mocks/mock_rootcoord.go
View File

@ -311,6 +311,61 @@ func (_c *RootCoord_AlterDatabase_Call) RunAndReturn(run func(context.Context, *
return _c
}
// BackupRBAC provides a mock function with given fields: _a0, _a1
func (_m *RootCoord) 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
}
// RootCoord_BackupRBAC_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'BackupRBAC'
type RootCoord_BackupRBAC_Call struct {
*mock.Call
}
// BackupRBAC is a helper method to define mock.On call
// - _a0 context.Context
// - _a1 *milvuspb.BackupRBACMetaRequest
func (_e *RootCoord_Expecter) BackupRBAC(_a0 interface{}, _a1 interface{}) *RootCoord_BackupRBAC_Call {
return &RootCoord_BackupRBAC_Call{Call: _e.mock.On("BackupRBAC", _a0, _a1)}
}
func (_c *RootCoord_BackupRBAC_Call) Run(run func(_a0 context.Context, _a1 *milvuspb.BackupRBACMetaRequest)) *RootCoord_BackupRBAC_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(*milvuspb.BackupRBACMetaRequest))
})
return _c
}
func (_c *RootCoord_BackupRBAC_Call) Return(_a0 *milvuspb.BackupRBACMetaResponse, _a1 error) *RootCoord_BackupRBAC_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *RootCoord_BackupRBAC_Call) RunAndReturn(run func(context.Context, *milvuspb.BackupRBACMetaRequest) (*milvuspb.BackupRBACMetaResponse, error)) *RootCoord_BackupRBAC_Call {
_c.Call.Return(run)
return _c
}
// CheckHealth provides a mock function with given fields: _a0, _a1
func (_m *RootCoord) CheckHealth(_a0 context.Context, _a1 *milvuspb.CheckHealthRequest) (*milvuspb.CheckHealthResponse, error) {
ret := _m.Called(_a0, _a1)
@ -2208,6 +2263,61 @@ func (_c *RootCoord_RenameCollection_Call) RunAndReturn(run func(context.Context
return _c
}
// RestoreRBAC provides a mock function with given fields: _a0, _a1
func (_m *RootCoord) 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
}
// RootCoord_RestoreRBAC_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RestoreRBAC'
type RootCoord_RestoreRBAC_Call struct {
*mock.Call
}
// RestoreRBAC is a helper method to define mock.On call
// - _a0 context.Context
// - _a1 *milvuspb.RestoreRBACMetaRequest
func (_e *RootCoord_Expecter) RestoreRBAC(_a0 interface{}, _a1 interface{}) *RootCoord_RestoreRBAC_Call {
return &RootCoord_RestoreRBAC_Call{Call: _e.mock.On("RestoreRBAC", _a0, _a1)}
}
func (_c *RootCoord_RestoreRBAC_Call) Run(run func(_a0 context.Context, _a1 *milvuspb.RestoreRBACMetaRequest)) *RootCoord_RestoreRBAC_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(*milvuspb.RestoreRBACMetaRequest))
})
return _c
}
func (_c *RootCoord_RestoreRBAC_Call) Return(_a0 *commonpb.Status, _a1 error) *RootCoord_RestoreRBAC_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *RootCoord_RestoreRBAC_Call) RunAndReturn(run func(context.Context, *milvuspb.RestoreRBACMetaRequest) (*commonpb.Status, error)) *RootCoord_RestoreRBAC_Call {
_c.Call.Return(run)
return _c
}
// SelectGrant provides a mock function with given fields: _a0, _a1
func (_m *RootCoord) SelectGrant(_a0 context.Context, _a1 *milvuspb.SelectGrantRequest) (*milvuspb.SelectGrantResponse, error) {
ret := _m.Called(_a0, _a1)

140
internal/mocks/mock_rootcoord_client.go
View File

@ -383,6 +383,76 @@ func (_c *MockRootCoordClient_AlterDatabase_Call) RunAndReturn(run func(context.
return _c
}
// BackupRBAC provides a mock function with given fields: ctx, in, opts
func (_m *MockRootCoordClient) BackupRBAC(ctx context.Context, in *milvuspb.BackupRBACMetaRequest, opts ...grpc.CallOption) (*milvuspb.BackupRBACMetaResponse, error) {
_va := make([]interface{}, len(opts))
for _i := range opts {
_va[_i] = opts[_i]
}
var _ca []interface{}
_ca = append(_ca, ctx, in)
_ca = append(_ca, _va...)
ret := _m.Called(_ca...)
var r0 *milvuspb.BackupRBACMetaResponse
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.BackupRBACMetaRequest, ...grpc.CallOption) (*milvuspb.BackupRBACMetaResponse, error)); ok {
return rf(ctx, in, opts...)
}
if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.BackupRBACMetaRequest, ...grpc.CallOption) *milvuspb.BackupRBACMetaResponse); ok {
r0 = rf(ctx, in, opts...)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*milvuspb.BackupRBACMetaResponse)
}
}
if rf, ok := ret.Get(1).(func(context.Context, *milvuspb.BackupRBACMetaRequest, ...grpc.CallOption) error); ok {
r1 = rf(ctx, in, opts...)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockRootCoordClient_BackupRBAC_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'BackupRBAC'
type MockRootCoordClient_BackupRBAC_Call struct {
*mock.Call
}
// BackupRBAC is a helper method to define mock.On call
// - ctx context.Context
// - in *milvuspb.BackupRBACMetaRequest
// - opts ...grpc.CallOption
func (_e *MockRootCoordClient_Expecter) BackupRBAC(ctx interface{}, in interface{}, opts ...interface{}) *MockRootCoordClient_BackupRBAC_Call {
return &MockRootCoordClient_BackupRBAC_Call{Call: _e.mock.On("BackupRBAC",
append([]interface{}{ctx, in}, opts...)...)}
}
func (_c *MockRootCoordClient_BackupRBAC_Call) Run(run func(ctx context.Context, in *milvuspb.BackupRBACMetaRequest, opts ...grpc.CallOption)) *MockRootCoordClient_BackupRBAC_Call {
_c.Call.Run(func(args mock.Arguments) {
variadicArgs := make([]grpc.CallOption, len(args)-2)
for i, a := range args[2:] {
if a != nil {
variadicArgs[i] = a.(grpc.CallOption)
}
}
run(args[0].(context.Context), args[1].(*milvuspb.BackupRBACMetaRequest), variadicArgs...)
})
return _c
}
func (_c *MockRootCoordClient_BackupRBAC_Call) Return(_a0 *milvuspb.BackupRBACMetaResponse, _a1 error) *MockRootCoordClient_BackupRBAC_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockRootCoordClient_BackupRBAC_Call) RunAndReturn(run func(context.Context, *milvuspb.BackupRBACMetaRequest, ...grpc.CallOption) (*milvuspb.BackupRBACMetaResponse, error)) *MockRootCoordClient_BackupRBAC_Call {
_c.Call.Return(run)
return _c
}
// CheckHealth provides a mock function with given fields: ctx, in, opts
func (_m *MockRootCoordClient) CheckHealth(ctx context.Context, in *milvuspb.CheckHealthRequest, opts ...grpc.CallOption) (*milvuspb.CheckHealthResponse, error) {
_va := make([]interface{}, len(opts))
@ -2734,6 +2804,76 @@ func (_c *MockRootCoordClient_RenameCollection_Call) RunAndReturn(run func(conte
return _c
}
// RestoreRBAC provides a mock function with given fields: ctx, in, opts
func (_m *MockRootCoordClient) RestoreRBAC(ctx context.Context, in *milvuspb.RestoreRBACMetaRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
_va := make([]interface{}, len(opts))
for _i := range opts {
_va[_i] = opts[_i]
}
var _ca []interface{}
_ca = append(_ca, ctx, in)
_ca = append(_ca, _va...)
ret := _m.Called(_ca...)
var r0 *commonpb.Status
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.RestoreRBACMetaRequest, ...grpc.CallOption) (*commonpb.Status, error)); ok {
return rf(ctx, in, opts...)
}
if rf, ok := ret.Get(0).(func(context.Context, *milvuspb.RestoreRBACMetaRequest, ...grpc.CallOption) *commonpb.Status); ok {
r0 = rf(ctx, in, opts...)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*commonpb.Status)
}
}
if rf, ok := ret.Get(1).(func(context.Context, *milvuspb.RestoreRBACMetaRequest, ...grpc.CallOption) error); ok {
r1 = rf(ctx, in, opts...)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// MockRootCoordClient_RestoreRBAC_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RestoreRBAC'
type MockRootCoordClient_RestoreRBAC_Call struct {
*mock.Call
}
// RestoreRBAC is a helper method to define mock.On call
// - ctx context.Context
// - in *milvuspb.RestoreRBACMetaRequest
// - opts ...grpc.CallOption
func (_e *MockRootCoordClient_Expecter) RestoreRBAC(ctx interface{}, in interface{}, opts ...interface{}) *MockRootCoordClient_RestoreRBAC_Call {
return &MockRootCoordClient_RestoreRBAC_Call{Call: _e.mock.On("RestoreRBAC",
append([]interface{}{ctx, in}, opts...)...)}
}
func (_c *MockRootCoordClient_RestoreRBAC_Call) Run(run func(ctx context.Context, in *milvuspb.RestoreRBACMetaRequest, opts ...grpc.CallOption)) *MockRootCoordClient_RestoreRBAC_Call {
_c.Call.Run(func(args mock.Arguments) {
variadicArgs := make([]grpc.CallOption, len(args)-2)
for i, a := range args[2:] {
if a != nil {
variadicArgs[i] = a.(grpc.CallOption)
}
}
run(args[0].(context.Context), args[1].(*milvuspb.RestoreRBACMetaRequest), variadicArgs...)
})
return _c
}
func (_c *MockRootCoordClient_RestoreRBAC_Call) Return(_a0 *commonpb.Status, _a1 error) *MockRootCoordClient_RestoreRBAC_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *MockRootCoordClient_RestoreRBAC_Call) RunAndReturn(run func(context.Context, *milvuspb.RestoreRBACMetaRequest, ...grpc.CallOption) (*commonpb.Status, error)) *MockRootCoordClient_RestoreRBAC_Call {
_c.Call.Return(run)
return _c
}
// SelectGrant provides a mock function with given fields: ctx, in, opts
func (_m *MockRootCoordClient) SelectGrant(ctx context.Context, in *milvuspb.SelectGrantRequest, opts ...grpc.CallOption) (*milvuspb.SelectGrantResponse, error) {
_va := make([]interface{}, len(opts))

2
internal/proto/root_coord.proto
View File

@ -124,6 +124,8 @@ service RootCoord {
rpc OperatePrivilege(milvus.OperatePrivilegeRequest) returns (common.Status) {}
rpc SelectGrant(milvus.SelectGrantRequest) returns (milvus.SelectGrantResponse) {}
rpc ListPolicy(internal.ListPolicyRequest) returns (internal.ListPolicyResponse) {}
rpc BackupRBAC(milvus.BackupRBACMetaRequest) returns (milvus.BackupRBACMetaResponse){}
rpc RestoreRBAC(milvus.RestoreRBACMetaRequest) returns (common.Status){}
rpc CheckHealth(milvus.CheckHealthRequest) returns (milvus.CheckHealthResponse) {}

40
internal/proxy/impl.go
View File

@ -5325,6 +5325,46 @@ func (node *Proxy) SelectGrant(ctx context.Context, req *milvuspb.SelectGrantReq
return result, nil
}
func (node *Proxy) BackupRBAC(ctx context.Context, req *milvuspb.BackupRBACMetaRequest) (*milvuspb.BackupRBACMetaResponse, error) {
ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-BackupRBAC")
defer sp.End()
log := log.Ctx(ctx)
log.Debug("BackupRBAC", zap.Any("req", req))
if err := merr.CheckHealthy(node.GetStateCode()); err != nil {
return &milvuspb.BackupRBACMetaResponse{Status: merr.Status(err)}, nil
}
result, err := node.rootCoord.BackupRBAC(ctx, req)
if err != nil {
log.Warn("fail to backup rbac", zap.Error(err))
return &milvuspb.BackupRBACMetaResponse{
Status: merr.Status(err),
}, nil
}
return result, nil
}
func (node *Proxy) RestoreRBAC(ctx context.Context, req *milvuspb.RestoreRBACMetaRequest) (*commonpb.Status, error) {
ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-RestoreRBAC")
defer sp.End()
log := log.Ctx(ctx)
log.Debug("RestoreRBAC", zap.Any("req", req))
if err := merr.CheckHealthy(node.GetStateCode()); err != nil {
return merr.Status(err), nil
}
result, err := node.rootCoord.RestoreRBAC(ctx, req)
if err != nil {
log.Warn("fail to restore rbac", zap.Error(err))
return merr.Status(err), nil
}
return result, nil
}
func (node *Proxy) RefreshPolicyInfoCache(ctx context.Context, req *proxypb.RefreshPolicyInfoCacheRequest) (*commonpb.Status, error) {
ctx, sp := otel.Tracer(typeutil.ProxyRole).Start(ctx, "Proxy-RefreshPolicyInfoCache")
defer sp.End()

8
internal/proxy/rootcoord_mock_test.go
View File

@ -1126,6 +1126,14 @@ func (coord *RootCoordMock) AlterDatabase(ctx context.Context, in *rootcoordpb.A
return &commonpb.Status{}, nil
}
func (coord *RootCoordMock) BackupRBAC(ctx context.Context, in *milvuspb.BackupRBACMetaRequest, opts ...grpc.CallOption) (*milvuspb.BackupRBACMetaResponse, error) {
return &milvuspb.BackupRBACMetaResponse{}, nil
}
func (coord *RootCoordMock) RestoreRBAC(ctx context.Context, in *milvuspb.RestoreRBACMetaRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
return &commonpb.Status{}, nil
}
type DescribeCollectionFunc func(ctx context.Context, request *milvuspb.DescribeCollectionRequest, opts ...grpc.CallOption) (*milvuspb.DescribeCollectionResponse, error)
type ShowPartitionsFunc func(ctx context.Context, request *milvuspb.ShowPartitionsRequest, opts ...grpc.CallOption) (*milvuspb.ShowPartitionsResponse, error)

16
internal/rootcoord/meta_table.go
View File

@ -97,6 +97,8 @@ type IMetaTable interface {
DropGrant(tenant string, role *milvuspb.RoleEntity) error
ListPolicy(tenant string) ([]string, error)
ListUserRole(tenant string) ([]string, error)
BackupRBAC(ctx context.Context, tenant string) (*milvuspb.RBACMeta, error)
RestoreRBAC(ctx context.Context, tenant string, meta *milvuspb.RBACMeta) error
}
// MetaTable is a persistent meta set of all databases, collections and partitions.
@ -1435,3 +1437,17 @@ func (mt *MetaTable) ListUserRole(tenant string) ([]string, error) {
return mt.catalog.ListUserRole(mt.ctx, tenant)
}
func (mt *MetaTable) BackupRBAC(ctx context.Context, tenant string) (*milvuspb.RBACMeta, error) {
mt.permissionLock.RLock()
defer mt.permissionLock.RUnlock()
return mt.catalog.BackupRBAC(mt.ctx, tenant)
}
func (mt *MetaTable) RestoreRBAC(ctx context.Context, tenant string, meta *milvuspb.RBACMeta) error {
mt.permissionLock.Lock()
defer mt.permissionLock.Unlock()
return mt.catalog.RestoreRBAC(mt.ctx, tenant, meta)
}

41
internal/rootcoord/meta_table_test.go
View File

@ -2000,3 +2000,44 @@ func TestMetaTable_DropDatabase(t *testing.T) {
assert.False(t, mt.aliases.exist("not_commit"))
})
}
func TestMetaTable_BackupRBAC(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.EXPECT().BackupRBAC(mock.Anything, mock.Anything).Return(&milvuspb.RBACMeta{}, nil)
mt := &MetaTable{
dbName2Meta: map[string]*model.Database{
"not_commit": model.NewDatabase(1, "not_commit", pb.DatabaseState_DatabaseCreated, nil),
},
names: newNameDb(),
aliases: newNameDb(),
catalog: catalog,
}
_, err := mt.BackupRBAC(context.TODO(), util.DefaultTenant)
assert.NoError(t, err)
catalog.ExpectedCalls = nil
catalog.EXPECT().BackupRBAC(mock.Anything, mock.Anything).Return(nil, errors.New("error mock BackupRBAC"))
_, err = mt.BackupRBAC(context.TODO(), util.DefaultTenant)
assert.Error(t, err)
}
func TestMetaTable_RestoreRBAC(t *testing.T) {
catalog := mocks.NewRootCoordCatalog(t)
catalog.EXPECT().RestoreRBAC(mock.Anything, mock.Anything, mock.Anything).Return(nil)
mt := &MetaTable{
dbName2Meta: map[string]*model.Database{
"not_commit": model.NewDatabase(1, "not_commit", pb.DatabaseState_DatabaseCreated, nil),
},
names: newNameDb(),
aliases: newNameDb(),
catalog: catalog,
}
err := mt.RestoreRBAC(context.TODO(), util.DefaultTenant, &milvuspb.RBACMeta{})
assert.NoError(t, err)
catalog.ExpectedCalls = nil
catalog.EXPECT().RestoreRBAC(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("error mock RestoreRBAC"))
err = mt.RestoreRBAC(context.TODO(), util.DefaultTenant, &milvuspb.RBACMeta{})
assert.Error(t, err)
}

99
internal/rootcoord/mocks/meta_table.go
View File

@ -336,6 +336,61 @@ func (_c *IMetaTable_AlterDatabase_Call) RunAndReturn(run func(context.Context,
return _c
}
// BackupRBAC provides a mock function with given fields: ctx, tenant
func (_m *IMetaTable) BackupRBAC(ctx context.Context, tenant string) (*milvuspb.RBACMeta, error) {
ret := _m.Called(ctx, tenant)
var r0 *milvuspb.RBACMeta
var r1 error
if rf, ok := ret.Get(0).(func(context.Context, string) (*milvuspb.RBACMeta, error)); ok {
return rf(ctx, tenant)
}
if rf, ok := ret.Get(0).(func(context.Context, string) *milvuspb.RBACMeta); ok {
r0 = rf(ctx, tenant)
} else {
if ret.Get(0) != nil {
r0 = ret.Get(0).(*milvuspb.RBACMeta)
}
}
if rf, ok := ret.Get(1).(func(context.Context, string) error); ok {
r1 = rf(ctx, tenant)
} else {
r1 = ret.Error(1)
}
return r0, r1
}
// IMetaTable_BackupRBAC_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'BackupRBAC'
type IMetaTable_BackupRBAC_Call struct {
*mock.Call
}
// BackupRBAC is a helper method to define mock.On call
// - ctx context.Context
// - tenant string
func (_e *IMetaTable_Expecter) BackupRBAC(ctx interface{}, tenant interface{}) *IMetaTable_BackupRBAC_Call {
return &IMetaTable_BackupRBAC_Call{Call: _e.mock.On("BackupRBAC", ctx, tenant)}
}
func (_c *IMetaTable_BackupRBAC_Call) Run(run func(ctx context.Context, tenant string)) *IMetaTable_BackupRBAC_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(string))
})
return _c
}
func (_c *IMetaTable_BackupRBAC_Call) Return(_a0 *milvuspb.RBACMeta, _a1 error) *IMetaTable_BackupRBAC_Call {
_c.Call.Return(_a0, _a1)
return _c
}
func (_c *IMetaTable_BackupRBAC_Call) RunAndReturn(run func(context.Context, string) (*milvuspb.RBACMeta, error)) *IMetaTable_BackupRBAC_Call {
_c.Call.Return(run)
return _c
}
// ChangeCollectionState provides a mock function with given fields: ctx, collectionID, state, ts
func (_m *IMetaTable) ChangeCollectionState(ctx context.Context, collectionID int64, state etcdpb.CollectionState, ts uint64) error {
ret := _m.Called(ctx, collectionID, state, ts)
@ -1986,6 +2041,50 @@ func (_c *IMetaTable_RenameCollection_Call) RunAndReturn(run func(context.Contex
return _c
}
// RestoreRBAC provides a mock function with given fields: ctx, tenant, meta
func (_m *IMetaTable) RestoreRBAC(ctx context.Context, tenant string, meta *milvuspb.RBACMeta) error {
ret := _m.Called(ctx, tenant, meta)
var r0 error
if rf, ok := ret.Get(0).(func(context.Context, string, *milvuspb.RBACMeta) error); ok {
r0 = rf(ctx, tenant, meta)
} else {
r0 = ret.Error(0)
}
return r0
}
// IMetaTable_RestoreRBAC_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'RestoreRBAC'
type IMetaTable_RestoreRBAC_Call struct {
*mock.Call
}
// RestoreRBAC is a helper method to define mock.On call
// - ctx context.Context
// - tenant string
// - meta *milvuspb.RBACMeta
func (_e *IMetaTable_Expecter) RestoreRBAC(ctx interface{}, tenant interface{}, meta interface{}) *IMetaTable_RestoreRBAC_Call {
return &IMetaTable_RestoreRBAC_Call{Call: _e.mock.On("RestoreRBAC", ctx, tenant, meta)}
}
func (_c *IMetaTable_RestoreRBAC_Call) Run(run func(ctx context.Context, tenant string, meta *milvuspb.RBACMeta)) *IMetaTable_RestoreRBAC_Call {
_c.Call.Run(func(args mock.Arguments) {
run(args[0].(context.Context), args[1].(string), args[2].(*milvuspb.RBACMeta))
})
return _c
}
func (_c *IMetaTable_RestoreRBAC_Call) Return(_a0 error) *IMetaTable_RestoreRBAC_Call {
_c.Call.Return(_a0)
return _c
}
func (_c *IMetaTable_RestoreRBAC_Call) RunAndReturn(run func(context.Context, string, *milvuspb.RBACMeta) error) *IMetaTable_RestoreRBAC_Call {
_c.Call.Return(run)
return _c
}
// SelectGrant provides a mock function with given fields: tenant, entity
func (_m *IMetaTable) SelectGrant(tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error) {
ret := _m.Called(tenant, entity)

51
internal/rootcoord/root_coord.go
View File

@ -2706,6 +2706,57 @@ func (c *Core) ListPolicy(ctx context.Context, in *internalpb.ListPolicyRequest)
}, nil
}
func (c *Core) BackupRBAC(ctx context.Context, in *milvuspb.BackupRBACMetaRequest) (*milvuspb.BackupRBACMetaResponse, error) {
method := "BackupRBAC"
metrics.RootCoordDDLReqCounter.WithLabelValues(method, metrics.TotalLabel).Inc()
tr := timerecord.NewTimeRecorder(method)
ctxLog := log.Ctx(ctx).With(zap.String("role", typeutil.RootCoordRole), zap.Any("in", in))
ctxLog.Debug(method)
if err := merr.CheckHealthy(c.GetStateCode()); err != nil {
return &milvuspb.BackupRBACMetaResponse{
Status: merr.Status(err),
}, nil
}
rbacMeta, err := c.meta.BackupRBAC(ctx, util.DefaultTenant)
if err != nil {
return &milvuspb.BackupRBACMetaResponse{
Status: merr.Status(err),
}, nil
}
ctxLog.Debug(method + " success")
metrics.RootCoordDDLReqCounter.WithLabelValues(method, metrics.SuccessLabel).Inc()
metrics.RootCoordDDLReqLatency.WithLabelValues(method).Observe(float64(tr.ElapseSpan().Milliseconds()))
return &milvuspb.BackupRBACMetaResponse{
Status: merr.Success(),
RBACMeta: rbacMeta,
}, nil
}
func (c *Core) RestoreRBAC(ctx context.Context, in *milvuspb.RestoreRBACMetaRequest) (*commonpb.Status, error) {
method := "RestoreRBAC"
metrics.RootCoordDDLReqCounter.WithLabelValues(method, metrics.TotalLabel).Inc()
tr := timerecord.NewTimeRecorder(method)
ctxLog := log.Ctx(ctx).With(zap.String("role", typeutil.RootCoordRole), zap.Any("in", in))
ctxLog.Debug(method)
if err := merr.CheckHealthy(c.GetStateCode()); err != nil {
return merr.Status(err), nil
}
if err := c.meta.RestoreRBAC(ctx, util.DefaultTenant, in.RBACMeta); err != nil {
return merr.Status(err), nil
}
ctxLog.Debug(method + " success")
metrics.RootCoordDDLReqCounter.WithLabelValues(method, metrics.SuccessLabel).Inc()
metrics.RootCoordDDLReqLatency.WithLabelValues(method).Observe(float64(tr.ElapseSpan().Milliseconds()))
return merr.Success(), nil
}
func (c *Core) RenameCollection(ctx context.Context, req *milvuspb.RenameCollectionRequest) (*commonpb.Status, error) {
if err := merr.CheckHealthy(c.GetStateCode()); err != nil {
return merr.Status(err), nil

32
internal/rootcoord/root_coord_test.go
View File

@ -1971,6 +1971,38 @@ func TestCore_InitRBAC(t *testing.T) {
})
}
func TestCore_BackupRBAC(t *testing.T) {
meta := mockrootcoord.NewIMetaTable(t)
c := newTestCore(withHealthyCode(), withMeta(meta))
meta.EXPECT().BackupRBAC(mock.Anything, mock.Anything).Return(&milvuspb.RBACMeta{}, nil)
resp, err := c.BackupRBAC(context.Background(), &milvuspb.BackupRBACMetaRequest{})
assert.NoError(t, err)
assert.True(t, merr.Ok(resp.GetStatus()))
meta.ExpectedCalls = nil
meta.EXPECT().BackupRBAC(mock.Anything, mock.Anything).Return(nil, errors.New("mock error"))
resp, err = c.BackupRBAC(context.Background(), &milvuspb.BackupRBACMetaRequest{})
assert.NoError(t, err)
assert.False(t, merr.Ok(resp.GetStatus()))
}
func TestCore_RestoreRBAC(t *testing.T) {
meta := mockrootcoord.NewIMetaTable(t)
c := newTestCore(withHealthyCode(), withMeta(meta))
meta.EXPECT().RestoreRBAC(mock.Anything, mock.Anything, mock.Anything).Return(nil)
resp, err := c.RestoreRBAC(context.Background(), &milvuspb.RestoreRBACMetaRequest{})
assert.NoError(t, err)
assert.True(t, merr.Ok(resp))
meta.ExpectedCalls = nil
meta.EXPECT().RestoreRBAC(mock.Anything, mock.Anything, mock.Anything).Return(errors.New("mock error"))
resp, err = c.RestoreRBAC(context.Background(), &milvuspb.RestoreRBACMetaRequest{})
assert.NoError(t, err)
assert.False(t, merr.Ok(resp))
}
type RootCoordSuite struct {
suite.Suite
}

8
internal/util/mock/grpc_rootcoord_client.go
View File

@ -266,6 +266,14 @@ func (m *GrpcRootCoordClient) AlterDatabase(ctx context.Context, in *rootcoordpb
return &commonpb.Status{}, m.Err
}
func (m *GrpcRootCoordClient) BackupRBAC(ctx context.Context, in *milvuspb.BackupRBACMetaRequest, opts ...grpc.CallOption) (*milvuspb.BackupRBACMetaResponse, error) {
return &milvuspb.BackupRBACMetaResponse{}, m.Err
}
func (m *GrpcRootCoordClient) RestoreRBAC(ctx context.Context, in *milvuspb.RestoreRBACMetaRequest, opts ...grpc.CallOption) (*commonpb.Status, error) {
return &commonpb.Status{}, m.Err
}
func (m *GrpcRootCoordClient) Close() error {
return nil
}

6
pkg/go.mod
View File

@ -6,13 +6,15 @@ require (
github.com/apache/pulsar-client-go v0.6.1-0.20210728062540-29414db801a7
github.com/benesch/cgosymbolizer v0.0.0-20190515212042-bec6fe6e597b
github.com/blang/semver/v4 v4.0.0
github.com/cenkalti/backoff/v4 v4.2.1
github.com/cockroachdb/errors v1.9.1
github.com/confluentinc/confluent-kafka-go v1.9.1
github.com/containerd/cgroups/v3 v3.0.3
github.com/expr-lang/expr v1.15.7
github.com/golang/protobuf v1.5.4
github.com/grpc-ecosystem/go-grpc-middleware v1.3.0
github.com/klauspost/compress v1.17.7
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240717062137-3ffb1db01632
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240815113856-e2789dca8b59
github.com/nats-io/nats-server/v2 v2.10.12
github.com/nats-io/nats.go v1.34.1
github.com/panjf2000/ants/v2 v2.7.2
@ -65,7 +67,6 @@ require (
github.com/ardielle/ardielle-go v1.5.2 // indirect
github.com/benbjohnson/clock v1.1.0 // indirect
github.com/beorn7/perks v1.0.1 // indirect
github.com/cenkalti/backoff/v4 v4.2.1 // indirect
github.com/cespare/xxhash/v2 v2.3.0 // indirect
github.com/cilium/ebpf v0.11.0 // indirect
github.com/cockroachdb/logtags v0.0.0-20211118104740-dabe8e521a4f // indirect
@ -92,7 +93,6 @@ require (
github.com/godbus/dbus/v5 v5.0.4 // indirect
github.com/gogo/protobuf v1.3.2 // indirect
github.com/golang-jwt/jwt v3.2.2+incompatible // indirect
github.com/golang/protobuf v1.5.4 // indirect
github.com/golang/snappy v0.0.4 // indirect
github.com/google/btree v1.1.2 // indirect
github.com/google/uuid v1.6.0 // indirect

4
pkg/go.sum
View File

@ -494,8 +494,8 @@ github.com/milvus-io/cgosymbolizer v0.0.0-20240722103217-b7dee0e50119 h1:9VXijWu
github.com/milvus-io/cgosymbolizer v0.0.0-20240722103217-b7dee0e50119/go.mod h1:DvXTE/K/RtHehxU8/GtDs4vFtfw64jJ3PaCnFri8CRg=
github.com/milvus-io/gorocksdb v0.0.0-20220624081344-8c5f4212846b h1:TfeY0NxYxZzUfIfYe5qYDBzt4ZYRqzUjTR6CvUzjat8=
github.com/milvus-io/gorocksdb v0.0.0-20220624081344-8c5f4212846b/go.mod h1:iwW+9cWfIzzDseEBCCeDSN5SD16Tidvy8cwQ7ZY8Qj4=
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240717062137-3ffb1db01632 h1:CXig0DNtUsCLzchCFe3PR2KgOdobbz9gK2nSV7195PM=
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240717062137-3ffb1db01632/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs=
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240815113856-e2789dca8b59 h1:mKekr0GmCKMpIQh9OJ67TlKVKxDt08600ltARc/JUXY=
github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240815113856-e2789dca8b59/go.mod h1:/6UT4zZl6awVeXLeE7UGDWZvXj3IWkRsh3mqsn0DiAs=
github.com/milvus-io/pulsar-client-go v0.6.10 h1:eqpJjU+/QX0iIhEo3nhOqMNXL+TyInAs1IAHZCrCM/A=
github.com/milvus-io/pulsar-client-go v0.6.10/go.mod h1:lQqCkgwDF8YFYjKA+zOheTk1tev2B+bKj5j7+nm8M1w=
github.com/minio/highwayhash v1.0.2 h1:Aak5U0nElisjDCfPSG79Tgzkn2gl66NxOMspRrKnA/g=

2
pkg/util/constant.go
View File

@ -123,6 +123,8 @@ var (
MetaStore2API(commonpb.ObjectPrivilege_PrivilegeDropOwnership.String()),
MetaStore2API(commonpb.ObjectPrivilege_PrivilegeSelectOwnership.String()),
MetaStore2API(commonpb.ObjectPrivilege_PrivilegeManageOwnership.String()),
MetaStore2API(commonpb.ObjectPrivilege_PrivilegeBackupRBAC.String()),
MetaStore2API(commonpb.ObjectPrivilege_PrivilegeRestoreRBAC.String()),
MetaStore2API(commonpb.ObjectPrivilege_PrivilegeCreateResourceGroup.String()),
MetaStore2API(commonpb.ObjectPrivilege_PrivilegeUpdateResourceGroups.String()),

166
tests/integration/rbac/rbac_backup_test.go Normal file
View File

@ -0,0 +1,166 @@
// 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 rbac
import (
"context"
"strings"
"testing"
"github.com/stretchr/testify/suite"
"google.golang.org/grpc/metadata"
"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/pkg/util"
"github.com/milvus-io/milvus/pkg/util/crypto"
"github.com/milvus-io/milvus/pkg/util/merr"
"github.com/milvus-io/milvus/pkg/util/paramtable"
"github.com/milvus-io/milvus/tests/integration"
)
const (
dim = 128
dbName = ""
collectionName = "test_load_collection"
)
type RBACBackupTestSuite struct {
integration.MiniClusterSuite
}
func (s *RBACBackupTestSuite) SetupSuite() {
paramtable.Init()
paramtable.Get().Save(paramtable.Get().QueryCoordCfg.BalanceCheckInterval.Key, "1000")
paramtable.Get().Save(paramtable.Get().QueryNodeCfg.GracefulStopTimeout.Key, "1")
paramtable.Get().Save(paramtable.Get().CommonCfg.AuthorizationEnabled.Key, "true")
s.Require().NoError(s.SetupEmbedEtcd())
}
func GetContext(ctx context.Context, originValue string) context.Context {
authKey := strings.ToLower(util.HeaderAuthorize)
authValue := crypto.Base64Encode(originValue)
contextMap := map[string]string{
authKey: authValue,
}
md := metadata.New(contextMap)
return metadata.NewIncomingContext(ctx, md)
}
func (s *RBACBackupTestSuite) TestBackup() {
ctx := GetContext(context.Background(), "root:123456")
// test empty rbac content
resp, err := s.Cluster.Proxy.BackupRBAC(ctx, &milvuspb.BackupRBACMetaRequest{})
s.NoError(err)
s.True(merr.Ok(resp.GetStatus()))
s.Equal("", resp.GetRBACMeta().String())
// generate some rbac content
roleName := "test_role"
resp1, err := s.Cluster.Proxy.CreateRole(ctx, &milvuspb.CreateRoleRequest{
Entity: &milvuspb.RoleEntity{
Name: roleName,
},
})
s.NoError(err)
s.True(merr.Ok(resp1))
resp2, err := s.Cluster.Proxy.OperatePrivilege(ctx, &milvuspb.OperatePrivilegeRequest{
Type: milvuspb.OperatePrivilegeType_Grant,
Entity: &milvuspb.GrantEntity{
Role: &milvuspb.RoleEntity{Name: roleName},
Object: &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Collection.String()},
ObjectName: util.AnyWord,
DbName: util.AnyWord,
Grantor: &milvuspb.GrantorEntity{
User: &milvuspb.UserEntity{Name: util.UserRoot},
Privilege: &milvuspb.PrivilegeEntity{Name: util.AnyWord},
},
},
})
s.NoError(err)
s.True(merr.Ok(resp2))
s.Equal("", resp2.GetReason())
userName := "test_user"
passwd := "test_passwd"
resp3, err := s.Cluster.Proxy.CreateCredential(ctx, &milvuspb.CreateCredentialRequest{
Username: userName,
Password: crypto.Base64Encode(passwd),
})
s.NoError(err)
s.True(merr.Ok(resp3))
resp4, err := s.Cluster.Proxy.OperateUserRole(ctx, &milvuspb.OperateUserRoleRequest{
Username: userName,
RoleName: roleName,
})
s.NoError(err)
s.True(merr.Ok(resp4))
// test back up rbac
resp5, err := s.Cluster.Proxy.BackupRBAC(ctx, &milvuspb.BackupRBACMetaRequest{})
s.NoError(err)
s.True(merr.Ok(resp5.GetStatus()))
// test restore, expect to failed due to role/user already exist
resp6, err := s.Cluster.Proxy.RestoreRBAC(ctx, &milvuspb.RestoreRBACMetaRequest{
RBACMeta: resp5.GetRBACMeta(),
})
s.NoError(err)
s.False(merr.Ok(resp6))
// drop exist role/user, successful to restore
resp7, err := s.Cluster.Proxy.OperatePrivilege(ctx, &milvuspb.OperatePrivilegeRequest{
Type: milvuspb.OperatePrivilegeType_Revoke,
Entity: &milvuspb.GrantEntity{
Role: &milvuspb.RoleEntity{Name: roleName},
Object: &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Collection.String()},
ObjectName: util.AnyWord,
DbName: util.AnyWord,
Grantor: &milvuspb.GrantorEntity{
User: &milvuspb.UserEntity{Name: util.UserRoot},
Privilege: &milvuspb.PrivilegeEntity{Name: util.AnyWord},
},
},
})
s.NoError(err)
s.True(merr.Ok(resp7))
resp8, err := s.Cluster.Proxy.DropRole(ctx, &milvuspb.DropRoleRequest{
RoleName: roleName,
})
s.NoError(err)
s.True(merr.Ok(resp8))
resp9, err := s.Cluster.Proxy.DeleteCredential(ctx, &milvuspb.DeleteCredentialRequest{
Username: userName,
})
s.NoError(err)
s.True(merr.Ok(resp9))
resp10, err := s.Cluster.Proxy.RestoreRBAC(ctx, &milvuspb.RestoreRBACMetaRequest{
RBACMeta: resp5.GetRBACMeta(),
})
s.NoError(err)
s.True(merr.Ok(resp10))
// check the restored rbac, should be same as the original one
resp11, err := s.Cluster.Proxy.BackupRBAC(ctx, &milvuspb.BackupRBACMetaRequest{})
s.NoError(err)
s.True(merr.Ok(resp11.GetStatus()))
s.Equal(resp11.GetRBACMeta().String(), resp5.GetRBACMeta().String())
}
func TestRBACBackup(t *testing.T) {
suite.Run(t, new(RBACBackupTestSuite))
}