diff --git a/go.mod b/go.mod index 220b3610a7d38..b2e2b8e215a3f 100644 --- a/go.mod +++ b/go.mod @@ -26,7 +26,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.4.9 + github.com/milvus-io/milvus-proto/go-api/v2 v2.4.9-0.20240816062459-7b3eea802adb 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 @@ -73,6 +73,7 @@ require ( github.com/pkg/errors v0.9.1 github.com/valyala/fastjson v1.6.4 github.com/zeebo/xxh3 v1.0.2 + google.golang.org/protobuf v1.33.0 gopkg.in/yaml.v3 v3.0.1 ) @@ -235,7 +236,6 @@ require ( google.golang.org/genproto v0.0.0-20230822172742-b8732ec3820d // indirect google.golang.org/genproto/googleapis/api v0.0.0-20230822172742-b8732ec3820d // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20240401170217-c3f982113cda // indirect - google.golang.org/protobuf v1.33.0 // indirect gopkg.in/inf.v0 v0.9.1 // indirect gopkg.in/ini.v1 v1.67.0 // indirect gopkg.in/natefinch/lumberjack.v2 v2.0.0 // indirect diff --git a/go.sum b/go.sum index ad3807c450e7c..14390589b945f 100644 --- a/go.sum +++ b/go.sum @@ -606,8 +606,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.4.9 h1:GhdkA7L4hE/TpUc6CzC9QufoaaF7StLRnISLsXexVLo= -github.com/milvus-io/milvus-proto/go-api/v2 v2.4.9/go.mod h1:1OIl0v5PQeNxIJhCvY+K55CBUOYDZevw9g9380u1Wek= +github.com/milvus-io/milvus-proto/go-api/v2 v2.4.9-0.20240816062459-7b3eea802adb h1:DA2PyjTVCA7c23x0UBqzL2+lBrLVH6zP2VPR8yDagr4= +github.com/milvus-io/milvus-proto/go-api/v2 v2.4.9-0.20240816062459-7b3eea802adb/go.mod h1:1OIl0v5PQeNxIJhCvY+K55CBUOYDZevw9g9380u1Wek= github.com/milvus-io/milvus-storage/go v0.0.0-20231227072638-ebd0b8e56d70 h1:Z+sp64fmAOxAG7mU0dfVOXvAXlwRB0c8a96rIM5HevI= github.com/milvus-io/milvus-storage/go v0.0.0-20231227072638-ebd0b8e56d70/go.mod h1:GPETMcTZq1gLY1WA6Na5kiNAKnq8SEMMiVKUZrM3sho= github.com/milvus-io/pulsar-client-go v0.6.10 h1:eqpJjU+/QX0iIhEo3nhOqMNXL+TyInAs1IAHZCrCM/A= diff --git a/internal/datacoord/mock_test.go b/internal/datacoord/mock_test.go index 787c0f977b1f2..aa0aee9c27e94 100644 --- a/internal/datacoord/mock_test.go +++ b/internal/datacoord/mock_test.go @@ -620,6 +620,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{} } diff --git a/internal/distributed/proxy/service.go b/internal/distributed/proxy/service.go index 898d7e5a823f3..9105cddf8b80f 100644 --- a/internal/distributed/proxy/service.go +++ b/internal/distributed/proxy/service.go @@ -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) } diff --git a/internal/distributed/rootcoord/client/client.go b/internal/distributed/rootcoord/client/client.go index 1379df5ce59b6..fbf8d2657f16e 100644 --- a/internal/distributed/rootcoord/client/client.go +++ b/internal/distributed/rootcoord/client/client.go @@ -670,3 +670,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) + }) +} diff --git a/internal/distributed/rootcoord/service.go b/internal/distributed/rootcoord/service.go index 42b326fab5716..150d232d77fe6 100644 --- a/internal/distributed/rootcoord/service.go +++ b/internal/distributed/rootcoord/service.go @@ -528,3 +528,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) +} diff --git a/internal/metastore/catalog.go b/internal/metastore/catalog.go index 996805253e495..0c830117a702f 100644 --- a/internal/metastore/catalog.go +++ b/internal/metastore/catalog.go @@ -81,6 +81,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() } diff --git a/internal/metastore/kv/rootcoord/kv_catalog.go b/internal/metastore/kv/rootcoord/kv_catalog.go index 367ff07f42ff1..f3c8463926f4b 100644 --- a/internal/metastore/kv/rootcoord/kv_catalog.go +++ b/internal/metastore/kv/rootcoord/kv_catalog.go @@ -8,6 +8,7 @@ import ( "github.com/cockroachdb/errors" "github.com/golang/protobuf/proto" + "github.com/samber/lo" "go.uber.org/zap" "github.com/milvus-io/milvus-proto/go-api/v2/milvuspb" @@ -747,23 +748,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 { @@ -1217,6 +1232,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 } diff --git a/internal/metastore/kv/rootcoord/kv_catalog_test.go b/internal/metastore/kv/rootcoord/kv_catalog_test.go index 236ba03ad6aa7..fbdf8c244c0e7 100644 --- a/internal/metastore/kv/rootcoord/kv_catalog_test.go +++ b/internal/metastore/kv/rootcoord/kv_catalog_test.go @@ -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" @@ -1462,7 +1464,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() @@ -1920,7 +1935,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 @@ -2533,6 +2555,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} diff --git a/internal/metastore/mocks/mock_rootcoord_catalog.go b/internal/metastore/mocks/mock_rootcoord_catalog.go index 20208d9cd6ff7..ca2ed2f27f18a 100644 --- a/internal/metastore/mocks/mock_rootcoord_catalog.go +++ b/internal/metastore/mocks/mock_rootcoord_catalog.go @@ -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 { diff --git a/internal/mocks/mock_proxy.go b/internal/mocks/mock_proxy.go index ab28fd9f87ac4..36555c92ef429 100644 --- a/internal/mocks/mock_proxy.go +++ b/internal/mocks/mock_proxy.go @@ -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) diff --git a/internal/mocks/mock_rootcoord.go b/internal/mocks/mock_rootcoord.go index d41d2cbd9ab1a..00c997d90e9e0 100644 --- a/internal/mocks/mock_rootcoord.go +++ b/internal/mocks/mock_rootcoord.go @@ -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) @@ -2153,6 +2208,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) diff --git a/internal/mocks/mock_rootcoord_client.go b/internal/mocks/mock_rootcoord_client.go index 4a1ba8099bf32..dbb566bc3d4c3 100644 --- a/internal/mocks/mock_rootcoord_client.go +++ b/internal/mocks/mock_rootcoord_client.go @@ -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)) @@ -2664,6 +2734,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)) diff --git a/internal/proto/root_coord.proto b/internal/proto/root_coord.proto index fb65591f4e6e5..a1408f6e29656 100644 --- a/internal/proto/root_coord.proto +++ b/internal/proto/root_coord.proto @@ -132,6 +132,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) {} diff --git a/internal/proxy/impl.go b/internal/proxy/impl.go index b08437db162ce..c32e096766da3 100644 --- a/internal/proxy/impl.go +++ b/internal/proxy/impl.go @@ -5261,6 +5261,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() diff --git a/internal/proxy/rootcoord_mock_test.go b/internal/proxy/rootcoord_mock_test.go index 97f466f2e9df0..166a64e5c859a 100644 --- a/internal/proxy/rootcoord_mock_test.go +++ b/internal/proxy/rootcoord_mock_test.go @@ -1122,6 +1122,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) diff --git a/internal/rootcoord/meta_table.go b/internal/rootcoord/meta_table.go index 49bca49368780..675bfa342bd7f 100644 --- a/internal/rootcoord/meta_table.go +++ b/internal/rootcoord/meta_table.go @@ -94,6 +94,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 } type MetaTable struct { @@ -1404,3 +1406,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) +} diff --git a/internal/rootcoord/meta_table_test.go b/internal/rootcoord/meta_table_test.go index b916facaf1009..1aed09cc80970 100644 --- a/internal/rootcoord/meta_table_test.go +++ b/internal/rootcoord/meta_table_test.go @@ -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) +} diff --git a/internal/rootcoord/mocks/meta_table.go b/internal/rootcoord/mocks/meta_table.go index f67bc65c8fcb8..9dc0aee9c9e08 100644 --- a/internal/rootcoord/mocks/meta_table.go +++ b/internal/rootcoord/mocks/meta_table.go @@ -334,6 +334,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) @@ -1940,6 +1995,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) diff --git a/internal/rootcoord/root_coord.go b/internal/rootcoord/root_coord.go index 8198a2a3d5cb6..caea15ed0b991 100644 --- a/internal/rootcoord/root_coord.go +++ b/internal/rootcoord/root_coord.go @@ -2696,6 +2696,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 diff --git a/internal/rootcoord/root_coord_test.go b/internal/rootcoord/root_coord_test.go index 5fea892a722a6..11dce7e364370 100644 --- a/internal/rootcoord/root_coord_test.go +++ b/internal/rootcoord/root_coord_test.go @@ -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 } diff --git a/internal/util/mock/grpc_rootcoord_client.go b/internal/util/mock/grpc_rootcoord_client.go index 097303b65178a..bc80d954ae250 100644 --- a/internal/util/mock/grpc_rootcoord_client.go +++ b/internal/util/mock/grpc_rootcoord_client.go @@ -262,6 +262,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 } diff --git a/pkg/go.mod b/pkg/go.mod index 323f3198e6c62..d2cde85140844 100644 --- a/pkg/go.mod +++ b/pkg/go.mod @@ -13,7 +13,7 @@ require ( 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.4.9 + github.com/milvus-io/milvus-proto/go-api/v2 v2.4.9-0.20240816062459-7b3eea802adb 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 diff --git a/pkg/go.sum b/pkg/go.sum index d4b15b5b36348..ba27b833a416f 100644 --- a/pkg/go.sum +++ b/pkg/go.sum @@ -503,8 +503,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.4.9 h1:GhdkA7L4hE/TpUc6CzC9QufoaaF7StLRnISLsXexVLo= -github.com/milvus-io/milvus-proto/go-api/v2 v2.4.9/go.mod h1:1OIl0v5PQeNxIJhCvY+K55CBUOYDZevw9g9380u1Wek= +github.com/milvus-io/milvus-proto/go-api/v2 v2.4.9-0.20240816062459-7b3eea802adb h1:DA2PyjTVCA7c23x0UBqzL2+lBrLVH6zP2VPR8yDagr4= +github.com/milvus-io/milvus-proto/go-api/v2 v2.4.9-0.20240816062459-7b3eea802adb/go.mod h1:1OIl0v5PQeNxIJhCvY+K55CBUOYDZevw9g9380u1Wek= 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= diff --git a/pkg/util/constant.go b/pkg/util/constant.go index 4c73ebedac19d..1abefb3a00ecc 100644 --- a/pkg/util/constant.go +++ b/pkg/util/constant.go @@ -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()), diff --git a/tests/integration/rbac/rbac_backup_test.go b/tests/integration/rbac/rbac_backup_test.go new file mode 100644 index 0000000000000..9b95b747ded49 --- /dev/null +++ b/tests/integration/rbac/rbac_backup_test.go @@ -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)) +}