admin/milvus
1
0
Fork 0
mirror of https://gitee.com/milvus-io/milvus.git synced 2025-12-07 01:28:27 +08:00
Code Issues Packages Projects Releases Wiki Activity

enhance: refine rootcoord/metatable interfaces to ensure that each method includes a ctx parameter (#37846)

Browse Source 
issue: #35917
Before enhancing log trace information, it's necessary to pass the
context to the method entry point.
This PR first refine the rootcoord/metatable interfaces to ensure that
each method includes a ctx parameter.

Signed-off-by: tinswzy <zhenyuan.wei@zilliz.com>
This commit is contained in:
tinswzy 2024-11-21 19:30:32 +08:00 committed by GitHub
parent 965bda6e60
commit e247ff9ee7
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
16 changed files with 656 additions and 625 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
internal/rootcoord/alter_collection_task.go
View File

@ -84,7 +84,7 @@ func (a *alterCollectionTask) Execute(ctx context.Context) error {
})
// properties needs to be refreshed in the cache
aliases := a.core.meta.ListAliasesByID(oldColl.CollectionID)
aliases := a.core.meta.ListAliasesByID(ctx, oldColl.CollectionID)
redoTask.AddSyncStep(&expireCacheStep{
baseStep: baseStep{core: a.core},
dbName: a.Req.GetDbName(),

8
internal/rootcoord/alter_collection_task_test.go
View File

@ -92,7 +92,7 @@ func Test_alterCollectionTask_Execute(t *testing.T) {
mock.Anything,
mock.Anything,
).Return(errors.New("err"))
meta.On("ListAliasesByID", mock.Anything).Return([]string{})
meta.On("ListAliasesByID", mock.Anything, mock.Anything).Return([]string{})
core := newTestCore(withValidProxyManager(), withMeta(meta))
task := &alterCollectionTask{
@ -122,7 +122,7 @@ func Test_alterCollectionTask_Execute(t *testing.T) {
mock.Anything,
mock.Anything,
).Return(nil)
meta.On("ListAliasesByID", mock.Anything).Return([]string{})
meta.On("ListAliasesByID", mock.Anything, mock.Anything).Return([]string{})
broker := newMockBroker()
broker.BroadcastAlteredCollectionFunc = func(ctx context.Context, req *milvuspb.AlterCollectionRequest) error {
@ -157,7 +157,7 @@ func Test_alterCollectionTask_Execute(t *testing.T) {
mock.Anything,
mock.Anything,
).Return(nil)
meta.On("ListAliasesByID", mock.Anything).Return([]string{})
meta.On("ListAliasesByID", mock.Anything, mock.Anything).Return([]string{})
broker := newMockBroker()
broker.BroadcastAlteredCollectionFunc = func(ctx context.Context, req *milvuspb.AlterCollectionRequest) error {
@ -231,7 +231,7 @@ func Test_alterCollectionTask_Execute(t *testing.T) {
mock.Anything,
mock.Anything,
).Return(nil)
meta.On("ListAliasesByID", mock.Anything).Return([]string{})
meta.On("ListAliasesByID", mock.Anything, mock.Anything).Return([]string{})
broker := newMockBroker()
broker.BroadcastAlteredCollectionFunc = func(ctx context.Context, req *milvuspb.AlterCollectionRequest) error {

2
internal/rootcoord/describe_collection_task.go
View File

@ -45,7 +45,7 @@ func (t *describeCollectionTask) Execute(ctx context.Context) (err error) {
return err
}
aliases := t.core.meta.ListAliasesByID(coll.CollectionID)
aliases := t.core.meta.ListAliasesByID(ctx, coll.CollectionID)
db, err := t.core.meta.GetDatabaseByID(ctx, coll.DBID, t.GetTs())
if err != nil {
return err

1
internal/rootcoord/describe_collection_task_test.go
View File

@ -105,6 +105,7 @@ func Test_describeCollectionTask_Execute(t *testing.T) {
}, nil)
meta.On("ListAliasesByID",
mock.Anything,
mock.Anything,
).Return([]string{alias1, alias2})
meta.EXPECT().GetDatabaseByID(mock.Anything, mock.Anything, mock.Anything).Return(&model.Database{
ID: 1,

8
internal/rootcoord/drop_collection_task.go
View File

@ -37,18 +37,18 @@ type dropCollectionTask struct {
Req *milvuspb.DropCollectionRequest
}
func (t *dropCollectionTask) validate() error {
func (t *dropCollectionTask) validate(ctx context.Context) error {
if err := CheckMsgType(t.Req.GetBase().GetMsgType(), commonpb.MsgType_DropCollection); err != nil {
return err
}
if t.core.meta.IsAlias(t.Req.GetDbName(), t.Req.GetCollectionName()) {
if t.core.meta.IsAlias(ctx, t.Req.GetDbName(), t.Req.GetCollectionName()) {
return fmt.Errorf("cannot drop the collection via alias = %s", t.Req.CollectionName)
}
return nil
}
func (t *dropCollectionTask) Prepare(ctx context.Context) error {
return t.validate()
return t.validate(ctx)
}
func (t *dropCollectionTask) Execute(ctx context.Context) error {
@ -68,7 +68,7 @@ func (t *dropCollectionTask) Execute(ctx context.Context) error {
}
// meta cache of all aliases should also be cleaned.
aliases := t.core.meta.ListAliasesByID(collMeta.CollectionID)
aliases := t.core.meta.ListAliasesByID(ctx, collMeta.CollectionID)
ts := t.GetTs()

5
internal/rootcoord/drop_collection_task_test.go
View File

@ -51,6 +51,7 @@ func Test_dropCollectionTask_Prepare(t *testing.T) {
meta.On("IsAlias",
mock.Anything,
mock.Anything,
mock.Anything,
).Return(true)
core := newTestCore(withMeta(meta))
@ -72,6 +73,7 @@ func Test_dropCollectionTask_Prepare(t *testing.T) {
meta.On("IsAlias",
mock.Anything,
mock.Anything,
mock.Anything,
).Return(false)
core := newTestCore(withMeta(meta))
@ -129,6 +131,7 @@ func Test_dropCollectionTask_Execute(t *testing.T) {
mock.Anything,
).Return(coll.Clone(), nil)
meta.On("ListAliasesByID",
mock.Anything,
mock.AnythingOfType("int64"),
).Return([]string{})
@ -163,6 +166,7 @@ func Test_dropCollectionTask_Execute(t *testing.T) {
).Return(errors.New("error mock ChangeCollectionState"))
meta.On("ListAliasesByID",
mock.Anything,
mock.Anything,
).Return([]string{})
core := newTestCore(withValidProxyManager(), withMeta(meta))
@ -207,6 +211,7 @@ func Test_dropCollectionTask_Execute(t *testing.T) {
).Return(nil)
meta.On("ListAliasesByID",
mock.Anything,
mock.Anything,
).Return([]string{})
removeCollectionMetaCalled := false
removeCollectionMetaChan := make(chan struct{}, 1)

4
internal/rootcoord/list_db_task.go
View File

@ -58,7 +58,7 @@ func (t *listDatabaseTask) Execute(ctx context.Context) error {
privilegeDBs.Insert(util.AnyWord)
return privilegeDBs, nil
}
userRoles, err := t.core.meta.SelectUser("", &milvuspb.UserEntity{
userRoles, err := t.core.meta.SelectUser(ctx, "", &milvuspb.UserEntity{
Name: curUser,
}, true)
if err != nil {
@ -72,7 +72,7 @@ func (t *listDatabaseTask) Execute(ctx context.Context) error {
privilegeDBs.Insert(util.AnyWord)
return privilegeDBs, nil
}
entities, err := t.core.meta.SelectGrant("", &milvuspb.GrantEntity{
entities, err := t.core.meta.SelectGrant(ctx, "", &milvuspb.GrantEntity{
Role: role,
DbName: util.AnyWord,
})

18
internal/rootcoord/list_db_task_test.go
View File

@ -132,7 +132,7 @@ func Test_ListDBTask(t *testing.T) {
{
// select role fail
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything).
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
Return(nil, errors.New("mock select user error")).Once()
ctx := GetContext(context.Background(), "foo:root")
task := getTask()
@ -142,7 +142,7 @@ func Test_ListDBTask(t *testing.T) {
{
// select role, empty result
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything).
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
Return([]*milvuspb.UserResult{}, nil).Once()
ctx := GetContext(context.Background(), "foo:root")
task := getTask()
@ -153,7 +153,7 @@ func Test_ListDBTask(t *testing.T) {
{
// select role, the user is added to admin role
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything).
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
Return([]*milvuspb.UserResult{
{
User: &milvuspb.UserEntity{
@ -176,7 +176,7 @@ func Test_ListDBTask(t *testing.T) {
{
// select grant fail
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything).
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
Return([]*milvuspb.UserResult{
{
User: &milvuspb.UserEntity{
@ -189,7 +189,7 @@ func Test_ListDBTask(t *testing.T) {
},
},
}, nil).Once()
meta.EXPECT().SelectGrant(mock.Anything, mock.Anything).
meta.EXPECT().SelectGrant(mock.Anything, mock.Anything, mock.Anything).
Return(nil, errors.New("mock select grant error")).Once()
ctx := GetContext(context.Background(), "foo:root")
task := getTask()
@ -199,7 +199,7 @@ func Test_ListDBTask(t *testing.T) {
{
// normal user
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything).
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
Return([]*milvuspb.UserResult{
{
User: &milvuspb.UserEntity{
@ -220,7 +220,7 @@ func Test_ListDBTask(t *testing.T) {
Name: "default",
},
}, nil).Once()
meta.EXPECT().SelectGrant(mock.Anything, mock.Anything).
meta.EXPECT().SelectGrant(mock.Anything, mock.Anything, mock.Anything).
Return([]*milvuspb.GrantEntity{
{
DbName: "fooDB",
@ -236,7 +236,7 @@ func Test_ListDBTask(t *testing.T) {
{
// normal user with any db privilege
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything).
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
Return([]*milvuspb.UserResult{
{
User: &milvuspb.UserEntity{
@ -257,7 +257,7 @@ func Test_ListDBTask(t *testing.T) {
Name: "default",
},
}, nil).Once()
meta.EXPECT().SelectGrant(mock.Anything, mock.Anything).
meta.EXPECT().SelectGrant(mock.Anything, mock.Anything, mock.Anything).
Return([]*milvuspb.GrantEntity{
{
DbName: "*",

174
internal/rootcoord/meta_table.go
View File

@ -61,9 +61,9 @@ type IMetaTable interface {
GetCollectionByIDWithMaxTs(ctx context.Context, collectionID UniqueID) (*model.Collection, error)
ListCollections(ctx context.Context, dbName string, ts Timestamp, onlyAvail bool) ([]*model.Collection, error)
ListAllAvailCollections(ctx context.Context) map[int64][]int64
ListCollectionPhysicalChannels() map[typeutil.UniqueID][]string
GetCollectionVirtualChannels(colID int64) []string
GetPChannelInfo(pchannel string) *rootcoordpb.GetPChannelInfoResponse
ListCollectionPhysicalChannels(ctx context.Context) map[typeutil.UniqueID][]string
GetCollectionVirtualChannels(ctx context.Context, colID int64) []string
GetPChannelInfo(ctx context.Context, pchannel string) *rootcoordpb.GetPChannelInfoResponse
AddPartition(ctx context.Context, partition *model.Partition) error
ChangePartitionState(ctx context.Context, collectionID UniqueID, partitionID UniqueID, state pb.PartitionState, ts Timestamp) error
RemovePartition(ctx context.Context, dbID int64, collectionID UniqueID, partitionID UniqueID, ts Timestamp) error
@ -76,35 +76,33 @@ type IMetaTable interface {
RenameCollection(ctx context.Context, dbName string, oldName string, newDBName string, newName string, ts Timestamp) error
// TODO: it'll be a big cost if we handle the time travel logic, since we should always list all aliases in catalog.
IsAlias(db, name string) bool
ListAliasesByID(collID UniqueID) []string
IsAlias(ctx context.Context, db, name string) bool
ListAliasesByID(ctx context.Context, collID UniqueID) []string
// TODO: better to accept ctx.
AddCredential(credInfo *internalpb.CredentialInfo) error
GetCredential(username string) (*internalpb.CredentialInfo, error)
DeleteCredential(username string) error
AlterCredential(credInfo *internalpb.CredentialInfo) error
ListCredentialUsernames() (*milvuspb.ListCredUsersResponse, error)
AddCredential(ctx context.Context, credInfo *internalpb.CredentialInfo) error
GetCredential(ctx context.Context, username string) (*internalpb.CredentialInfo, error)
DeleteCredential(ctx context.Context, username string) error
AlterCredential(ctx context.Context, credInfo *internalpb.CredentialInfo) error
ListCredentialUsernames(ctx context.Context) (*milvuspb.ListCredUsersResponse, error)
// TODO: better to accept ctx.
CreateRole(tenant string, entity *milvuspb.RoleEntity) error
DropRole(tenant string, roleName string) error
OperateUserRole(tenant string, userEntity *milvuspb.UserEntity, roleEntity *milvuspb.RoleEntity, operateType milvuspb.OperateUserRoleType) error
SelectRole(tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error)
SelectUser(tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error)
OperatePrivilege(tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error
SelectGrant(tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error)
DropGrant(tenant string, role *milvuspb.RoleEntity) error
ListPolicy(tenant string) ([]string, error)
ListUserRole(tenant string) ([]string, error)
CreateRole(ctx context.Context, tenant string, entity *milvuspb.RoleEntity) error
DropRole(ctx context.Context, tenant string, roleName string) error
OperateUserRole(ctx context.Context, tenant string, userEntity *milvuspb.UserEntity, roleEntity *milvuspb.RoleEntity, operateType milvuspb.OperateUserRoleType) error
SelectRole(ctx context.Context, tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error)
SelectUser(ctx context.Context, tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error)
OperatePrivilege(ctx context.Context, tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error
SelectGrant(ctx context.Context, tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error)
DropGrant(ctx context.Context, tenant string, role *milvuspb.RoleEntity) error
ListPolicy(ctx context.Context, tenant string) ([]string, error)
ListUserRole(ctx context.Context, tenant string) ([]string, error)
BackupRBAC(ctx context.Context, tenant string) (*milvuspb.RBACMeta, error)
RestoreRBAC(ctx context.Context, tenant string, meta *milvuspb.RBACMeta) error
IsCustomPrivilegeGroup(groupName string) (bool, error)
CreatePrivilegeGroup(groupName string) error
DropPrivilegeGroup(groupName string) error
ListPrivilegeGroups() ([]*milvuspb.PrivilegeGroupInfo, error)
OperatePrivilegeGroup(groupName string, privileges []*milvuspb.PrivilegeEntity, operateType milvuspb.OperatePrivilegeGroupType) error
GetPrivilegeGroupRoles(groupName string) ([]*milvuspb.RoleEntity, error)
IsCustomPrivilegeGroup(ctx context.Context, groupName string) (bool, error)
CreatePrivilegeGroup(ctx context.Context, groupName string) error
DropPrivilegeGroup(ctx context.Context, groupName string) error
ListPrivilegeGroups(ctx context.Context) ([]*milvuspb.PrivilegeGroupInfo, error)
OperatePrivilegeGroup(ctx context.Context, groupName string, privileges []*milvuspb.PrivilegeEntity, operateType milvuspb.OperatePrivilegeGroupType) error
GetPrivilegeGroupRoles(ctx context.Context, groupName string) ([]*milvuspb.RoleEntity, error)
}
// MetaTable is a persistent meta set of all databases, collections and partitions.
@ -729,7 +727,7 @@ func (mt *MetaTable) listCollectionFromCache(dbName string, onlyAvail bool) ([]*
}
// ListCollectionPhysicalChannels list physical channels of all collections.
func (mt *MetaTable) ListCollectionPhysicalChannels() map[typeutil.UniqueID][]string {
func (mt *MetaTable) ListCollectionPhysicalChannels(ctx context.Context) map[typeutil.UniqueID][]string {
mt.ddLock.RLock()
defer mt.ddLock.RUnlock()
@ -838,7 +836,7 @@ func (mt *MetaTable) RenameCollection(ctx context.Context, dbName string, oldNam
}
// GetCollectionVirtualChannels returns virtual channels of a given collection.
func (mt *MetaTable) GetCollectionVirtualChannels(colID int64) []string {
func (mt *MetaTable) GetCollectionVirtualChannels(ctx context.Context, colID int64) []string {
mt.ddLock.RLock()
defer mt.ddLock.RUnlock()
for id, collInfo := range mt.collID2Meta {
@ -850,7 +848,7 @@ func (mt *MetaTable) GetCollectionVirtualChannels(colID int64) []string {
}
// GetPChannelInfo returns infos on pchannel.
func (mt *MetaTable) GetPChannelInfo(pchannel string) *rootcoordpb.GetPChannelInfoResponse {
func (mt *MetaTable) GetPChannelInfo(ctx context.Context, pchannel string) *rootcoordpb.GetPChannelInfoResponse {
mt.ddLock.RLock()
defer mt.ddLock.RUnlock()
resp := &rootcoordpb.GetPChannelInfoResponse{
@ -1199,7 +1197,7 @@ func (mt *MetaTable) ListAliases(ctx context.Context, dbName string, collectionN
return aliases, nil
}
func (mt *MetaTable) IsAlias(db, name string) bool {
func (mt *MetaTable) IsAlias(ctx context.Context, db, name string) bool {
mt.ddLock.RLock()
defer mt.ddLock.RUnlock()
@ -1218,7 +1216,7 @@ func (mt *MetaTable) listAliasesByID(collID UniqueID) []string {
return ret
}
func (mt *MetaTable) ListAliasesByID(collID UniqueID) []string {
func (mt *MetaTable) ListAliasesByID(ctx context.Context, collID UniqueID) []string {
mt.ddLock.RLock()
defer mt.ddLock.RUnlock()
@ -1226,14 +1224,14 @@ func (mt *MetaTable) ListAliasesByID(collID UniqueID) []string {
}
// AddCredential add credential
func (mt *MetaTable) AddCredential(credInfo *internalpb.CredentialInfo) error {
func (mt *MetaTable) AddCredential(ctx context.Context, credInfo *internalpb.CredentialInfo) error {
if credInfo.Username == "" {
return fmt.Errorf("username is empty")
}
mt.permissionLock.Lock()
defer mt.permissionLock.Unlock()
usernames, err := mt.catalog.ListCredentials(mt.ctx)
usernames, err := mt.catalog.ListCredentials(ctx)
if err != nil {
return err
}
@ -1243,7 +1241,7 @@ func (mt *MetaTable) AddCredential(credInfo *internalpb.CredentialInfo) error {
return errors.New(errMsg)
}
if origin, _ := mt.catalog.GetCredential(mt.ctx, credInfo.Username); origin != nil {
if origin, _ := mt.catalog.GetCredential(ctx, credInfo.Username); origin != nil {
return fmt.Errorf("user already exists: %s", credInfo.Username)
}
@ -1251,11 +1249,11 @@ func (mt *MetaTable) AddCredential(credInfo *internalpb.CredentialInfo) error {
Username: credInfo.Username,
EncryptedPassword: credInfo.EncryptedPassword,
}
return mt.catalog.CreateCredential(mt.ctx, credential)
return mt.catalog.CreateCredential(ctx, credential)
}
// AlterCredential update credential
func (mt *MetaTable) AlterCredential(credInfo *internalpb.CredentialInfo) error {
func (mt *MetaTable) AlterCredential(ctx context.Context, credInfo *internalpb.CredentialInfo) error {
if credInfo.Username == "" {
return fmt.Errorf("username is empty")
}
@ -1267,32 +1265,32 @@ func (mt *MetaTable) AlterCredential(credInfo *internalpb.CredentialInfo) error
Username: credInfo.Username,
EncryptedPassword: credInfo.EncryptedPassword,
}
return mt.catalog.AlterCredential(mt.ctx, credential)
return mt.catalog.AlterCredential(ctx, credential)
}
// GetCredential get credential by username
func (mt *MetaTable) GetCredential(username string) (*internalpb.CredentialInfo, error) {
func (mt *MetaTable) GetCredential(ctx context.Context, username string) (*internalpb.CredentialInfo, error) {
mt.permissionLock.RLock()
defer mt.permissionLock.RUnlock()
credential, err := mt.catalog.GetCredential(mt.ctx, username)
credential, err := mt.catalog.GetCredential(ctx, username)
return model.MarshalCredentialModel(credential), err
}
// DeleteCredential delete credential
func (mt *MetaTable) DeleteCredential(username string) error {
func (mt *MetaTable) DeleteCredential(ctx context.Context, username string) error {
mt.permissionLock.Lock()
defer mt.permissionLock.Unlock()
return mt.catalog.DropCredential(mt.ctx, username)
return mt.catalog.DropCredential(ctx, username)
}
// ListCredentialUsernames list credential usernames
func (mt *MetaTable) ListCredentialUsernames() (*milvuspb.ListCredUsersResponse, error) {
func (mt *MetaTable) ListCredentialUsernames(ctx context.Context) (*milvuspb.ListCredUsersResponse, error) {
mt.permissionLock.RLock()
defer mt.permissionLock.RUnlock()
usernames, err := mt.catalog.ListCredentials(mt.ctx)
usernames, err := mt.catalog.ListCredentials(ctx)
if err != nil {
return nil, fmt.Errorf("list credential usernames err:%w", err)
}
@ -1300,14 +1298,14 @@ func (mt *MetaTable) ListCredentialUsernames() (*milvuspb.ListCredUsersResponse,
}
// CreateRole create role
func (mt *MetaTable) CreateRole(tenant string, entity *milvuspb.RoleEntity) error {
func (mt *MetaTable) CreateRole(ctx context.Context, tenant string, entity *milvuspb.RoleEntity) error {
if funcutil.IsEmptyString(entity.Name) {
return fmt.Errorf("the role name in the role info is empty")
}
mt.permissionLock.Lock()
defer mt.permissionLock.Unlock()
results, err := mt.catalog.ListRole(mt.ctx, tenant, nil, false)
results, err := mt.catalog.ListRole(ctx, tenant, nil, false)
if err != nil {
log.Warn("fail to list roles", zap.Error(err))
return err
@ -1324,19 +1322,19 @@ func (mt *MetaTable) CreateRole(tenant string, entity *milvuspb.RoleEntity) erro
return errors.New(errMsg)
}
return mt.catalog.CreateRole(mt.ctx, tenant, entity)
return mt.catalog.CreateRole(ctx, tenant, entity)
}
// DropRole drop role info
func (mt *MetaTable) DropRole(tenant string, roleName string) error {
func (mt *MetaTable) DropRole(ctx context.Context, tenant string, roleName string) error {
mt.permissionLock.Lock()
defer mt.permissionLock.Unlock()
return mt.catalog.DropRole(mt.ctx, tenant, roleName)
return mt.catalog.DropRole(ctx, tenant, roleName)
}
// OperateUserRole operate the relationship between a user and a role, including adding a user to a role and removing a user from a role
func (mt *MetaTable) OperateUserRole(tenant string, userEntity *milvuspb.UserEntity, roleEntity *milvuspb.RoleEntity, operateType milvuspb.OperateUserRoleType) error {
func (mt *MetaTable) OperateUserRole(ctx context.Context, tenant string, userEntity *milvuspb.UserEntity, roleEntity *milvuspb.RoleEntity, operateType milvuspb.OperateUserRoleType) error {
if funcutil.IsEmptyString(userEntity.Name) {
return fmt.Errorf("username in the user entity is empty")
}
@ -1347,31 +1345,31 @@ func (mt *MetaTable) OperateUserRole(tenant string, userEntity *milvuspb.UserEnt
mt.permissionLock.Lock()
defer mt.permissionLock.Unlock()
return mt.catalog.AlterUserRole(mt.ctx, tenant, userEntity, roleEntity, operateType)
return mt.catalog.AlterUserRole(ctx, tenant, userEntity, roleEntity, operateType)
}
// SelectRole select role.
// Enter the role condition by the entity param. And this param is nil, which means selecting all roles.
// Get all users that are added to the role by setting the includeUserInfo param to true.
func (mt *MetaTable) SelectRole(tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
func (mt *MetaTable) SelectRole(ctx context.Context, tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
mt.permissionLock.RLock()
defer mt.permissionLock.RUnlock()
return mt.catalog.ListRole(mt.ctx, tenant, entity, includeUserInfo)
return mt.catalog.ListRole(ctx, tenant, entity, includeUserInfo)
}
// SelectUser select user.
// Enter the user condition by the entity param. And this param is nil, which means selecting all users.
// Get all roles that are added the user to by setting the includeRoleInfo param to true.
func (mt *MetaTable) SelectUser(tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error) {
func (mt *MetaTable) SelectUser(ctx context.Context, tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error) {
mt.permissionLock.RLock()
defer mt.permissionLock.RUnlock()
return mt.catalog.ListUser(mt.ctx, tenant, entity, includeRoleInfo)
return mt.catalog.ListUser(ctx, tenant, entity, includeRoleInfo)
}
// OperatePrivilege grant or revoke privilege by setting the operateType param
func (mt *MetaTable) OperatePrivilege(tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error {
func (mt *MetaTable) OperatePrivilege(ctx context.Context, tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error {
if funcutil.IsEmptyString(entity.ObjectName) {
return fmt.Errorf("the object name in the grant entity is empty")
}
@ -1400,13 +1398,13 @@ func (mt *MetaTable) OperatePrivilege(tenant string, entity *milvuspb.GrantEntit
mt.permissionLock.Lock()
defer mt.permissionLock.Unlock()
return mt.catalog.AlterGrant(mt.ctx, tenant, entity, operateType)
return mt.catalog.AlterGrant(ctx, tenant, entity, operateType)
}
// SelectGrant select grant
// The principal entity MUST be not empty in the grant entity
// The resource entity and the resource name are optional, and the two params should be not empty together when you select some grants about the resource kind.
func (mt *MetaTable) SelectGrant(tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error) {
func (mt *MetaTable) SelectGrant(ctx context.Context, tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error) {
var entities []*milvuspb.GrantEntity
if entity == nil {
return entities, fmt.Errorf("the grant entity is nil")
@ -1422,50 +1420,50 @@ func (mt *MetaTable) SelectGrant(tenant string, entity *milvuspb.GrantEntity) ([
mt.permissionLock.RLock()
defer mt.permissionLock.RUnlock()
return mt.catalog.ListGrant(mt.ctx, tenant, entity)
return mt.catalog.ListGrant(ctx, tenant, entity)
}
func (mt *MetaTable) DropGrant(tenant string, role *milvuspb.RoleEntity) error {
func (mt *MetaTable) DropGrant(ctx context.Context, tenant string, role *milvuspb.RoleEntity) error {
if role == nil || funcutil.IsEmptyString(role.Name) {
return fmt.Errorf("the role entity is invalid when dropping the grant")
}
mt.permissionLock.Lock()
defer mt.permissionLock.Unlock()
return mt.catalog.DeleteGrant(mt.ctx, tenant, role)
return mt.catalog.DeleteGrant(ctx, tenant, role)
}
func (mt *MetaTable) ListPolicy(tenant string) ([]string, error) {
func (mt *MetaTable) ListPolicy(ctx context.Context, tenant string) ([]string, error) {
mt.permissionLock.RLock()
defer mt.permissionLock.RUnlock()
return mt.catalog.ListPolicy(mt.ctx, tenant)
return mt.catalog.ListPolicy(ctx, tenant)
}
func (mt *MetaTable) ListUserRole(tenant string) ([]string, error) {
func (mt *MetaTable) ListUserRole(ctx context.Context, tenant string) ([]string, error) {
mt.permissionLock.RLock()
defer mt.permissionLock.RUnlock()
return mt.catalog.ListUserRole(mt.ctx, tenant)
return mt.catalog.ListUserRole(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)
return mt.catalog.BackupRBAC(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)
return mt.catalog.RestoreRBAC(ctx, tenant, meta)
}
// check if the privielge group name is defined by users
func (mt *MetaTable) IsCustomPrivilegeGroup(groupName string) (bool, error) {
privGroups, err := mt.catalog.ListPrivilegeGroups(mt.ctx)
func (mt *MetaTable) IsCustomPrivilegeGroup(ctx context.Context, groupName string) (bool, error) {
privGroups, err := mt.catalog.ListPrivilegeGroups(ctx)
if err != nil {
return false, err
}
@ -1477,14 +1475,14 @@ func (mt *MetaTable) IsCustomPrivilegeGroup(groupName string) (bool, error) {
return false, nil
}
func (mt *MetaTable) CreatePrivilegeGroup(groupName string) error {
func (mt *MetaTable) CreatePrivilegeGroup(ctx context.Context, groupName string) error {
if funcutil.IsEmptyString(groupName) {
return fmt.Errorf("the privilege group name is empty")
}
mt.permissionLock.Lock()
defer mt.permissionLock.Unlock()
definedByUsers, err := mt.IsCustomPrivilegeGroup(groupName)
definedByUsers, err := mt.IsCustomPrivilegeGroup(ctx, groupName)
if err != nil {
return err
}
@ -1498,17 +1496,17 @@ func (mt *MetaTable) CreatePrivilegeGroup(groupName string) error {
GroupName: groupName,
Privileges: make([]*milvuspb.PrivilegeEntity, 0),
}
return mt.catalog.SavePrivilegeGroup(mt.ctx, data)
return mt.catalog.SavePrivilegeGroup(ctx, data)
}
func (mt *MetaTable) DropPrivilegeGroup(groupName string) error {
func (mt *MetaTable) DropPrivilegeGroup(ctx context.Context, groupName string) error {
if funcutil.IsEmptyString(groupName) {
return fmt.Errorf("the privilege group name is empty")
}
mt.permissionLock.Lock()
defer mt.permissionLock.Unlock()
definedByUsers, err := mt.IsCustomPrivilegeGroup(groupName)
definedByUsers, err := mt.IsCustomPrivilegeGroup(ctx, groupName)
if err != nil {
return err
}
@ -1516,7 +1514,7 @@ func (mt *MetaTable) DropPrivilegeGroup(groupName string) error {
return nil
}
// check if the group is used by any role
roles, err := mt.catalog.ListRole(mt.ctx, util.DefaultTenant, nil, false)
roles, err := mt.catalog.ListRole(ctx, util.DefaultTenant, nil, false)
if err != nil {
return err
}
@ -1524,7 +1522,7 @@ func (mt *MetaTable) DropPrivilegeGroup(groupName string) error {
return entity.GetRole()
})
for _, role := range roleEntity {
grants, err := mt.catalog.ListGrant(mt.ctx, util.DefaultTenant, &milvuspb.GrantEntity{
grants, err := mt.catalog.ListGrant(ctx, util.DefaultTenant, &milvuspb.GrantEntity{
Role: role,
DbName: util.AnyWord,
})
@ -1537,17 +1535,17 @@ func (mt *MetaTable) DropPrivilegeGroup(groupName string) error {
}
}
}
return mt.catalog.DropPrivilegeGroup(mt.ctx, groupName)
return mt.catalog.DropPrivilegeGroup(ctx, groupName)
}
func (mt *MetaTable) ListPrivilegeGroups() ([]*milvuspb.PrivilegeGroupInfo, error) {
func (mt *MetaTable) ListPrivilegeGroups(ctx context.Context) ([]*milvuspb.PrivilegeGroupInfo, error) {
mt.permissionLock.Lock()
defer mt.permissionLock.Unlock()
return mt.catalog.ListPrivilegeGroups(mt.ctx)
return mt.catalog.ListPrivilegeGroups(ctx)
}
func (mt *MetaTable) OperatePrivilegeGroup(groupName string, privileges []*milvuspb.PrivilegeEntity, operateType milvuspb.OperatePrivilegeGroupType) error {
func (mt *MetaTable) OperatePrivilegeGroup(ctx context.Context, groupName string, privileges []*milvuspb.PrivilegeEntity, operateType milvuspb.OperatePrivilegeGroupType) error {
if funcutil.IsEmptyString(groupName) {
return fmt.Errorf("the privilege group name is empty")
}
@ -1555,14 +1553,14 @@ func (mt *MetaTable) OperatePrivilegeGroup(groupName string, privileges []*milvu
defer mt.permissionLock.Unlock()
// validate input params
definedByUsers, err := mt.IsCustomPrivilegeGroup(groupName)
definedByUsers, err := mt.IsCustomPrivilegeGroup(ctx, groupName)
if err != nil {
return err
}
if !definedByUsers {
return merr.WrapErrParameterInvalidMsg("there is no privilege group name [%s] to operate", groupName)
}
groups, err := mt.catalog.ListPrivilegeGroups(mt.ctx)
groups, err := mt.catalog.ListPrivilegeGroups(ctx)
if err != nil {
return err
}
@ -1581,7 +1579,7 @@ func (mt *MetaTable) OperatePrivilegeGroup(groupName string, privileges []*milvu
}
// merge with current privileges
group, err := mt.catalog.GetPrivilegeGroup(mt.ctx, groupName)
group, err := mt.catalog.GetPrivilegeGroup(ctx, groupName)
if err != nil {
log.Warn("fail to get privilege group", zap.String("privilege_group", groupName), zap.Error(err))
return err
@ -1610,10 +1608,10 @@ func (mt *MetaTable) OperatePrivilegeGroup(groupName string, privileges []*milvu
GroupName: groupName,
Privileges: mergedPrivs,
}
return mt.catalog.SavePrivilegeGroup(mt.ctx, data)
return mt.catalog.SavePrivilegeGroup(ctx, data)
}
func (mt *MetaTable) GetPrivilegeGroupRoles(groupName string) ([]*milvuspb.RoleEntity, error) {
func (mt *MetaTable) GetPrivilegeGroupRoles(ctx context.Context, groupName string) ([]*milvuspb.RoleEntity, error) {
if funcutil.IsEmptyString(groupName) {
return nil, fmt.Errorf("the privilege group name is empty")
}
@ -1621,7 +1619,7 @@ func (mt *MetaTable) GetPrivilegeGroupRoles(groupName string) ([]*milvuspb.RoleE
defer mt.permissionLock.RUnlock()
// get all roles
roles, err := mt.catalog.ListRole(mt.ctx, util.DefaultTenant, nil, false)
roles, err := mt.catalog.ListRole(ctx, util.DefaultTenant, nil, false)
if err != nil {
return nil, err
}
@ -1631,7 +1629,7 @@ func (mt *MetaTable) GetPrivilegeGroupRoles(groupName string) ([]*milvuspb.RoleE
rolesMap := make(map[*milvuspb.RoleEntity]struct{})
for _, role := range roleEntity {
grants, err := mt.catalog.ListGrant(mt.ctx, util.DefaultTenant, &milvuspb.GrantEntity{
grants, err := mt.catalog.ListGrant(ctx, util.DefaultTenant, &milvuspb.GrantEntity{
Role: role,
DbName: util.AnyWord,
})

61
internal/rootcoord/meta_table_test.go
View File

@ -48,7 +48,7 @@ func generateMetaTable(t *testing.T) *MetaTable {
func TestRbacAddCredential(t *testing.T) {
mt := generateMetaTable(t)
err := mt.AddCredential(&internalpb.CredentialInfo{
err := mt.AddCredential(context.TODO(), &internalpb.CredentialInfo{
Username: "user1",
Tenant: util.DefaultTenant,
})
@ -73,7 +73,7 @@ func TestRbacAddCredential(t *testing.T) {
paramtable.Get().Save(Params.ProxyCfg.MaxUserNum.Key, "3")
}
defer paramtable.Get().Reset(Params.ProxyCfg.MaxUserNum.Key)
err := mt.AddCredential(test.info)
err := mt.AddCredential(context.TODO(), test.info)
assert.Error(t, err)
})
}
@ -84,9 +84,9 @@ func TestRbacCreateRole(t *testing.T) {
paramtable.Get().Save(Params.ProxyCfg.MaxRoleNum.Key, "2")
defer paramtable.Get().Reset(Params.ProxyCfg.MaxRoleNum.Key)
err := mt.CreateRole(util.DefaultTenant, &milvuspb.RoleEntity{Name: "role1"})
err := mt.CreateRole(context.TODO(), util.DefaultTenant, &milvuspb.RoleEntity{Name: "role1"})
require.NoError(t, err)
err = mt.CreateRole(util.DefaultTenant, &milvuspb.RoleEntity{Name: "role2"})
err = mt.CreateRole(context.TODO(), util.DefaultTenant, &milvuspb.RoleEntity{Name: "role2"})
require.NoError(t, err)
tests := []struct {
@ -100,12 +100,12 @@ func TestRbacCreateRole(t *testing.T) {
for _, test := range tests {
t.Run(test.description, func(t *testing.T) {
err := mt.CreateRole(util.DefaultTenant, test.inEntity)
err := mt.CreateRole(context.TODO(), util.DefaultTenant, test.inEntity)
assert.Error(t, err)
})
}
t.Run("role has existed", func(t *testing.T) {
err := mt.CreateRole(util.DefaultTenant, &milvuspb.RoleEntity{Name: "role1"})
err := mt.CreateRole(context.TODO(), util.DefaultTenant, &milvuspb.RoleEntity{Name: "role1"})
assert.Error(t, err)
assert.True(t, common.IsIgnorableError(err))
})
@ -119,7 +119,7 @@ func TestRbacCreateRole(t *testing.T) {
mock.Anything,
).Return(nil, errors.New("error mock list role"))
mockMt := &MetaTable{catalog: mockCata}
err := mockMt.CreateRole(util.DefaultTenant, &milvuspb.RoleEntity{Name: "role1"})
err := mockMt.CreateRole(context.TODO(), util.DefaultTenant, &milvuspb.RoleEntity{Name: "role1"})
assert.Error(t, err)
}
}
@ -127,7 +127,7 @@ func TestRbacCreateRole(t *testing.T) {
func TestRbacDropRole(t *testing.T) {
mt := generateMetaTable(t)
err := mt.CreateRole(util.DefaultTenant, &milvuspb.RoleEntity{Name: "role1"})
err := mt.CreateRole(context.TODO(), util.DefaultTenant, &milvuspb.RoleEntity{Name: "role1"})
require.NoError(t, err)
tests := []struct {
@ -141,7 +141,7 @@ func TestRbacDropRole(t *testing.T) {
for _, test := range tests {
t.Run(test.description, func(t *testing.T) {
err := mt.DropRole(util.DefaultTenant, test.roleName)
err := mt.DropRole(context.TODO(), util.DefaultTenant, test.roleName)
assert.NoError(t, err)
})
}
@ -149,7 +149,7 @@ func TestRbacDropRole(t *testing.T) {
func TestRbacOperateRole(t *testing.T) {
mt := generateMetaTable(t)
err := mt.CreateRole(util.DefaultTenant, &milvuspb.RoleEntity{Name: "role1"})
err := mt.CreateRole(context.TODO(), util.DefaultTenant, &milvuspb.RoleEntity{Name: "role1"})
require.NoError(t, err)
tests := []struct {
@ -168,7 +168,7 @@ func TestRbacOperateRole(t *testing.T) {
for _, test := range tests {
t.Run(test.description, func(t *testing.T) {
err := mt.OperateUserRole(util.DefaultTenant, &milvuspb.UserEntity{Name: test.user}, &milvuspb.RoleEntity{Name: test.role}, test.oType)
err := mt.OperateUserRole(context.TODO(), util.DefaultTenant, &milvuspb.UserEntity{Name: test.user}, &milvuspb.RoleEntity{Name: test.role}, test.oType)
assert.Error(t, err)
})
}
@ -185,7 +185,7 @@ func TestRbacSelect(t *testing.T) {
}
for _, role := range roles {
err := mt.CreateRole(util.DefaultTenant, &milvuspb.RoleEntity{Name: role})
err := mt.CreateRole(context.TODO(), util.DefaultTenant, &milvuspb.RoleEntity{Name: role})
require.NoError(t, err)
}
@ -198,6 +198,7 @@ func TestRbacSelect(t *testing.T) {
require.NoError(t, err)
for _, r := range rs {
err := mt.OperateUserRole(
context.TODO(),
util.DefaultTenant,
&milvuspb.UserEntity{Name: user},
&milvuspb.RoleEntity{Name: r},
@ -226,7 +227,7 @@ func TestRbacSelect(t *testing.T) {
for _, test := range tests {
t.Run(test.description, func(t *testing.T) {
res, err := mt.SelectUser(util.DefaultTenant, test.inEntity, test.includeRoleInfo)
res, err := mt.SelectUser(context.TODO(), util.DefaultTenant, test.inEntity, test.includeRoleInfo)
if test.isValid {
assert.NoError(t, err)
@ -264,7 +265,7 @@ func TestRbacSelect(t *testing.T) {
for _, test := range testRoles {
t.Run(test.description, func(t *testing.T) {
res, err := mt.SelectRole(util.DefaultTenant, test.inEntity, test.includeUserInfo)
res, err := mt.SelectRole(context.TODO(), util.DefaultTenant, test.inEntity, test.includeUserInfo)
if test.isValid {
assert.NoError(t, err)
@ -357,7 +358,7 @@ func TestRbacOperatePrivilege(t *testing.T) {
for _, test := range tests {
t.Run(test.description, func(t *testing.T) {
err := mt.OperatePrivilege(util.DefaultTenant, test.entity, test.oType)
err := mt.OperatePrivilege(context.TODO(), util.DefaultTenant, test.entity, test.oType)
assert.Error(t, err)
})
}
@ -372,7 +373,7 @@ func TestRbacOperatePrivilege(t *testing.T) {
ObjectName: "obj_name",
}
err := mt.OperatePrivilege(util.DefaultTenant, &validEntity, milvuspb.OperatePrivilegeType_Grant)
err := mt.OperatePrivilege(context.TODO(), util.DefaultTenant, &validEntity, milvuspb.OperatePrivilegeType_Grant)
assert.NoError(t, err)
}
@ -399,7 +400,7 @@ func TestRbacSelectGrant(t *testing.T) {
for _, test := range tests {
t.Run(test.description, func(t *testing.T) {
entities, err := mt.SelectGrant(util.DefaultTenant, test.entity)
entities, err := mt.SelectGrant(context.TODO(), util.DefaultTenant, test.entity)
if test.isValid {
assert.NoError(t, err)
assert.Equal(t, 0, len(entities))
@ -426,7 +427,7 @@ func TestRbacDropGrant(t *testing.T) {
for _, test := range tests {
t.Run(test.description, func(t *testing.T) {
err := mt.DropGrant(util.DefaultTenant, test.role)
err := mt.DropGrant(context.TODO(), util.DefaultTenant, test.role)
if test.isValid {
assert.NoError(t, err)
} else {
@ -439,11 +440,11 @@ func TestRbacDropGrant(t *testing.T) {
func TestRbacListPolicy(t *testing.T) {
mt := generateMetaTable(t)
policies, err := mt.ListPolicy(util.DefaultTenant)
policies, err := mt.ListPolicy(context.TODO(), util.DefaultTenant)
assert.NoError(t, err)
assert.Empty(t, policies)
userRoles, err := mt.ListUserRole(util.DefaultTenant)
userRoles, err := mt.ListUserRole(context.TODO(), util.DefaultTenant)
assert.NoError(t, err)
assert.Equal(t, 0, len(userRoles))
}
@ -2091,24 +2092,24 @@ func TestMetaTable_PrivilegeGroup(t *testing.T) {
aliases: newNameDb(),
catalog: catalog,
}
err := mt.CreatePrivilegeGroup("pg1")
err := mt.CreatePrivilegeGroup(context.TODO(), "pg1")
assert.Error(t, err)
err = mt.CreatePrivilegeGroup("")
err = mt.CreatePrivilegeGroup(context.TODO(), "")
assert.Error(t, err)
err = mt.CreatePrivilegeGroup("Insert")
err = mt.CreatePrivilegeGroup(context.TODO(), "Insert")
assert.Error(t, err)
err = mt.CreatePrivilegeGroup("pg2")
err = mt.CreatePrivilegeGroup(context.TODO(), "pg2")
assert.NoError(t, err)
err = mt.DropPrivilegeGroup("")
err = mt.DropPrivilegeGroup(context.TODO(), "")
assert.Error(t, err)
err = mt.DropPrivilegeGroup("pg1")
err = mt.DropPrivilegeGroup(context.TODO(), "pg1")
assert.NoError(t, err)
err = mt.OperatePrivilegeGroup("", []*milvuspb.PrivilegeEntity{}, milvuspb.OperatePrivilegeGroupType_AddPrivilegesToGroup)
err = mt.OperatePrivilegeGroup(context.TODO(), "", []*milvuspb.PrivilegeEntity{}, milvuspb.OperatePrivilegeGroupType_AddPrivilegesToGroup)
assert.Error(t, err)
err = mt.OperatePrivilegeGroup("pg3", []*milvuspb.PrivilegeEntity{}, milvuspb.OperatePrivilegeGroupType_AddPrivilegesToGroup)
err = mt.OperatePrivilegeGroup(context.TODO(), "pg3", []*milvuspb.PrivilegeEntity{}, milvuspb.OperatePrivilegeGroupType_AddPrivilegesToGroup)
assert.Error(t, err)
_, err = mt.GetPrivilegeGroupRoles("")
_, err = mt.GetPrivilegeGroupRoles(context.TODO(), "")
assert.Error(t, err)
_, err = mt.ListPrivilegeGroups()
_, err = mt.ListPrivilegeGroups(context.TODO())
assert.NoError(t, err)
}

178
internal/rootcoord/mock_test.go
View File

@ -72,36 +72,36 @@ type mockMetaTable struct {
CreateAliasFunc func(ctx context.Context, dbName string, alias string, collectionName string, ts Timestamp) error
AlterAliasFunc func(ctx context.Context, dbName string, alias string, collectionName string, ts Timestamp) error
DropAliasFunc func(ctx context.Context, dbName string, alias string, ts Timestamp) error
IsAliasFunc func(dbName, name string) bool
IsAliasFunc func(ctx context.Context, dbName, name string) bool
DescribeAliasFunc func(ctx context.Context, dbName, alias string, ts Timestamp) (string, error)
ListAliasesFunc func(ctx context.Context, dbName, collectionName string, ts Timestamp) ([]string, error)
ListAliasesByIDFunc func(collID UniqueID) []string
ListAliasesByIDFunc func(ctx context.Context, collID UniqueID) []string
GetCollectionIDByNameFunc func(name string) (UniqueID, error)
GetPartitionByNameFunc func(collID UniqueID, partitionName string, ts Timestamp) (UniqueID, error)
GetCollectionVirtualChannelsFunc func(colID int64) []string
GetCollectionVirtualChannelsFunc func(ctx context.Context, colID int64) []string
AlterCollectionFunc func(ctx context.Context, oldColl *model.Collection, newColl *model.Collection, ts Timestamp) error
RenameCollectionFunc func(ctx context.Context, oldName string, newName string, ts Timestamp) error
AddCredentialFunc func(credInfo *internalpb.CredentialInfo) error
GetCredentialFunc func(username string) (*internalpb.CredentialInfo, error)
DeleteCredentialFunc func(username string) error
AlterCredentialFunc func(credInfo *internalpb.CredentialInfo) error
ListCredentialUsernamesFunc func() (*milvuspb.ListCredUsersResponse, error)
CreateRoleFunc func(tenant string, entity *milvuspb.RoleEntity) error
DropRoleFunc func(tenant string, roleName string) error
OperateUserRoleFunc func(tenant string, userEntity *milvuspb.UserEntity, roleEntity *milvuspb.RoleEntity, operateType milvuspb.OperateUserRoleType) error
SelectRoleFunc func(tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error)
SelectUserFunc func(tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error)
OperatePrivilegeFunc func(tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error
SelectGrantFunc func(tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error)
DropGrantFunc func(tenant string, role *milvuspb.RoleEntity) error
ListPolicyFunc func(tenant string) ([]string, error)
ListUserRoleFunc func(tenant string) ([]string, error)
AddCredentialFunc func(ctx context.Context, credInfo *internalpb.CredentialInfo) error
GetCredentialFunc func(ctx context.Context, username string) (*internalpb.CredentialInfo, error)
DeleteCredentialFunc func(ctx context.Context, username string) error
AlterCredentialFunc func(ctx context.Context, credInfo *internalpb.CredentialInfo) error
ListCredentialUsernamesFunc func(ctx context.Context) (*milvuspb.ListCredUsersResponse, error)
CreateRoleFunc func(ctx context.Context, tenant string, entity *milvuspb.RoleEntity) error
DropRoleFunc func(ctx context.Context, tenant string, roleName string) error
OperateUserRoleFunc func(ctx context.Context, tenant string, userEntity *milvuspb.UserEntity, roleEntity *milvuspb.RoleEntity, operateType milvuspb.OperateUserRoleType) error
SelectRoleFunc func(ctx context.Context, tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error)
SelectUserFunc func(ctx context.Context, tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error)
OperatePrivilegeFunc func(ctx context.Context, tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error
SelectGrantFunc func(ctx context.Context, tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error)
DropGrantFunc func(ctx context.Context, tenant string, role *milvuspb.RoleEntity) error
ListPolicyFunc func(ctx context.Context, tenant string) ([]string, error)
ListUserRoleFunc func(ctx context.Context, tenant string) ([]string, error)
DescribeDatabaseFunc func(ctx context.Context, dbName string) (*model.Database, error)
CreatePrivilegeGroupFunc func(groupName string) error
DropPrivilegeGroupFunc func(groupName string) error
ListPrivilegeGroupsFunc func() ([]*milvuspb.PrivilegeGroupInfo, error)
OperatePrivilegeGroupFunc func(groupName string, privileges []*milvuspb.PrivilegeEntity, operateType milvuspb.OperatePrivilegeGroupType) error
GetPrivilegeGroupRolesFunc func(groupName string) ([]*milvuspb.RoleEntity, error)
CreatePrivilegeGroupFunc func(ctx context.Context, groupName string) error
DropPrivilegeGroupFunc func(ctx context.Context, groupName string) error
ListPrivilegeGroupsFunc func(ctx context.Context) ([]*milvuspb.PrivilegeGroupInfo, error)
OperatePrivilegeGroupFunc func(ctx context.Context, groupName string, privileges []*milvuspb.PrivilegeEntity, operateType milvuspb.OperatePrivilegeGroupType) error
GetPrivilegeGroupRolesFunc func(ctx context.Context, groupName string) ([]*milvuspb.RoleEntity, error)
}
func (m mockMetaTable) GetDatabaseByName(ctx context.Context, dbName string, ts Timestamp) (*model.Database, error) {
@ -160,8 +160,8 @@ func (m mockMetaTable) DropAlias(ctx context.Context, dbName, alias string, ts T
return m.DropAliasFunc(ctx, dbName, alias, ts)
}
func (m mockMetaTable) IsAlias(dbName, name string) bool {
return m.IsAliasFunc(dbName, name)
func (m mockMetaTable) IsAlias(ctx context.Context, dbName, name string) bool {
return m.IsAliasFunc(ctx, dbName, name)
}
func (m mockMetaTable) DescribeAlias(ctx context.Context, dbName, alias string, ts Timestamp) (string, error) {
@ -172,8 +172,8 @@ func (m mockMetaTable) ListAliases(ctx context.Context, dbName, collectionName s
return m.ListAliasesFunc(ctx, dbName, collectionName, ts)
}
func (m mockMetaTable) ListAliasesByID(collID UniqueID) []string {
return m.ListAliasesByIDFunc(collID)
func (m mockMetaTable) ListAliasesByID(ctx context.Context, collID UniqueID) []string {
return m.ListAliasesByIDFunc(ctx, collID)
}
func (m mockMetaTable) AlterCollection(ctx context.Context, oldColl *model.Collection, newColl *model.Collection, ts Timestamp) error {
@ -192,88 +192,88 @@ func (m mockMetaTable) GetPartitionByName(collID UniqueID, partitionName string,
return m.GetPartitionByNameFunc(collID, partitionName, ts)
}
func (m mockMetaTable) GetCollectionVirtualChannels(colID int64) []string {
return m.GetCollectionVirtualChannelsFunc(colID)
func (m mockMetaTable) GetCollectionVirtualChannels(ctx context.Context, colID int64) []string {
return m.GetCollectionVirtualChannelsFunc(ctx, colID)
}
func (m mockMetaTable) AddCredential(credInfo *internalpb.CredentialInfo) error {
return m.AddCredentialFunc(credInfo)
func (m mockMetaTable) AddCredential(ctx context.Context, credInfo *internalpb.CredentialInfo) error {
return m.AddCredentialFunc(ctx, credInfo)
}
func (m mockMetaTable) GetCredential(username string) (*internalpb.CredentialInfo, error) {
return m.GetCredentialFunc(username)
func (m mockMetaTable) GetCredential(ctx context.Context, username string) (*internalpb.CredentialInfo, error) {
return m.GetCredentialFunc(ctx, username)
}
func (m mockMetaTable) DeleteCredential(username string) error {
return m.DeleteCredentialFunc(username)
func (m mockMetaTable) DeleteCredential(ctx context.Context, username string) error {
return m.DeleteCredentialFunc(ctx, username)
}
func (m mockMetaTable) AlterCredential(credInfo *internalpb.CredentialInfo) error {
return m.AlterCredentialFunc(credInfo)
func (m mockMetaTable) AlterCredential(ctx context.Context, credInfo *internalpb.CredentialInfo) error {
return m.AlterCredentialFunc(ctx, credInfo)
}
func (m mockMetaTable) ListCredentialUsernames() (*milvuspb.ListCredUsersResponse, error) {
return m.ListCredentialUsernamesFunc()
func (m mockMetaTable) ListCredentialUsernames(ctx context.Context) (*milvuspb.ListCredUsersResponse, error) {
return m.ListCredentialUsernamesFunc(ctx)
}
func (m mockMetaTable) CreateRole(tenant string, entity *milvuspb.RoleEntity) error {
return m.CreateRoleFunc(tenant, entity)
func (m mockMetaTable) CreateRole(ctx context.Context, tenant string, entity *milvuspb.RoleEntity) error {
return m.CreateRoleFunc(ctx, tenant, entity)
}
func (m mockMetaTable) DropRole(tenant string, roleName string) error {
return m.DropRoleFunc(tenant, roleName)
func (m mockMetaTable) DropRole(ctx context.Context, tenant string, roleName string) error {
return m.DropRoleFunc(ctx, tenant, roleName)
}
func (m mockMetaTable) OperateUserRole(tenant string, userEntity *milvuspb.UserEntity, roleEntity *milvuspb.RoleEntity, operateType milvuspb.OperateUserRoleType) error {
return m.OperateUserRoleFunc(tenant, userEntity, roleEntity, operateType)
func (m mockMetaTable) OperateUserRole(ctx context.Context, tenant string, userEntity *milvuspb.UserEntity, roleEntity *milvuspb.RoleEntity, operateType milvuspb.OperateUserRoleType) error {
return m.OperateUserRoleFunc(ctx, tenant, userEntity, roleEntity, operateType)
}
func (m mockMetaTable) SelectRole(tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
return m.SelectRoleFunc(tenant, entity, includeUserInfo)
func (m mockMetaTable) SelectRole(ctx context.Context, tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
return m.SelectRoleFunc(ctx, tenant, entity, includeUserInfo)
}
func (m mockMetaTable) SelectUser(tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error) {
return m.SelectUserFunc(tenant, entity, includeRoleInfo)
func (m mockMetaTable) SelectUser(ctx context.Context, tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error) {
return m.SelectUserFunc(ctx, tenant, entity, includeRoleInfo)
}
func (m mockMetaTable) OperatePrivilege(tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error {
return m.OperatePrivilegeFunc(tenant, entity, operateType)
func (m mockMetaTable) OperatePrivilege(ctx context.Context, tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error {
return m.OperatePrivilegeFunc(ctx, tenant, entity, operateType)
}
func (m mockMetaTable) SelectGrant(tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error) {
return m.SelectGrantFunc(tenant, entity)
func (m mockMetaTable) SelectGrant(ctx context.Context, tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error) {
return m.SelectGrantFunc(ctx, tenant, entity)
}
func (m mockMetaTable) DropGrant(tenant string, role *milvuspb.RoleEntity) error {
return m.DropGrantFunc(tenant, role)
func (m mockMetaTable) DropGrant(ctx context.Context, tenant string, role *milvuspb.RoleEntity) error {
return m.DropGrantFunc(ctx, tenant, role)
}
func (m mockMetaTable) ListPolicy(tenant string) ([]string, error) {
return m.ListPolicyFunc(tenant)
func (m mockMetaTable) ListPolicy(ctx context.Context, tenant string) ([]string, error) {
return m.ListPolicyFunc(ctx, tenant)
}
func (m mockMetaTable) ListUserRole(tenant string) ([]string, error) {
return m.ListUserRoleFunc(tenant)
func (m mockMetaTable) ListUserRole(ctx context.Context, tenant string) ([]string, error) {
return m.ListUserRoleFunc(ctx, tenant)
}
func (m mockMetaTable) CreatePrivilegeGroup(groupName string) error {
return m.CreatePrivilegeGroupFunc(groupName)
func (m mockMetaTable) CreatePrivilegeGroup(ctx context.Context, groupName string) error {
return m.CreatePrivilegeGroupFunc(ctx, groupName)
}
func (m mockMetaTable) DropPrivilegeGroup(groupName string) error {
return m.DropPrivilegeGroupFunc(groupName)
func (m mockMetaTable) DropPrivilegeGroup(ctx context.Context, groupName string) error {
return m.DropPrivilegeGroupFunc(ctx, groupName)
}
func (m mockMetaTable) ListPrivilegeGroups() ([]*milvuspb.PrivilegeGroupInfo, error) {
return m.ListPrivilegeGroupsFunc()
func (m mockMetaTable) ListPrivilegeGroups(ctx context.Context) ([]*milvuspb.PrivilegeGroupInfo, error) {
return m.ListPrivilegeGroupsFunc(ctx)
}
func (m mockMetaTable) OperatePrivilegeGroup(groupName string, privileges []*milvuspb.PrivilegeEntity, operateType milvuspb.OperatePrivilegeGroupType) error {
return m.OperatePrivilegeGroupFunc(groupName, privileges, operateType)
func (m mockMetaTable) OperatePrivilegeGroup(ctx context.Context, groupName string, privileges []*milvuspb.PrivilegeEntity, operateType milvuspb.OperatePrivilegeGroupType) error {
return m.OperatePrivilegeGroupFunc(ctx, groupName, privileges, operateType)
}
func (m mockMetaTable) GetPrivilegeGroupRoles(groupName string) ([]*milvuspb.RoleEntity, error) {
return m.GetPrivilegeGroupRolesFunc(groupName)
func (m mockMetaTable) GetPrivilegeGroupRoles(ctx context.Context, groupName string) ([]*milvuspb.RoleEntity, error) {
return m.GetPrivilegeGroupRolesFunc(ctx, groupName)
}
func newMockMetaTable() *mockMetaTable {
@ -498,49 +498,49 @@ func withInvalidMeta() Opt {
meta.DropAliasFunc = func(ctx context.Context, dbName string, alias string, ts Timestamp) error {
return errors.New("error mock DropAlias")
}
meta.AddCredentialFunc = func(credInfo *internalpb.CredentialInfo) error {
meta.AddCredentialFunc = func(ctx context.Context, credInfo *internalpb.CredentialInfo) error {
return errors.New("error mock AddCredential")
}
meta.GetCredentialFunc = func(username string) (*internalpb.CredentialInfo, error) {
meta.GetCredentialFunc = func(ctx context.Context, username string) (*internalpb.CredentialInfo, error) {
return nil, errors.New("error mock GetCredential")
}
meta.DeleteCredentialFunc = func(username string) error {
meta.DeleteCredentialFunc = func(ctx context.Context, username string) error {
return errors.New("error mock DeleteCredential")
}
meta.AlterCredentialFunc = func(credInfo *internalpb.CredentialInfo) error {
meta.AlterCredentialFunc = func(ctx context.Context, credInfo *internalpb.CredentialInfo) error {
return errors.New("error mock AlterCredential")
}
meta.ListCredentialUsernamesFunc = func() (*milvuspb.ListCredUsersResponse, error) {
meta.ListCredentialUsernamesFunc = func(ctx context.Context) (*milvuspb.ListCredUsersResponse, error) {
return nil, errors.New("error mock ListCredentialUsernames")
}
meta.CreateRoleFunc = func(tenant string, entity *milvuspb.RoleEntity) error {
meta.CreateRoleFunc = func(ctx context.Context, tenant string, entity *milvuspb.RoleEntity) error {
return errors.New("error mock CreateRole")
}
meta.DropRoleFunc = func(tenant string, roleName string) error {
meta.DropRoleFunc = func(ctx context.Context, tenant string, roleName string) error {
return errors.New("error mock DropRole")
}
meta.OperateUserRoleFunc = func(tenant string, userEntity *milvuspb.UserEntity, roleEntity *milvuspb.RoleEntity, operateType milvuspb.OperateUserRoleType) error {
meta.OperateUserRoleFunc = func(ctx context.Context, tenant string, userEntity *milvuspb.UserEntity, roleEntity *milvuspb.RoleEntity, operateType milvuspb.OperateUserRoleType) error {
return errors.New("error mock OperateUserRole")
}
meta.SelectUserFunc = func(tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error) {
meta.SelectUserFunc = func(ctx context.Context, tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error) {
return nil, errors.New("error mock SelectUser")
}
meta.SelectRoleFunc = func(tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
meta.SelectRoleFunc = func(ctx context.Context, tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
return nil, errors.New("error mock SelectRole")
}
meta.OperatePrivilegeFunc = func(tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error {
meta.OperatePrivilegeFunc = func(ctx context.Context, tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error {
return errors.New("error mock OperatePrivilege")
}
meta.SelectGrantFunc = func(tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error) {
meta.SelectGrantFunc = func(ctx context.Context, tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error) {
return nil, errors.New("error mock SelectGrant")
}
meta.DropGrantFunc = func(tenant string, role *milvuspb.RoleEntity) error {
meta.DropGrantFunc = func(ctx context.Context, tenant string, role *milvuspb.RoleEntity) error {
return errors.New("error mock DropGrant")
}
meta.ListPolicyFunc = func(tenant string) ([]string, error) {
meta.ListPolicyFunc = func(ctx context.Context, tenant string) ([]string, error) {
return nil, errors.New("error mock ListPolicy")
}
meta.ListUserRoleFunc = func(tenant string) ([]string, error) {
meta.ListUserRoleFunc = func(ctx context.Context, tenant string) ([]string, error) {
return nil, errors.New("error mock ListUserRole")
}
meta.DescribeAliasFunc = func(ctx context.Context, dbName, alias string, ts Timestamp) (string, error) {
@ -552,19 +552,19 @@ func withInvalidMeta() Opt {
meta.DescribeDatabaseFunc = func(ctx context.Context, dbName string) (*model.Database, error) {
return nil, errors.New("error mock DescribeDatabase")
}
meta.CreatePrivilegeGroupFunc = func(groupName string) error {
meta.CreatePrivilegeGroupFunc = func(ctx context.Context, groupName string) error {
return errors.New("error mock CreatePrivilegeGroup")
}
meta.DropPrivilegeGroupFunc = func(groupName string) error {
meta.DropPrivilegeGroupFunc = func(ctx context.Context, groupName string) error {
return errors.New("error mock DropPrivilegeGroup")
}
meta.ListPrivilegeGroupsFunc = func() ([]*milvuspb.PrivilegeGroupInfo, error) {
meta.ListPrivilegeGroupsFunc = func(ctx context.Context) ([]*milvuspb.PrivilegeGroupInfo, error) {
return nil, errors.New("error mock ListPrivilegeGroups")
}
meta.OperatePrivilegeGroupFunc = func(groupName string, privileges []*milvuspb.PrivilegeEntity, operateType milvuspb.OperatePrivilegeGroupType) error {
meta.OperatePrivilegeGroupFunc = func(ctx context.Context, groupName string, privileges []*milvuspb.PrivilegeEntity, operateType milvuspb.OperatePrivilegeGroupType) error {
return errors.New("error mock OperatePrivilegeGroup")
}
meta.GetPrivilegeGroupRolesFunc = func(groupName string) ([]*milvuspb.RoleEntity, error) {
meta.GetPrivilegeGroupRolesFunc = func(ctx context.Context, groupName string) ([]*milvuspb.RoleEntity, error) {
return nil, errors.New("error mock GetPrivilegeGroupRoles")
}
return withMeta(meta)

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

File diff suppressed because it is too large Load Diff

108
internal/rootcoord/root_coord.go
View File

@ -452,7 +452,7 @@ func (c *Core) initInternal() error {
c.scheduler = newScheduler(c.ctx, c.idAllocator, c.tsoAllocator)
c.factory.Init(Params)
chanMap := c.meta.ListCollectionPhysicalChannels()
chanMap := c.meta.ListCollectionPhysicalChannels(c.ctx)
c.chanTimeTick = newTimeTickSync(c.ctx, c.session.ServerID, c.factory, chanMap)
log.Info("create TimeTick sync done")
@ -549,11 +549,11 @@ func (c *Core) Init() error {
}
func (c *Core) initCredentials() error {
credInfo, _ := c.meta.GetCredential(util.UserRoot)
credInfo, _ := c.meta.GetCredential(c.ctx, util.UserRoot)
if credInfo == nil {
log.Debug("RootCoord init user root")
encryptedRootPassword, _ := crypto.PasswordEncrypt(Params.CommonCfg.DefaultRootPassword.GetValue())
err := c.meta.AddCredential(&internalpb.CredentialInfo{Username: util.UserRoot, EncryptedPassword: encryptedRootPassword})
err := c.meta.AddCredential(c.ctx, &internalpb.CredentialInfo{Username: util.UserRoot, EncryptedPassword: encryptedRootPassword})
return err
}
return nil
@ -563,7 +563,7 @@ func (c *Core) initRbac() error {
var err error
// create default roles, including admin, public
for _, role := range util.DefaultRoles {
err = c.meta.CreateRole(util.DefaultTenant, &milvuspb.RoleEntity{Name: role})
err = c.meta.CreateRole(c.ctx, util.DefaultTenant, &milvuspb.RoleEntity{Name: role})
if err != nil && !common.IsIgnorableError(err) {
return errors.Wrap(err, "failed to create role")
}
@ -593,7 +593,7 @@ func (c *Core) initPublicRolePrivilege() error {
var err error
for _, globalPrivilege := range globalPrivileges {
err = c.meta.OperatePrivilege(util.DefaultTenant, &milvuspb.GrantEntity{
err = c.meta.OperatePrivilege(c.ctx, util.DefaultTenant, &milvuspb.GrantEntity{
Role: &milvuspb.RoleEntity{Name: util.RolePublic},
Object: &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Global.String()},
ObjectName: util.AnyWord,
@ -608,7 +608,7 @@ func (c *Core) initPublicRolePrivilege() error {
}
}
for _, collectionPrivilege := range collectionPrivileges {
err = c.meta.OperatePrivilege(util.DefaultTenant, &milvuspb.GrantEntity{
err = c.meta.OperatePrivilege(c.ctx, util.DefaultTenant, &milvuspb.GrantEntity{
Role: &milvuspb.RoleEntity{Name: util.RolePublic},
Object: &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Collection.String()},
ObjectName: util.AnyWord,
@ -672,7 +672,7 @@ func (c *Core) initBuiltinPrivilegeGroups() []*milvuspb.PrivilegeGroupInfo {
func (c *Core) initBuiltinRoles() error {
rolePrivilegesMap := Params.RoleCfg.Roles.GetAsRoleDetails()
for role, privilegesJSON := range rolePrivilegesMap {
err := c.meta.CreateRole(util.DefaultTenant, &milvuspb.RoleEntity{Name: role})
err := c.meta.CreateRole(c.ctx, util.DefaultTenant, &milvuspb.RoleEntity{Name: role})
if err != nil && !common.IsIgnorableError(err) {
log.Error("create a builtin role fail", zap.String("roleName", role), zap.Error(err))
return errors.Wrapf(err, "failed to create a builtin role: %s", role)
@ -680,13 +680,13 @@ func (c *Core) initBuiltinRoles() error {
for _, privilege := range privilegesJSON[util.RoleConfigPrivileges] {
privilegeName := privilege[util.RoleConfigPrivilege]
if !util.IsAnyWord(privilege[util.RoleConfigPrivilege]) {
dbPrivName, err := c.getMetastorePrivilegeName(privilege[util.RoleConfigPrivilege])
dbPrivName, err := c.getMetastorePrivilegeName(c.ctx, privilege[util.RoleConfigPrivilege])
if err != nil {
return errors.Wrapf(err, "failed to get metastore privilege name for: %s", privilege[util.RoleConfigPrivilege])
}
privilegeName = dbPrivName
}
err := c.meta.OperatePrivilege(util.DefaultTenant, &milvuspb.GrantEntity{
err := c.meta.OperatePrivilege(c.ctx, util.DefaultTenant, &milvuspb.GrantEntity{
Role: &milvuspb.RoleEntity{Name: role},
Object: &milvuspb.ObjectEntity{Name: privilege[util.RoleConfigObjectType]},
ObjectName: privilege[util.RoleConfigObjectName],
@ -1666,7 +1666,7 @@ func (c *Core) GetPChannelInfo(ctx context.Context, in *rootcoordpb.GetPChannelI
Status: merr.Status(err),
}, nil
}
return c.meta.GetPChannelInfo(in.GetPchannel()), nil
return c.meta.GetPChannelInfo(ctx, in.GetPchannel()), nil
}
// AllocTimestamp alloc timestamp
@ -2082,7 +2082,7 @@ func (c *Core) CreateCredential(ctx context.Context, credInfo *internalpb.Creden
}
// insert to db
err := c.meta.AddCredential(credInfo)
err := c.meta.AddCredential(ctx, credInfo)
if err != nil {
ctxLog.Warn("CreateCredential save credential failed", zap.Error(err))
metrics.RootCoordDDLReqCounter.WithLabelValues(method, metrics.FailLabel).Inc()
@ -2114,7 +2114,7 @@ func (c *Core) GetCredential(ctx context.Context, in *rootcoordpb.GetCredentialR
return &rootcoordpb.GetCredentialResponse{Status: merr.Status(err)}, nil
}
credInfo, err := c.meta.GetCredential(in.Username)
credInfo, err := c.meta.GetCredential(ctx, in.Username)
if err != nil {
ctxLog.Warn("GetCredential query credential failed", zap.Error(err))
metrics.RootCoordDDLReqCounter.WithLabelValues(method, metrics.FailLabel).Inc()
@ -2144,7 +2144,7 @@ func (c *Core) UpdateCredential(ctx context.Context, credInfo *internalpb.Creden
return merr.Status(err), nil
}
// update data on storage
err := c.meta.AlterCredential(credInfo)
err := c.meta.AlterCredential(ctx, credInfo)
if err != nil {
ctxLog.Warn("UpdateCredential save credential failed", zap.Error(err))
metrics.RootCoordDDLReqCounter.WithLabelValues(method, metrics.FailLabel).Inc()
@ -2183,7 +2183,7 @@ func (c *Core) DeleteCredential(ctx context.Context, in *milvuspb.DeleteCredenti
redoTask := newBaseRedoTask(c.stepExecutor)
redoTask.AddSyncStep(NewSimpleStep("delete credential meta data", func(ctx context.Context) ([]nestedStep, error) {
err := c.meta.DeleteCredential(in.Username)
err := c.meta.DeleteCredential(ctx, in.Username)
if err != nil {
ctxLog.Warn("delete credential meta data failed", zap.Error(err))
}
@ -2234,7 +2234,7 @@ func (c *Core) ListCredUsers(ctx context.Context, in *milvuspb.ListCredUsersRequ
return &milvuspb.ListCredUsersResponse{Status: merr.Status(err)}, nil
}
credInfo, err := c.meta.ListCredentialUsernames()
credInfo, err := c.meta.ListCredentialUsernames(ctx)
if err != nil {
ctxLog.Warn("ListCredUsers query usernames failed", zap.Error(err))
metrics.RootCoordDDLReqCounter.WithLabelValues(method, metrics.FailLabel).Inc()
@ -2269,7 +2269,7 @@ func (c *Core) CreateRole(ctx context.Context, in *milvuspb.CreateRoleRequest) (
}
entity := in.Entity
err := c.meta.CreateRole(util.DefaultTenant, &milvuspb.RoleEntity{Name: entity.Name})
err := c.meta.CreateRole(ctx, util.DefaultTenant, &milvuspb.RoleEntity{Name: entity.Name})
if err != nil {
errMsg := "fail to create role"
ctxLog.Warn(errMsg, zap.Error(err))
@ -2305,14 +2305,14 @@ func (c *Core) DropRole(ctx context.Context, in *milvuspb.DropRoleRequest) (*com
err := merr.WrapErrPrivilegeNotPermitted("the role[%s] is a builtin role, which can't be dropped", in.GetRoleName())
return merr.Status(err), nil
}
if _, err := c.meta.SelectRole(util.DefaultTenant, &milvuspb.RoleEntity{Name: in.RoleName}, false); err != nil {
if _, err := c.meta.SelectRole(ctx, util.DefaultTenant, &milvuspb.RoleEntity{Name: in.RoleName}, false); err != nil {
errMsg := "not found the role, maybe the role isn't existed or internal system error"
ctxLog.Warn(errMsg, zap.Error(err))
return merr.StatusWithErrorCode(errors.New(errMsg), commonpb.ErrorCode_DropRoleFailure), nil
}
if !in.ForceDrop {
grantEntities, err := c.meta.SelectGrant(util.DefaultTenant, &milvuspb.GrantEntity{
grantEntities, err := c.meta.SelectGrant(ctx, util.DefaultTenant, &milvuspb.GrantEntity{
Role: &milvuspb.RoleEntity{Name: in.RoleName},
})
if len(grantEntities) != 0 {
@ -2323,7 +2323,7 @@ func (c *Core) DropRole(ctx context.Context, in *milvuspb.DropRoleRequest) (*com
}
redoTask := newBaseRedoTask(c.stepExecutor)
redoTask.AddSyncStep(NewSimpleStep("drop role meta data", func(ctx context.Context) ([]nestedStep, error) {
err := c.meta.DropRole(util.DefaultTenant, in.RoleName)
err := c.meta.DropRole(ctx, util.DefaultTenant, in.RoleName)
if err != nil {
ctxLog.Warn("drop role mata data failed", zap.Error(err))
}
@ -2333,7 +2333,7 @@ func (c *Core) DropRole(ctx context.Context, in *milvuspb.DropRoleRequest) (*com
if !in.ForceDrop {
return nil, nil
}
err := c.meta.DropGrant(util.DefaultTenant, &milvuspb.RoleEntity{Name: in.RoleName})
err := c.meta.DropGrant(ctx, util.DefaultTenant, &milvuspb.RoleEntity{Name: in.RoleName})
if err != nil {
ctxLog.Warn("drop the privilege list failed for the role", zap.Error(err))
}
@ -2380,13 +2380,13 @@ func (c *Core) OperateUserRole(ctx context.Context, in *milvuspb.OperateUserRole
return merr.Status(err), nil
}
if _, err := c.meta.SelectRole(util.DefaultTenant, &milvuspb.RoleEntity{Name: in.RoleName}, false); err != nil {
if _, err := c.meta.SelectRole(ctx, util.DefaultTenant, &milvuspb.RoleEntity{Name: in.RoleName}, false); err != nil {
errMsg := "not found the role, maybe the role isn't existed or internal system error"
ctxLog.Warn(errMsg, zap.Error(err))
return merr.StatusWithErrorCode(errors.New(errMsg), commonpb.ErrorCode_OperateUserRoleFailure), nil
}
if in.Type != milvuspb.OperateUserRoleType_RemoveUserFromRole {
if _, err := c.meta.SelectUser(util.DefaultTenant, &milvuspb.UserEntity{Name: in.Username}, false); err != nil {
if _, err := c.meta.SelectUser(ctx, util.DefaultTenant, &milvuspb.UserEntity{Name: in.Username}, false); err != nil {
errMsg := "not found the user, maybe the user isn't existed or internal system error"
ctxLog.Warn(errMsg, zap.Error(err))
return merr.StatusWithErrorCode(errors.New(errMsg), commonpb.ErrorCode_OperateUserRoleFailure), nil
@ -2395,7 +2395,7 @@ func (c *Core) OperateUserRole(ctx context.Context, in *milvuspb.OperateUserRole
redoTask := newBaseRedoTask(c.stepExecutor)
redoTask.AddSyncStep(NewSimpleStep("operate user role meta data", func(ctx context.Context) ([]nestedStep, error) {
err := c.meta.OperateUserRole(util.DefaultTenant, &milvuspb.UserEntity{Name: in.Username}, &milvuspb.RoleEntity{Name: in.RoleName}, in.Type)
err := c.meta.OperateUserRole(ctx, util.DefaultTenant, &milvuspb.UserEntity{Name: in.Username}, &milvuspb.RoleEntity{Name: in.RoleName}, in.Type)
if err != nil && !common.IsIgnorableError(err) {
log.Warn("operate user role mata data failed", zap.Error(err))
return nil, err
@ -2452,7 +2452,7 @@ func (c *Core) SelectRole(ctx context.Context, in *milvuspb.SelectRoleRequest) (
}
if in.Role != nil {
if _, err := c.meta.SelectRole(util.DefaultTenant, &milvuspb.RoleEntity{Name: in.Role.Name}, false); err != nil {
if _, err := c.meta.SelectRole(ctx, util.DefaultTenant, &milvuspb.RoleEntity{Name: in.Role.Name}, false); err != nil {
if errors.Is(err, merr.ErrIoKeyNotFound) {
return &milvuspb.SelectRoleResponse{
Status: merr.Success(),
@ -2465,7 +2465,7 @@ func (c *Core) SelectRole(ctx context.Context, in *milvuspb.SelectRoleRequest) (
}, nil
}
}
roleResults, err := c.meta.SelectRole(util.DefaultTenant, in.Role, in.IncludeUserInfo)
roleResults, err := c.meta.SelectRole(ctx, util.DefaultTenant, in.Role, in.IncludeUserInfo)
if err != nil {
errMsg := "fail to select the role"
ctxLog.Warn(errMsg, zap.Error(err))
@ -2499,7 +2499,7 @@ func (c *Core) SelectUser(ctx context.Context, in *milvuspb.SelectUserRequest) (
}
if in.User != nil {
if _, err := c.meta.SelectUser(util.DefaultTenant, &milvuspb.UserEntity{Name: in.User.Name}, false); err != nil {
if _, err := c.meta.SelectUser(ctx, util.DefaultTenant, &milvuspb.UserEntity{Name: in.User.Name}, false); err != nil {
if errors.Is(err, merr.ErrIoKeyNotFound) {
return &milvuspb.SelectUserResponse{
Status: merr.Success(),
@ -2512,7 +2512,7 @@ func (c *Core) SelectUser(ctx context.Context, in *milvuspb.SelectUserRequest) (
}, nil
}
}
userResults, err := c.meta.SelectUser(util.DefaultTenant, in.User, in.IncludeRoleInfo)
userResults, err := c.meta.SelectUser(ctx, util.DefaultTenant, in.User, in.IncludeRoleInfo)
if err != nil {
errMsg := "fail to select the user"
ctxLog.Warn(errMsg, zap.Error(err))
@ -2537,7 +2537,7 @@ func (c *Core) isValidRole(entity *milvuspb.RoleEntity) error {
if entity.Name == "" {
return errors.New("the name in the role entity is empty")
}
if _, err := c.meta.SelectRole(util.DefaultTenant, &milvuspb.RoleEntity{Name: entity.Name}, false); err != nil {
if _, err := c.meta.SelectRole(c.ctx, util.DefaultTenant, &milvuspb.RoleEntity{Name: entity.Name}, false); err != nil {
log.Warn("fail to select the role", zap.String("role_name", entity.Name), zap.Error(err))
return errors.New("not found the role, maybe the role isn't existed or internal system error")
}
@ -2554,14 +2554,14 @@ func (c *Core) isValidObject(entity *milvuspb.ObjectEntity) error {
return nil
}
func (c *Core) isValidGrantor(entity *milvuspb.GrantorEntity, object string) error {
func (c *Core) isValidGrantor(ctx context.Context, entity *milvuspb.GrantorEntity, object string) error {
if entity == nil {
return errors.New("the grantor entity is nil")
}
if entity.User == nil || entity.User.Name == "" {
return errors.New("the user entity in the grantor entity is nil or empty")
}
if _, err := c.meta.SelectUser(util.DefaultTenant, &milvuspb.UserEntity{Name: entity.User.Name}, false); err != nil {
if _, err := c.meta.SelectUser(ctx, util.DefaultTenant, &milvuspb.UserEntity{Name: entity.User.Name}, false); err != nil {
log.Warn("fail to select the user", zap.String("username", entity.User.Name), zap.Error(err))
return errors.New("not found the user, maybe the user isn't existed or internal system error")
}
@ -2584,7 +2584,7 @@ func (c *Core) isValidGrantor(entity *milvuspb.GrantorEntity, object string) err
}
}
// check if it is a custom privilege group
customPrivGroup, err := c.meta.IsCustomPrivilegeGroup(entity.Privilege.Name)
customPrivGroup, err := c.meta.IsCustomPrivilegeGroup(ctx, entity.Privilege.Name)
if err != nil {
return err
}
@ -2629,7 +2629,7 @@ func (c *Core) OperatePrivilege(ctx context.Context, in *milvuspb.OperatePrivile
ctxLog.Warn("", zap.Error(err))
return merr.StatusWithErrorCode(err, commonpb.ErrorCode_OperatePrivilegeFailure), nil
}
if err := c.isValidGrantor(in.Entity.Grantor, in.Entity.Object.Name); err != nil {
if err := c.isValidGrantor(ctx, in.Entity.Grantor, in.Entity.Object.Name); err != nil {
ctxLog.Error("", zap.Error(err))
return merr.StatusWithErrorCode(err, commonpb.ErrorCode_OperatePrivilegeFailure), nil
}
@ -2645,14 +2645,14 @@ func (c *Core) OperatePrivilege(ctx context.Context, in *milvuspb.OperatePrivile
redoTask.AddSyncStep(NewSimpleStep("operate privilege meta data", func(ctx context.Context) ([]nestedStep, error) {
if !util.IsAnyWord(privName) {
// set up privilege name for metastore
dbPrivName, err := c.getMetastorePrivilegeName(privName)
dbPrivName, err := c.getMetastorePrivilegeName(ctx, privName)
if err != nil {
return nil, err
}
in.Entity.Grantor.Privilege.Name = dbPrivName
}
err := c.meta.OperatePrivilege(util.DefaultTenant, in.Entity, in.Type)
err := c.meta.OperatePrivilege(ctx, util.DefaultTenant, in.Entity, in.Type)
if err != nil && !common.IsIgnorableError(err) {
log.Warn("fail to operate the privilege", zap.Any("in", in), zap.Error(err))
return nil, err
@ -2674,7 +2674,7 @@ func (c *Core) OperatePrivilege(ctx context.Context, in *milvuspb.OperatePrivile
}
grants := []*milvuspb.GrantEntity{in.Entity}
allGroups, err := c.meta.ListPrivilegeGroups()
allGroups, err := c.meta.ListPrivilegeGroups(ctx)
allGroups = append(allGroups, c.initBuiltinPrivilegeGroups()...)
if err != nil {
return nil, err
@ -2682,7 +2682,7 @@ func (c *Core) OperatePrivilege(ctx context.Context, in *milvuspb.OperatePrivile
groups := lo.SliceToMap(allGroups, func(group *milvuspb.PrivilegeGroupInfo) (string, []*milvuspb.PrivilegeEntity) {
return group.GroupName, group.Privileges
})
expandGrants, err := c.expandPrivilegeGroups(grants, groups)
expandGrants, err := c.expandPrivilegeGroups(ctx, grants, groups)
if err != nil {
return nil, err
}
@ -2709,13 +2709,13 @@ func (c *Core) OperatePrivilege(ctx context.Context, in *milvuspb.OperatePrivile
return merr.Success(), nil
}
func (c *Core) getMetastorePrivilegeName(privName string) (string, error) {
func (c *Core) getMetastorePrivilegeName(ctx context.Context, privName string) (string, error) {
// if it is built-in privilege, return the privilege name directly
if util.IsPrivilegeNameDefined(privName) {
return util.PrivilegeNameForMetastore(privName), nil
}
// return the privilege group name if it is a custom privilege group
customGroup, err := c.meta.IsCustomPrivilegeGroup(privName)
customGroup, err := c.meta.IsCustomPrivilegeGroup(ctx, privName)
if err != nil {
return "", err
}
@ -2764,7 +2764,7 @@ func (c *Core) SelectGrant(ctx context.Context, in *milvuspb.SelectGrantRequest)
}
}
grantEntities, err := c.meta.SelectGrant(util.DefaultTenant, in.Entity)
grantEntities, err := c.meta.SelectGrant(ctx, util.DefaultTenant, in.Entity)
if errors.Is(err, merr.ErrIoKeyNotFound) {
return &milvuspb.SelectGrantResponse{
Status: merr.Success(),
@ -2801,7 +2801,7 @@ func (c *Core) ListPolicy(ctx context.Context, in *internalpb.ListPolicyRequest)
}, nil
}
policies, err := c.meta.ListPolicy(util.DefaultTenant)
policies, err := c.meta.ListPolicy(ctx, util.DefaultTenant)
if err != nil {
errMsg := "fail to list policy"
ctxLog.Warn(errMsg, zap.Error(err))
@ -2809,7 +2809,7 @@ func (c *Core) ListPolicy(ctx context.Context, in *internalpb.ListPolicyRequest)
Status: merr.StatusWithErrorCode(errors.New(errMsg), commonpb.ErrorCode_ListPolicyFailure),
}, nil
}
userRoles, err := c.meta.ListUserRole(util.DefaultTenant)
userRoles, err := c.meta.ListUserRole(ctx, util.DefaultTenant)
if err != nil {
errMsg := "fail to list user-role"
ctxLog.Warn(errMsg, zap.Any("in", in), zap.Error(err))
@ -2817,7 +2817,7 @@ func (c *Core) ListPolicy(ctx context.Context, in *internalpb.ListPolicyRequest)
Status: merr.StatusWithErrorCode(errors.New(errMsg), commonpb.ErrorCode_ListPolicyFailure),
}, nil
}
privGroups, err := c.meta.ListPrivilegeGroups()
privGroups, err := c.meta.ListPrivilegeGroups(ctx)
if err != nil {
errMsg := "fail to list privilege groups"
ctxLog.Warn(errMsg, zap.Error(err))
@ -3046,7 +3046,7 @@ func (c *Core) CreatePrivilegeGroup(ctx context.Context, in *milvuspb.CreatePriv
return merr.Status(err), nil
}
if err := c.meta.CreatePrivilegeGroup(in.GroupName); err != nil {
if err := c.meta.CreatePrivilegeGroup(ctx, in.GroupName); err != nil {
ctxLog.Warn("fail to create privilege group", zap.Error(err))
return merr.Status(err), nil
}
@ -3069,7 +3069,7 @@ func (c *Core) DropPrivilegeGroup(ctx context.Context, in *milvuspb.DropPrivileg
return merr.Status(err), nil
}
if err := c.meta.DropPrivilegeGroup(in.GroupName); err != nil {
if err := c.meta.DropPrivilegeGroup(ctx, in.GroupName); err != nil {
ctxLog.Warn("fail to drop privilege group", zap.Error(err))
return merr.Status(err), nil
}
@ -3094,7 +3094,7 @@ func (c *Core) ListPrivilegeGroups(ctx context.Context, in *milvuspb.ListPrivile
}, nil
}
privGroups, err := c.meta.ListPrivilegeGroups()
privGroups, err := c.meta.ListPrivilegeGroups(ctx)
if err != nil {
ctxLog.Warn("fail to list privilege group", zap.Error(err))
return &milvuspb.ListPrivilegeGroupsResponse{
@ -3124,7 +3124,7 @@ func (c *Core) OperatePrivilegeGroup(ctx context.Context, in *milvuspb.OperatePr
redoTask := newBaseRedoTask(c.stepExecutor)
redoTask.AddSyncStep(NewSimpleStep("operate privilege group", func(ctx context.Context) ([]nestedStep, error) {
groups, err := c.meta.ListPrivilegeGroups()
groups, err := c.meta.ListPrivilegeGroups(ctx)
if err != nil && !common.IsIgnorableError(err) {
log.Warn("fail to list privilege groups", zap.Error(err))
return nil, err
@ -3134,7 +3134,7 @@ func (c *Core) OperatePrivilegeGroup(ctx context.Context, in *milvuspb.OperatePr
})
// get roles granted to the group
roles, err := c.meta.GetPrivilegeGroupRoles(in.GroupName)
roles, err := c.meta.GetPrivilegeGroupRoles(ctx, in.GroupName)
if err != nil {
return nil, err
}
@ -3170,18 +3170,18 @@ func (c *Core) OperatePrivilegeGroup(ctx context.Context, in *milvuspb.OperatePr
a.DbName == b.DbName
}
for _, role := range roles {
grants, err := c.meta.SelectGrant(util.DefaultTenant, &milvuspb.GrantEntity{
grants, err := c.meta.SelectGrant(ctx, util.DefaultTenant, &milvuspb.GrantEntity{
Role: role,
DbName: util.AnyWord,
})
if err != nil {
return nil, err
}
currGrants, err := c.expandPrivilegeGroups(grants, currGroups)
currGrants, err := c.expandPrivilegeGroups(ctx, grants, currGroups)
if err != nil {
return nil, err
}
newGrants, err := c.expandPrivilegeGroups(grants, newGroups)
newGrants, err := c.expandPrivilegeGroups(ctx, grants, newGroups)
if err != nil {
return nil, err
}
@ -3227,7 +3227,7 @@ func (c *Core) OperatePrivilegeGroup(ctx context.Context, in *milvuspb.OperatePr
}))
redoTask.AddSyncStep(NewSimpleStep("operate privilege group meta data", func(ctx context.Context) ([]nestedStep, error) {
err := c.meta.OperatePrivilegeGroup(in.GroupName, in.Privileges, in.Type)
err := c.meta.OperatePrivilegeGroup(ctx, in.GroupName, in.Privileges, in.Type)
if err != nil && !common.IsIgnorableError(err) {
log.Warn("fail to operate privilege group", zap.Error(err))
}
@ -3248,12 +3248,12 @@ func (c *Core) OperatePrivilegeGroup(ctx context.Context, in *milvuspb.OperatePr
return merr.Success(), nil
}
func (c *Core) expandPrivilegeGroups(grants []*milvuspb.GrantEntity, groups map[string][]*milvuspb.PrivilegeEntity) ([]*milvuspb.GrantEntity, error) {
func (c *Core) expandPrivilegeGroups(ctx context.Context, grants []*milvuspb.GrantEntity, groups map[string][]*milvuspb.PrivilegeEntity) ([]*milvuspb.GrantEntity, error) {
newGrants := []*milvuspb.GrantEntity{}
for _, grant := range grants {
privName := grant.Grantor.Privilege.Name
if privGroup, exists := groups[privName]; !exists {
metaName, err := c.getMetastorePrivilegeName(privName)
metaName, err := c.getMetastorePrivilegeName(ctx, privName)
if err != nil {
return nil, err
}
@ -3271,7 +3271,7 @@ func (c *Core) expandPrivilegeGroups(grants []*milvuspb.GrantEntity, groups map[
})
} else {
for _, priv := range privGroup {
metaName, err := c.getMetastorePrivilegeName(priv.Name)
metaName, err := c.getMetastorePrivilegeName(ctx, priv.Name)
if err != nil {
return nil, err
}

60
internal/rootcoord/root_coord_test.go
View File

@ -1682,19 +1682,19 @@ func TestRootCoord_RBACError(t *testing.T) {
})
t.Run("operate user role failed", func(t *testing.T) {
mockMeta := c.meta.(*mockMetaTable)
mockMeta.SelectRoleFunc = func(tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
mockMeta.SelectRoleFunc = func(ctx context.Context, tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
return nil, nil
}
mockMeta.SelectUserFunc = func(tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error) {
mockMeta.SelectUserFunc = func(ctx context.Context, tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error) {
return nil, nil
}
resp, err := c.OperateUserRole(ctx, &milvuspb.OperateUserRoleRequest{RoleName: "foo", Username: "bar", Type: milvuspb.OperateUserRoleType_AddUserToRole})
assert.NoError(t, err)
assert.NotEqual(t, commonpb.ErrorCode_Success, resp.ErrorCode)
mockMeta.SelectRoleFunc = func(tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
mockMeta.SelectRoleFunc = func(ctx context.Context, tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
return nil, errors.New("mock error")
}
mockMeta.SelectUserFunc = func(tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error) {
mockMeta.SelectUserFunc = func(ctx context.Context, tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error) {
return nil, errors.New("mock error")
}
})
@ -1745,10 +1745,10 @@ func TestRootCoord_RBACError(t *testing.T) {
}
mockMeta := c.meta.(*mockMetaTable)
mockMeta.SelectRoleFunc = func(tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
mockMeta.SelectRoleFunc = func(ctx context.Context, tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
return nil, nil
}
mockMeta.ListPrivilegeGroupsFunc = func() ([]*milvuspb.PrivilegeGroupInfo, error) {
mockMeta.ListPrivilegeGroupsFunc = func(ctx context.Context) ([]*milvuspb.PrivilegeGroupInfo, error) {
return nil, nil
}
{
@ -1765,7 +1765,7 @@ func TestRootCoord_RBACError(t *testing.T) {
assert.NotEqual(t, commonpb.ErrorCode_Success, resp.ErrorCode)
}
mockMeta.SelectUserFunc = func(tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error) {
mockMeta.SelectUserFunc = func(ctx context.Context, tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error) {
return nil, nil
}
resp, err := c.OperatePrivilege(ctx, &milvuspb.OperatePrivilegeRequest{Entity: &milvuspb.GrantEntity{
@ -1779,23 +1779,23 @@ func TestRootCoord_RBACError(t *testing.T) {
}, Type: milvuspb.OperatePrivilegeType_Grant})
assert.NoError(t, err)
assert.NotEqual(t, commonpb.ErrorCode_Success, resp.ErrorCode)
mockMeta.SelectRoleFunc = func(tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
mockMeta.SelectRoleFunc = func(ctx context.Context, tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
return nil, errors.New("mock error")
}
mockMeta.SelectUserFunc = func(tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error) {
mockMeta.SelectUserFunc = func(ctx context.Context, tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error) {
return nil, errors.New("mock error")
}
})
t.Run("operate privilege group failed", func(t *testing.T) {
mockMeta := c.meta.(*mockMetaTable)
mockMeta.ListPrivilegeGroupsFunc = func() ([]*milvuspb.PrivilegeGroupInfo, error) {
mockMeta.ListPrivilegeGroupsFunc = func(ctx context.Context) ([]*milvuspb.PrivilegeGroupInfo, error) {
return nil, errors.New("mock error")
}
mockMeta.CreatePrivilegeGroupFunc = func(groupName string) error {
mockMeta.CreatePrivilegeGroupFunc = func(ctx context.Context, groupName string) error {
return errors.New("mock error")
}
mockMeta.GetPrivilegeGroupRolesFunc = func(groupName string) ([]*milvuspb.RoleEntity, error) {
mockMeta.GetPrivilegeGroupRolesFunc = func(ctx context.Context, groupName string) ([]*milvuspb.RoleEntity, error) {
return nil, errors.New("mock error")
}
{
@ -1832,7 +1832,7 @@ func TestRootCoord_RBACError(t *testing.T) {
assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
}
mockMeta := c.meta.(*mockMetaTable)
mockMeta.SelectRoleFunc = func(tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
mockMeta.SelectRoleFunc = func(ctx context.Context, tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
return nil, nil
}
{
@ -1845,21 +1845,21 @@ func TestRootCoord_RBACError(t *testing.T) {
assert.NoError(t, err)
assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
}
mockMeta.SelectRoleFunc = func(tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
mockMeta.SelectRoleFunc = func(ctx context.Context, tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
return nil, errors.New("mock error")
}
})
t.Run("select grant success", func(t *testing.T) {
mockMeta := c.meta.(*mockMetaTable)
mockMeta.SelectRoleFunc = func(tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
mockMeta.SelectRoleFunc = func(ctx context.Context, tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
return []*milvuspb.RoleResult{
{
Role: &milvuspb.RoleEntity{Name: "foo"},
},
}, nil
}
mockMeta.SelectGrantFunc = func(tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error) {
mockMeta.SelectGrantFunc = func(ctx context.Context, tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error) {
return []*milvuspb.GrantEntity{
{
Role: &milvuspb.RoleEntity{Name: "foo"},
@ -1874,11 +1874,11 @@ func TestRootCoord_RBACError(t *testing.T) {
assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
}
mockMeta.SelectRoleFunc = func(tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
mockMeta.SelectRoleFunc = func(ctx context.Context, tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
return nil, errors.New("mock error")
}
mockMeta.SelectGrantFunc = func(tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error) {
mockMeta.SelectGrantFunc = func(ctx context.Context, tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error) {
return nil, errors.New("mock error")
}
})
@ -1889,13 +1889,13 @@ func TestRootCoord_RBACError(t *testing.T) {
assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
mockMeta := c.meta.(*mockMetaTable)
mockMeta.ListPolicyFunc = func(tenant string) ([]string, error) {
mockMeta.ListPolicyFunc = func(ctx context.Context, tenant string) ([]string, error) {
return []string{}, nil
}
resp, err = c.ListPolicy(ctx, &internalpb.ListPolicyRequest{})
assert.NoError(t, err)
assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
mockMeta.ListPolicyFunc = func(tenant string) ([]string, error) {
mockMeta.ListPolicyFunc = func(ctx context.Context, tenant string) ([]string, error) {
return []string{}, errors.New("mock error")
}
})
@ -1909,13 +1909,13 @@ func TestRootCoord_BuiltinRoles(t *testing.T) {
t.Run("init builtin roles success", func(t *testing.T) {
c := newTestCore(withHealthyCode(), withInvalidMeta())
mockMeta := c.meta.(*mockMetaTable)
mockMeta.CreateRoleFunc = func(tenant string, entity *milvuspb.RoleEntity) error {
mockMeta.CreateRoleFunc = func(ctx context.Context, tenant string, entity *milvuspb.RoleEntity) error {
return nil
}
mockMeta.OperatePrivilegeFunc = func(tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error {
mockMeta.OperatePrivilegeFunc = func(ctx context.Context, tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error {
return nil
}
mockMeta.ListPrivilegeGroupsFunc = func() ([]*milvuspb.PrivilegeGroupInfo, error) {
mockMeta.ListPrivilegeGroupsFunc = func(ctx context.Context) ([]*milvuspb.PrivilegeGroupInfo, error) {
return nil, nil
}
err := c.initBuiltinRoles()
@ -1929,7 +1929,7 @@ func TestRootCoord_BuiltinRoles(t *testing.T) {
t.Run("init builtin roles fail to create role", func(t *testing.T) {
c := newTestCore(withHealthyCode(), withInvalidMeta())
mockMeta := c.meta.(*mockMetaTable)
mockMeta.CreateRoleFunc = func(tenant string, entity *milvuspb.RoleEntity) error {
mockMeta.CreateRoleFunc = func(ctx context.Context, tenant string, entity *milvuspb.RoleEntity) error {
return merr.ErrPrivilegeNotPermitted
}
err := c.initBuiltinRoles()
@ -1938,10 +1938,10 @@ func TestRootCoord_BuiltinRoles(t *testing.T) {
t.Run("init builtin roles fail to operate privileg", func(t *testing.T) {
c := newTestCore(withHealthyCode(), withInvalidMeta())
mockMeta := c.meta.(*mockMetaTable)
mockMeta.CreateRoleFunc = func(tenant string, entity *milvuspb.RoleEntity) error {
mockMeta.CreateRoleFunc = func(ctx context.Context, tenant string, entity *milvuspb.RoleEntity) error {
return nil
}
mockMeta.OperatePrivilegeFunc = func(tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error {
mockMeta.OperatePrivilegeFunc = func(ctx context.Context, tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error {
return merr.ErrPrivilegeNotPermitted
}
err := c.initBuiltinRoles()
@ -1974,8 +1974,8 @@ func TestCore_InitRBAC(t *testing.T) {
t.Run("init default role and public role privilege", func(t *testing.T) {
meta := mockrootcoord.NewIMetaTable(t)
c := newTestCore(withHealthyCode(), withMeta(meta))
meta.EXPECT().CreateRole(mock.Anything, mock.Anything).Return(nil).Twice()
meta.EXPECT().OperatePrivilege(mock.Anything, mock.Anything, mock.Anything).Return(nil).Twice()
meta.EXPECT().CreateRole(mock.Anything, mock.Anything, mock.Anything).Return(nil).Twice()
meta.EXPECT().OperatePrivilege(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil).Twice()
Params.Save(Params.RoleCfg.Enabled.Key, "false")
Params.Save(Params.ProxyCfg.EnablePublicPrivilege.Key, "true")
@ -1993,8 +1993,8 @@ func TestCore_InitRBAC(t *testing.T) {
builtinRoles := `{"db_admin": {"privileges": [{"object_type": "Global", "object_name": "*", "privilege": "CreateCollection", "db_name": "*"}]}}`
meta := mockrootcoord.NewIMetaTable(t)
c := newTestCore(withHealthyCode(), withMeta(meta))
meta.EXPECT().CreateRole(mock.Anything, mock.Anything).Return(nil).Times(3)
meta.EXPECT().OperatePrivilege(mock.Anything, mock.Anything, mock.Anything).Return(nil).Once()
meta.EXPECT().CreateRole(mock.Anything, mock.Anything, mock.Anything).Return(nil).Times(3)
meta.EXPECT().OperatePrivilege(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil).Once()
Params.Save(Params.RoleCfg.Enabled.Key, "true")
Params.Save(Params.RoleCfg.Roles.Key, builtinRoles)

4
internal/rootcoord/show_collection_task.go
View File

@ -65,7 +65,7 @@ func (t *showCollectionTask) Execute(ctx context.Context) error {
privilegeColls.Insert(util.AnyWord)
return privilegeColls, nil
}
userRoles, err := t.core.meta.SelectUser("", &milvuspb.UserEntity{
userRoles, err := t.core.meta.SelectUser(ctx, "", &milvuspb.UserEntity{
Name: curUser,
}, true)
if err != nil {
@ -79,7 +79,7 @@ func (t *showCollectionTask) Execute(ctx context.Context) error {
privilegeColls.Insert(util.AnyWord)
return privilegeColls, nil
}
entities, err := t.core.meta.SelectGrant("", &milvuspb.GrantEntity{
entities, err := t.core.meta.SelectGrant(ctx, "", &milvuspb.GrantEntity{
Role: role,
DbName: t.Req.GetDbName(),
})

22
internal/rootcoord/show_collection_task_test.go
View File

@ -169,7 +169,7 @@ func TestShowCollectionsAuth(t *testing.T) {
meta := mockrootcoord.NewIMetaTable(t)
core := newTestCore(withMeta(meta))
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything).
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
Return(nil, errors.New("mock error: select user")).Once()
task := &showCollectionTask{
@ -189,7 +189,7 @@ func TestShowCollectionsAuth(t *testing.T) {
meta := mockrootcoord.NewIMetaTable(t)
core := newTestCore(withMeta(meta))
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything).
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
Return([]*milvuspb.UserResult{}, nil).Once()
task := &showCollectionTask{
@ -210,7 +210,7 @@ func TestShowCollectionsAuth(t *testing.T) {
meta := mockrootcoord.NewIMetaTable(t)
core := newTestCore(withMeta(meta))
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything).
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
Return([]*milvuspb.UserResult{
{
User: &milvuspb.UserEntity{
@ -250,7 +250,7 @@ func TestShowCollectionsAuth(t *testing.T) {
meta := mockrootcoord.NewIMetaTable(t)
core := newTestCore(withMeta(meta))
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything).
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
Return([]*milvuspb.UserResult{
{
User: &milvuspb.UserEntity{
@ -263,7 +263,7 @@ func TestShowCollectionsAuth(t *testing.T) {
},
},
}, nil).Once()
meta.EXPECT().SelectGrant(mock.Anything, mock.Anything).Return(nil, errors.New("mock error: select grant")).Once()
meta.EXPECT().SelectGrant(mock.Anything, mock.Anything, mock.Anything).Return(nil, errors.New("mock error: select grant")).Once()
task := &showCollectionTask{
baseTask: newBaseTask(context.Background(), core),
@ -281,7 +281,7 @@ func TestShowCollectionsAuth(t *testing.T) {
meta := mockrootcoord.NewIMetaTable(t)
core := newTestCore(withMeta(meta))
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything).
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
Return([]*milvuspb.UserResult{
{
User: &milvuspb.UserEntity{
@ -294,7 +294,7 @@ func TestShowCollectionsAuth(t *testing.T) {
},
},
}, nil).Once()
meta.EXPECT().SelectGrant(mock.Anything, mock.Anything).Return([]*milvuspb.GrantEntity{
meta.EXPECT().SelectGrant(mock.Anything, mock.Anything, mock.Anything).Return([]*milvuspb.GrantEntity{
{
Object: &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Global.String()},
Grantor: &milvuspb.GrantorEntity{
@ -331,7 +331,7 @@ func TestShowCollectionsAuth(t *testing.T) {
meta := mockrootcoord.NewIMetaTable(t)
core := newTestCore(withMeta(meta))
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything).
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
Return([]*milvuspb.UserResult{
{
User: &milvuspb.UserEntity{
@ -344,7 +344,7 @@ func TestShowCollectionsAuth(t *testing.T) {
},
},
}, nil).Once()
meta.EXPECT().SelectGrant(mock.Anything, mock.Anything).Return([]*milvuspb.GrantEntity{
meta.EXPECT().SelectGrant(mock.Anything, mock.Anything, mock.Anything).Return([]*milvuspb.GrantEntity{
{
Object: &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Collection.String()},
ObjectName: util.AnyWord,
@ -376,7 +376,7 @@ func TestShowCollectionsAuth(t *testing.T) {
meta := mockrootcoord.NewIMetaTable(t)
core := newTestCore(withMeta(meta))
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything).
meta.EXPECT().SelectUser(mock.Anything, mock.Anything, mock.Anything, mock.Anything).
Return([]*milvuspb.UserResult{
{
User: &milvuspb.UserEntity{
@ -389,7 +389,7 @@ func TestShowCollectionsAuth(t *testing.T) {
},
},
}, nil).Once()
meta.EXPECT().SelectGrant(mock.Anything, mock.Anything).Return([]*milvuspb.GrantEntity{
meta.EXPECT().SelectGrant(mock.Anything, mock.Anything, mock.Anything).Return([]*milvuspb.GrantEntity{
{
Object: &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Collection.String()},
ObjectName: "a",