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 // 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{ redoTask.AddSyncStep(&expireCacheStep{
baseStep: baseStep{core: a.core}, baseStep: baseStep{core: a.core},
dbName: a.Req.GetDbName(), 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,
mock.Anything, mock.Anything,
).Return(errors.New("err")) ).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)) core := newTestCore(withValidProxyManager(), withMeta(meta))
task := &alterCollectionTask{ task := &alterCollectionTask{
@ -122,7 +122,7 @@ func Test_alterCollectionTask_Execute(t *testing.T) {
mock.Anything, mock.Anything,
mock.Anything, mock.Anything,
).Return(nil) ).Return(nil)
meta.On("ListAliasesByID", mock.Anything).Return([]string{}) meta.On("ListAliasesByID", mock.Anything, mock.Anything).Return([]string{})
broker := newMockBroker() broker := newMockBroker()
broker.BroadcastAlteredCollectionFunc = func(ctx context.Context, req *milvuspb.AlterCollectionRequest) error { 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,
mock.Anything, mock.Anything,
).Return(nil) ).Return(nil)
meta.On("ListAliasesByID", mock.Anything).Return([]string{}) meta.On("ListAliasesByID", mock.Anything, mock.Anything).Return([]string{})
broker := newMockBroker() broker := newMockBroker()
broker.BroadcastAlteredCollectionFunc = func(ctx context.Context, req *milvuspb.AlterCollectionRequest) error { 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,
mock.Anything, mock.Anything,
).Return(nil) ).Return(nil)
meta.On("ListAliasesByID", mock.Anything).Return([]string{}) meta.On("ListAliasesByID", mock.Anything, mock.Anything).Return([]string{})
broker := newMockBroker() broker := newMockBroker()
broker.BroadcastAlteredCollectionFunc = func(ctx context.Context, req *milvuspb.AlterCollectionRequest) error { 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 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()) db, err := t.core.meta.GetDatabaseByID(ctx, coll.DBID, t.GetTs())
if err != nil { if err != nil {
return err return err

1
internal/rootcoord/describe_collection_task_test.go
View File

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

8
internal/rootcoord/drop_collection_task.go
View File

@ -37,18 +37,18 @@ type dropCollectionTask struct {
Req *milvuspb.DropCollectionRequest 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 { if err := CheckMsgType(t.Req.GetBase().GetMsgType(), commonpb.MsgType_DropCollection); err != nil {
return err 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 fmt.Errorf("cannot drop the collection via alias = %s", t.Req.CollectionName)
} }
return nil return nil
} }
func (t *dropCollectionTask) Prepare(ctx context.Context) error { func (t *dropCollectionTask) Prepare(ctx context.Context) error {
return t.validate() return t.validate(ctx)
} }
func (t *dropCollectionTask) Execute(ctx context.Context) error { 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. // 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() 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", meta.On("IsAlias",
mock.Anything, mock.Anything,
mock.Anything, mock.Anything,
mock.Anything,
).Return(true) ).Return(true)
core := newTestCore(withMeta(meta)) core := newTestCore(withMeta(meta))
@ -72,6 +73,7 @@ func Test_dropCollectionTask_Prepare(t *testing.T) {
meta.On("IsAlias", meta.On("IsAlias",
mock.Anything, mock.Anything,
mock.Anything, mock.Anything,
mock.Anything,
).Return(false) ).Return(false)
core := newTestCore(withMeta(meta)) core := newTestCore(withMeta(meta))
@ -129,6 +131,7 @@ func Test_dropCollectionTask_Execute(t *testing.T) {
mock.Anything, mock.Anything,
).Return(coll.Clone(), nil) ).Return(coll.Clone(), nil)
meta.On("ListAliasesByID", meta.On("ListAliasesByID",
mock.Anything,
mock.AnythingOfType("int64"), mock.AnythingOfType("int64"),
).Return([]string{}) ).Return([]string{})
@ -163,6 +166,7 @@ func Test_dropCollectionTask_Execute(t *testing.T) {
).Return(errors.New("error mock ChangeCollectionState")) ).Return(errors.New("error mock ChangeCollectionState"))
meta.On("ListAliasesByID", meta.On("ListAliasesByID",
mock.Anything, mock.Anything,
mock.Anything,
).Return([]string{}) ).Return([]string{})
core := newTestCore(withValidProxyManager(), withMeta(meta)) core := newTestCore(withValidProxyManager(), withMeta(meta))
@ -207,6 +211,7 @@ func Test_dropCollectionTask_Execute(t *testing.T) {
).Return(nil) ).Return(nil)
meta.On("ListAliasesByID", meta.On("ListAliasesByID",
mock.Anything, mock.Anything,
mock.Anything,
).Return([]string{}) ).Return([]string{})
removeCollectionMetaCalled := false removeCollectionMetaCalled := false
removeCollectionMetaChan := make(chan struct{}, 1) 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) privilegeDBs.Insert(util.AnyWord)
return privilegeDBs, nil return privilegeDBs, nil
} }
userRoles, err := t.core.meta.SelectUser("", &milvuspb.UserEntity{ userRoles, err := t.core.meta.SelectUser(ctx, "", &milvuspb.UserEntity{
Name: curUser, Name: curUser,
}, true) }, true)
if err != nil { if err != nil {
@ -72,7 +72,7 @@ func (t *listDatabaseTask) Execute(ctx context.Context) error {
privilegeDBs.Insert(util.AnyWord) privilegeDBs.Insert(util.AnyWord)
return privilegeDBs, nil return privilegeDBs, nil
} }
entities, err := t.core.meta.SelectGrant("", &milvuspb.GrantEntity{ entities, err := t.core.meta.SelectGrant(ctx, "", &milvuspb.GrantEntity{
Role: role, Role: role,
DbName: util.AnyWord, 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 // 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() Return(nil, errors.New("mock select user error")).Once()
ctx := GetContext(context.Background(), "foo:root") ctx := GetContext(context.Background(), "foo:root")
task := getTask() task := getTask()
@ -142,7 +142,7 @@ func Test_ListDBTask(t *testing.T) {
{ {
// select role, empty result // 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() Return([]*milvuspb.UserResult{}, nil).Once()
ctx := GetContext(context.Background(), "foo:root") ctx := GetContext(context.Background(), "foo:root")
task := getTask() task := getTask()
@ -153,7 +153,7 @@ func Test_ListDBTask(t *testing.T) {
{ {
// select role, the user is added to admin role // 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{ Return([]*milvuspb.UserResult{
{ {
User: &milvuspb.UserEntity{ User: &milvuspb.UserEntity{
@ -176,7 +176,7 @@ func Test_ListDBTask(t *testing.T) {
{ {
// select grant fail // 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{ Return([]*milvuspb.UserResult{
{ {
User: &milvuspb.UserEntity{ User: &milvuspb.UserEntity{
@ -189,7 +189,7 @@ func Test_ListDBTask(t *testing.T) {
}, },
}, },
}, nil).Once() }, 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() Return(nil, errors.New("mock select grant error")).Once()
ctx := GetContext(context.Background(), "foo:root") ctx := GetContext(context.Background(), "foo:root")
task := getTask() task := getTask()
@ -199,7 +199,7 @@ func Test_ListDBTask(t *testing.T) {
{ {
// normal user // 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{ Return([]*milvuspb.UserResult{
{ {
User: &milvuspb.UserEntity{ User: &milvuspb.UserEntity{
@ -220,7 +220,7 @@ func Test_ListDBTask(t *testing.T) {
Name: "default", Name: "default",
}, },
}, nil).Once() }, nil).Once()
meta.EXPECT().SelectGrant(mock.Anything, mock.Anything). meta.EXPECT().SelectGrant(mock.Anything, mock.Anything, mock.Anything).
Return([]*milvuspb.GrantEntity{ Return([]*milvuspb.GrantEntity{
{ {
DbName: "fooDB", DbName: "fooDB",
@ -236,7 +236,7 @@ func Test_ListDBTask(t *testing.T) {
{ {
// normal user with any db privilege // 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{ Return([]*milvuspb.UserResult{
{ {
User: &milvuspb.UserEntity{ User: &milvuspb.UserEntity{
@ -257,7 +257,7 @@ func Test_ListDBTask(t *testing.T) {
Name: "default", Name: "default",
}, },
}, nil).Once() }, nil).Once()
meta.EXPECT().SelectGrant(mock.Anything, mock.Anything). meta.EXPECT().SelectGrant(mock.Anything, mock.Anything, mock.Anything).
Return([]*milvuspb.GrantEntity{ Return([]*milvuspb.GrantEntity{
{ {
DbName: "*", 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) GetCollectionByIDWithMaxTs(ctx context.Context, collectionID UniqueID) (*model.Collection, error)
ListCollections(ctx context.Context, dbName string, ts Timestamp, onlyAvail bool) ([]*model.Collection, error) ListCollections(ctx context.Context, dbName string, ts Timestamp, onlyAvail bool) ([]*model.Collection, error)
ListAllAvailCollections(ctx context.Context) map[int64][]int64 ListAllAvailCollections(ctx context.Context) map[int64][]int64
ListCollectionPhysicalChannels() map[typeutil.UniqueID][]string ListCollectionPhysicalChannels(ctx context.Context) map[typeutil.UniqueID][]string
GetCollectionVirtualChannels(colID int64) []string GetCollectionVirtualChannels(ctx context.Context, colID int64) []string
GetPChannelInfo(pchannel string) *rootcoordpb.GetPChannelInfoResponse GetPChannelInfo(ctx context.Context, pchannel string) *rootcoordpb.GetPChannelInfoResponse
AddPartition(ctx context.Context, partition *model.Partition) error AddPartition(ctx context.Context, partition *model.Partition) error
ChangePartitionState(ctx context.Context, collectionID UniqueID, partitionID UniqueID, state pb.PartitionState, ts Timestamp) 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 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 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. // 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 IsAlias(ctx context.Context, db, name string) bool
ListAliasesByID(collID UniqueID) []string ListAliasesByID(ctx context.Context, collID UniqueID) []string
// TODO: better to accept ctx. AddCredential(ctx context.Context, credInfo *internalpb.CredentialInfo) error
AddCredential(credInfo *internalpb.CredentialInfo) error GetCredential(ctx context.Context, username string) (*internalpb.CredentialInfo, error)
GetCredential(username string) (*internalpb.CredentialInfo, error) DeleteCredential(ctx context.Context, username string) error
DeleteCredential(username string) error AlterCredential(ctx context.Context, credInfo *internalpb.CredentialInfo) error
AlterCredential(credInfo *internalpb.CredentialInfo) error ListCredentialUsernames(ctx context.Context) (*milvuspb.ListCredUsersResponse, error)
ListCredentialUsernames() (*milvuspb.ListCredUsersResponse, error)
// TODO: better to accept ctx. CreateRole(ctx context.Context, tenant string, entity *milvuspb.RoleEntity) error
CreateRole(tenant string, entity *milvuspb.RoleEntity) error DropRole(ctx context.Context, tenant string, roleName string) error
DropRole(tenant string, roleName string) error OperateUserRole(ctx context.Context, tenant string, userEntity *milvuspb.UserEntity, roleEntity *milvuspb.RoleEntity, operateType milvuspb.OperateUserRoleType) error
OperateUserRole(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)
SelectRole(tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) SelectUser(ctx context.Context, tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error)
SelectUser(tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error) OperatePrivilege(ctx context.Context, tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error
OperatePrivilege(tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error SelectGrant(ctx context.Context, tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error)
SelectGrant(tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error) DropGrant(ctx context.Context, tenant string, role *milvuspb.RoleEntity) error
DropGrant(tenant string, role *milvuspb.RoleEntity) error ListPolicy(ctx context.Context, tenant string) ([]string, error)
ListPolicy(tenant string) ([]string, error) ListUserRole(ctx context.Context, tenant string) ([]string, error)
ListUserRole(tenant string) ([]string, error)
BackupRBAC(ctx context.Context, tenant string) (*milvuspb.RBACMeta, error) BackupRBAC(ctx context.Context, tenant string) (*milvuspb.RBACMeta, error)
RestoreRBAC(ctx context.Context, tenant string, meta *milvuspb.RBACMeta) error RestoreRBAC(ctx context.Context, tenant string, meta *milvuspb.RBACMeta) error
IsCustomPrivilegeGroup(groupName string) (bool, error) IsCustomPrivilegeGroup(ctx context.Context, groupName string) (bool, error)
CreatePrivilegeGroup(groupName string) error CreatePrivilegeGroup(ctx context.Context, groupName string) error
DropPrivilegeGroup(groupName string) error DropPrivilegeGroup(ctx context.Context, groupName string) error
ListPrivilegeGroups() ([]*milvuspb.PrivilegeGroupInfo, error) ListPrivilegeGroups(ctx context.Context) ([]*milvuspb.PrivilegeGroupInfo, error)
OperatePrivilegeGroup(groupName string, privileges []*milvuspb.PrivilegeEntity, operateType milvuspb.OperatePrivilegeGroupType) error OperatePrivilegeGroup(ctx context.Context, groupName string, privileges []*milvuspb.PrivilegeEntity, operateType milvuspb.OperatePrivilegeGroupType) error
GetPrivilegeGroupRoles(groupName string) ([]*milvuspb.RoleEntity, error) GetPrivilegeGroupRoles(ctx context.Context, groupName string) ([]*milvuspb.RoleEntity, error)
} }
// MetaTable is a persistent meta set of all databases, collections and partitions. // 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. // 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() mt.ddLock.RLock()
defer mt.ddLock.RUnlock() 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. // 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() mt.ddLock.RLock()
defer mt.ddLock.RUnlock() defer mt.ddLock.RUnlock()
for id, collInfo := range mt.collID2Meta { for id, collInfo := range mt.collID2Meta {
@ -850,7 +848,7 @@ func (mt *MetaTable) GetCollectionVirtualChannels(colID int64) []string {
} }
// GetPChannelInfo returns infos on pchannel. // 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() mt.ddLock.RLock()
defer mt.ddLock.RUnlock() defer mt.ddLock.RUnlock()
resp := &rootcoordpb.GetPChannelInfoResponse{ resp := &rootcoordpb.GetPChannelInfoResponse{
@ -1199,7 +1197,7 @@ func (mt *MetaTable) ListAliases(ctx context.Context, dbName string, collectionN
return aliases, nil 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() mt.ddLock.RLock()
defer mt.ddLock.RUnlock() defer mt.ddLock.RUnlock()
@ -1218,7 +1216,7 @@ func (mt *MetaTable) listAliasesByID(collID UniqueID) []string {
return ret return ret
} }
func (mt *MetaTable) ListAliasesByID(collID UniqueID) []string { func (mt *MetaTable) ListAliasesByID(ctx context.Context, collID UniqueID) []string {
mt.ddLock.RLock() mt.ddLock.RLock()
defer mt.ddLock.RUnlock() defer mt.ddLock.RUnlock()
@ -1226,14 +1224,14 @@ func (mt *MetaTable) ListAliasesByID(collID UniqueID) []string {
} }
// AddCredential add credential // 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 == "" { if credInfo.Username == "" {
return fmt.Errorf("username is empty") return fmt.Errorf("username is empty")
} }
mt.permissionLock.Lock() mt.permissionLock.Lock()
defer mt.permissionLock.Unlock() defer mt.permissionLock.Unlock()
usernames, err := mt.catalog.ListCredentials(mt.ctx) usernames, err := mt.catalog.ListCredentials(ctx)
if err != nil { if err != nil {
return err return err
} }
@ -1243,7 +1241,7 @@ func (mt *MetaTable) AddCredential(credInfo *internalpb.CredentialInfo) error {
return errors.New(errMsg) 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) return fmt.Errorf("user already exists: %s", credInfo.Username)
} }
@ -1251,11 +1249,11 @@ func (mt *MetaTable) AddCredential(credInfo *internalpb.CredentialInfo) error {
Username: credInfo.Username, Username: credInfo.Username,
EncryptedPassword: credInfo.EncryptedPassword, EncryptedPassword: credInfo.EncryptedPassword,
} }
return mt.catalog.CreateCredential(mt.ctx, credential) return mt.catalog.CreateCredential(ctx, credential)
} }
// AlterCredential update 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 == "" { if credInfo.Username == "" {
return fmt.Errorf("username is empty") return fmt.Errorf("username is empty")
} }
@ -1267,32 +1265,32 @@ func (mt *MetaTable) AlterCredential(credInfo *internalpb.CredentialInfo) error
Username: credInfo.Username, Username: credInfo.Username,
EncryptedPassword: credInfo.EncryptedPassword, EncryptedPassword: credInfo.EncryptedPassword,
} }
return mt.catalog.AlterCredential(mt.ctx, credential) return mt.catalog.AlterCredential(ctx, credential)
} }
// GetCredential get credential by username // 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() mt.permissionLock.RLock()
defer mt.permissionLock.RUnlock() defer mt.permissionLock.RUnlock()
credential, err := mt.catalog.GetCredential(mt.ctx, username) credential, err := mt.catalog.GetCredential(ctx, username)
return model.MarshalCredentialModel(credential), err return model.MarshalCredentialModel(credential), err
} }
// DeleteCredential delete credential // DeleteCredential delete credential
func (mt *MetaTable) DeleteCredential(username string) error { func (mt *MetaTable) DeleteCredential(ctx context.Context, username string) error {
mt.permissionLock.Lock() mt.permissionLock.Lock()
defer mt.permissionLock.Unlock() defer mt.permissionLock.Unlock()
return mt.catalog.DropCredential(mt.ctx, username) return mt.catalog.DropCredential(ctx, username)
} }
// ListCredentialUsernames list credential usernames // ListCredentialUsernames list credential usernames
func (mt *MetaTable) ListCredentialUsernames() (*milvuspb.ListCredUsersResponse, error) { func (mt *MetaTable) ListCredentialUsernames(ctx context.Context) (*milvuspb.ListCredUsersResponse, error) {
mt.permissionLock.RLock() mt.permissionLock.RLock()
defer mt.permissionLock.RUnlock() defer mt.permissionLock.RUnlock()
usernames, err := mt.catalog.ListCredentials(mt.ctx) usernames, err := mt.catalog.ListCredentials(ctx)
if err != nil { if err != nil {
return nil, fmt.Errorf("list credential usernames err:%w", err) return nil, fmt.Errorf("list credential usernames err:%w", err)
} }
@ -1300,14 +1298,14 @@ func (mt *MetaTable) ListCredentialUsernames() (*milvuspb.ListCredUsersResponse,
} }
// CreateRole create role // 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) { if funcutil.IsEmptyString(entity.Name) {
return fmt.Errorf("the role name in the role info is empty") return fmt.Errorf("the role name in the role info is empty")
} }
mt.permissionLock.Lock() mt.permissionLock.Lock()
defer mt.permissionLock.Unlock() 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 { if err != nil {
log.Warn("fail to list roles", zap.Error(err)) log.Warn("fail to list roles", zap.Error(err))
return err return err
@ -1324,19 +1322,19 @@ func (mt *MetaTable) CreateRole(tenant string, entity *milvuspb.RoleEntity) erro
return errors.New(errMsg) return errors.New(errMsg)
} }
return mt.catalog.CreateRole(mt.ctx, tenant, entity) return mt.catalog.CreateRole(ctx, tenant, entity)
} }
// DropRole drop role info // 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() mt.permissionLock.Lock()
defer mt.permissionLock.Unlock() 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 // 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) { if funcutil.IsEmptyString(userEntity.Name) {
return fmt.Errorf("username in the user entity is empty") 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() mt.permissionLock.Lock()
defer mt.permissionLock.Unlock() 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. // SelectRole select role.
// Enter the role condition by the entity param. And this param is nil, which means selecting all roles. // 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. // 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() mt.permissionLock.RLock()
defer mt.permissionLock.RUnlock() defer mt.permissionLock.RUnlock()
return mt.catalog.ListRole(mt.ctx, tenant, entity, includeUserInfo) return mt.catalog.ListRole(ctx, tenant, entity, includeUserInfo)
} }
// SelectUser select user. // SelectUser select user.
// Enter the user condition by the entity param. And this param is nil, which means selecting all users. // 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. // 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() mt.permissionLock.RLock()
defer mt.permissionLock.RUnlock() 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 // 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) { if funcutil.IsEmptyString(entity.ObjectName) {
return fmt.Errorf("the object name in the grant entity is empty") 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() mt.permissionLock.Lock()
defer mt.permissionLock.Unlock() defer mt.permissionLock.Unlock()
return mt.catalog.AlterGrant(mt.ctx, tenant, entity, operateType) return mt.catalog.AlterGrant(ctx, tenant, entity, operateType)
} }
// SelectGrant select grant // SelectGrant select grant
// The principal entity MUST be not empty in the grant entity // 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. // 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 var entities []*milvuspb.GrantEntity
if entity == nil { if entity == nil {
return entities, fmt.Errorf("the grant entity is 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() mt.permissionLock.RLock()
defer mt.permissionLock.RUnlock() 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) { if role == nil || funcutil.IsEmptyString(role.Name) {
return fmt.Errorf("the role entity is invalid when dropping the grant") return fmt.Errorf("the role entity is invalid when dropping the grant")
} }
mt.permissionLock.Lock() mt.permissionLock.Lock()
defer mt.permissionLock.Unlock() 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() mt.permissionLock.RLock()
defer mt.permissionLock.RUnlock() 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() mt.permissionLock.RLock()
defer mt.permissionLock.RUnlock() 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) { func (mt *MetaTable) BackupRBAC(ctx context.Context, tenant string) (*milvuspb.RBACMeta, error) {
mt.permissionLock.RLock() mt.permissionLock.RLock()
defer mt.permissionLock.RUnlock() 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 { func (mt *MetaTable) RestoreRBAC(ctx context.Context, tenant string, meta *milvuspb.RBACMeta) error {
mt.permissionLock.Lock() mt.permissionLock.Lock()
defer mt.permissionLock.Unlock() 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 // check if the privielge group name is defined by users
func (mt *MetaTable) IsCustomPrivilegeGroup(groupName string) (bool, error) { func (mt *MetaTable) IsCustomPrivilegeGroup(ctx context.Context, groupName string) (bool, error) {
privGroups, err := mt.catalog.ListPrivilegeGroups(mt.ctx) privGroups, err := mt.catalog.ListPrivilegeGroups(ctx)
if err != nil { if err != nil {
return false, err return false, err
} }
@ -1477,14 +1475,14 @@ func (mt *MetaTable) IsCustomPrivilegeGroup(groupName string) (bool, error) {
return false, nil return false, nil
} }
func (mt *MetaTable) CreatePrivilegeGroup(groupName string) error { func (mt *MetaTable) CreatePrivilegeGroup(ctx context.Context, groupName string) error {
if funcutil.IsEmptyString(groupName) { if funcutil.IsEmptyString(groupName) {
return fmt.Errorf("the privilege group name is empty") return fmt.Errorf("the privilege group name is empty")
} }
mt.permissionLock.Lock() mt.permissionLock.Lock()
defer mt.permissionLock.Unlock() defer mt.permissionLock.Unlock()
definedByUsers, err := mt.IsCustomPrivilegeGroup(groupName) definedByUsers, err := mt.IsCustomPrivilegeGroup(ctx, groupName)
if err != nil { if err != nil {
return err return err
} }
@ -1498,17 +1496,17 @@ func (mt *MetaTable) CreatePrivilegeGroup(groupName string) error {
GroupName: groupName, GroupName: groupName,
Privileges: make([]*milvuspb.PrivilegeEntity, 0), 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) { if funcutil.IsEmptyString(groupName) {
return fmt.Errorf("the privilege group name is empty") return fmt.Errorf("the privilege group name is empty")
} }
mt.permissionLock.Lock() mt.permissionLock.Lock()
defer mt.permissionLock.Unlock() defer mt.permissionLock.Unlock()
definedByUsers, err := mt.IsCustomPrivilegeGroup(groupName) definedByUsers, err := mt.IsCustomPrivilegeGroup(ctx, groupName)
if err != nil { if err != nil {
return err return err
} }
@ -1516,7 +1514,7 @@ func (mt *MetaTable) DropPrivilegeGroup(groupName string) error {
return nil return nil
} }
// check if the group is used by any role // 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 { if err != nil {
return err return err
} }
@ -1524,7 +1522,7 @@ func (mt *MetaTable) DropPrivilegeGroup(groupName string) error {
return entity.GetRole() return entity.GetRole()
}) })
for _, role := range roleEntity { 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, Role: role,
DbName: util.AnyWord, 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() mt.permissionLock.Lock()
defer mt.permissionLock.Unlock() 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) { if funcutil.IsEmptyString(groupName) {
return fmt.Errorf("the privilege group name is empty") 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() defer mt.permissionLock.Unlock()
// validate input params // validate input params
definedByUsers, err := mt.IsCustomPrivilegeGroup(groupName) definedByUsers, err := mt.IsCustomPrivilegeGroup(ctx, groupName)
if err != nil { if err != nil {
return err return err
} }
if !definedByUsers { if !definedByUsers {
return merr.WrapErrParameterInvalidMsg("there is no privilege group name [%s] to operate", groupName) 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 { if err != nil {
return err return err
} }
@ -1581,7 +1579,7 @@ func (mt *MetaTable) OperatePrivilegeGroup(groupName string, privileges []*milvu
} }
// merge with current privileges // merge with current privileges
group, err := mt.catalog.GetPrivilegeGroup(mt.ctx, groupName) group, err := mt.catalog.GetPrivilegeGroup(ctx, groupName)
if err != nil { if err != nil {
log.Warn("fail to get privilege group", zap.String("privilege_group", groupName), zap.Error(err)) log.Warn("fail to get privilege group", zap.String("privilege_group", groupName), zap.Error(err))
return err return err
@ -1610,10 +1608,10 @@ func (mt *MetaTable) OperatePrivilegeGroup(groupName string, privileges []*milvu
GroupName: groupName, GroupName: groupName,
Privileges: mergedPrivs, 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) { if funcutil.IsEmptyString(groupName) {
return nil, fmt.Errorf("the privilege group name is empty") 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() defer mt.permissionLock.RUnlock()
// get all roles // 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 { if err != nil {
return nil, err return nil, err
} }
@ -1631,7 +1629,7 @@ func (mt *MetaTable) GetPrivilegeGroupRoles(groupName string) ([]*milvuspb.RoleE
rolesMap := make(map[*milvuspb.RoleEntity]struct{}) rolesMap := make(map[*milvuspb.RoleEntity]struct{})
for _, role := range roleEntity { 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, Role: role,
DbName: util.AnyWord, 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) { func TestRbacAddCredential(t *testing.T) {
mt := generateMetaTable(t) mt := generateMetaTable(t)
err := mt.AddCredential(&internalpb.CredentialInfo{ err := mt.AddCredential(context.TODO(), &internalpb.CredentialInfo{
Username: "user1", Username: "user1",
Tenant: util.DefaultTenant, Tenant: util.DefaultTenant,
}) })
@ -73,7 +73,7 @@ func TestRbacAddCredential(t *testing.T) {
paramtable.Get().Save(Params.ProxyCfg.MaxUserNum.Key, "3") paramtable.Get().Save(Params.ProxyCfg.MaxUserNum.Key, "3")
} }
defer paramtable.Get().Reset(Params.ProxyCfg.MaxUserNum.Key) defer paramtable.Get().Reset(Params.ProxyCfg.MaxUserNum.Key)
err := mt.AddCredential(test.info) err := mt.AddCredential(context.TODO(), test.info)
assert.Error(t, err) assert.Error(t, err)
}) })
} }
@ -84,9 +84,9 @@ func TestRbacCreateRole(t *testing.T) {
paramtable.Get().Save(Params.ProxyCfg.MaxRoleNum.Key, "2") paramtable.Get().Save(Params.ProxyCfg.MaxRoleNum.Key, "2")
defer paramtable.Get().Reset(Params.ProxyCfg.MaxRoleNum.Key) 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) 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) require.NoError(t, err)
tests := []struct { tests := []struct {
@ -100,12 +100,12 @@ func TestRbacCreateRole(t *testing.T) {
for _, test := range tests { for _, test := range tests {
t.Run(test.description, func(t *testing.T) { 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) assert.Error(t, err)
}) })
} }
t.Run("role has existed", func(t *testing.T) { 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.Error(t, err)
assert.True(t, common.IsIgnorableError(err)) assert.True(t, common.IsIgnorableError(err))
}) })
@ -119,7 +119,7 @@ func TestRbacCreateRole(t *testing.T) {
mock.Anything, mock.Anything,
).Return(nil, errors.New("error mock list role")) ).Return(nil, errors.New("error mock list role"))
mockMt := &MetaTable{catalog: mockCata} 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) assert.Error(t, err)
} }
} }
@ -127,7 +127,7 @@ func TestRbacCreateRole(t *testing.T) {
func TestRbacDropRole(t *testing.T) { func TestRbacDropRole(t *testing.T) {
mt := generateMetaTable(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) require.NoError(t, err)
tests := []struct { tests := []struct {
@ -141,7 +141,7 @@ func TestRbacDropRole(t *testing.T) {
for _, test := range tests { for _, test := range tests {
t.Run(test.description, func(t *testing.T) { 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) assert.NoError(t, err)
}) })
} }
@ -149,7 +149,7 @@ func TestRbacDropRole(t *testing.T) {
func TestRbacOperateRole(t *testing.T) { func TestRbacOperateRole(t *testing.T) {
mt := generateMetaTable(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) require.NoError(t, err)
tests := []struct { tests := []struct {
@ -168,7 +168,7 @@ func TestRbacOperateRole(t *testing.T) {
for _, test := range tests { for _, test := range tests {
t.Run(test.description, func(t *testing.T) { 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) assert.Error(t, err)
}) })
} }
@ -185,7 +185,7 @@ func TestRbacSelect(t *testing.T) {
} }
for _, role := range roles { 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) require.NoError(t, err)
} }
@ -198,6 +198,7 @@ func TestRbacSelect(t *testing.T) {
require.NoError(t, err) require.NoError(t, err)
for _, r := range rs { for _, r := range rs {
err := mt.OperateUserRole( err := mt.OperateUserRole(
context.TODO(),
util.DefaultTenant, util.DefaultTenant,
&milvuspb.UserEntity{Name: user}, &milvuspb.UserEntity{Name: user},
&milvuspb.RoleEntity{Name: r}, &milvuspb.RoleEntity{Name: r},
@ -226,7 +227,7 @@ func TestRbacSelect(t *testing.T) {
for _, test := range tests { for _, test := range tests {
t.Run(test.description, func(t *testing.T) { 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 { if test.isValid {
assert.NoError(t, err) assert.NoError(t, err)
@ -264,7 +265,7 @@ func TestRbacSelect(t *testing.T) {
for _, test := range testRoles { for _, test := range testRoles {
t.Run(test.description, func(t *testing.T) { 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 { if test.isValid {
assert.NoError(t, err) assert.NoError(t, err)
@ -357,7 +358,7 @@ func TestRbacOperatePrivilege(t *testing.T) {
for _, test := range tests { for _, test := range tests {
t.Run(test.description, func(t *testing.T) { 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) assert.Error(t, err)
}) })
} }
@ -372,7 +373,7 @@ func TestRbacOperatePrivilege(t *testing.T) {
ObjectName: "obj_name", 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) assert.NoError(t, err)
} }
@ -399,7 +400,7 @@ func TestRbacSelectGrant(t *testing.T) {
for _, test := range tests { for _, test := range tests {
t.Run(test.description, func(t *testing.T) { 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 { if test.isValid {
assert.NoError(t, err) assert.NoError(t, err)
assert.Equal(t, 0, len(entities)) assert.Equal(t, 0, len(entities))
@ -426,7 +427,7 @@ func TestRbacDropGrant(t *testing.T) {
for _, test := range tests { for _, test := range tests {
t.Run(test.description, func(t *testing.T) { 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 { if test.isValid {
assert.NoError(t, err) assert.NoError(t, err)
} else { } else {
@ -439,11 +440,11 @@ func TestRbacDropGrant(t *testing.T) {
func TestRbacListPolicy(t *testing.T) { func TestRbacListPolicy(t *testing.T) {
mt := generateMetaTable(t) mt := generateMetaTable(t)
policies, err := mt.ListPolicy(util.DefaultTenant) policies, err := mt.ListPolicy(context.TODO(), util.DefaultTenant)
assert.NoError(t, err) assert.NoError(t, err)
assert.Empty(t, policies) assert.Empty(t, policies)
userRoles, err := mt.ListUserRole(util.DefaultTenant) userRoles, err := mt.ListUserRole(context.TODO(), util.DefaultTenant)
assert.NoError(t, err) assert.NoError(t, err)
assert.Equal(t, 0, len(userRoles)) assert.Equal(t, 0, len(userRoles))
} }
@ -2091,24 +2092,24 @@ func TestMetaTable_PrivilegeGroup(t *testing.T) {
aliases: newNameDb(), aliases: newNameDb(),
catalog: catalog, catalog: catalog,
} }
err := mt.CreatePrivilegeGroup("pg1") err := mt.CreatePrivilegeGroup(context.TODO(), "pg1")
assert.Error(t, err) assert.Error(t, err)
err = mt.CreatePrivilegeGroup("") err = mt.CreatePrivilegeGroup(context.TODO(), "")
assert.Error(t, err) assert.Error(t, err)
err = mt.CreatePrivilegeGroup("Insert") err = mt.CreatePrivilegeGroup(context.TODO(), "Insert")
assert.Error(t, err) assert.Error(t, err)
err = mt.CreatePrivilegeGroup("pg2") err = mt.CreatePrivilegeGroup(context.TODO(), "pg2")
assert.NoError(t, err) assert.NoError(t, err)
err = mt.DropPrivilegeGroup("") err = mt.DropPrivilegeGroup(context.TODO(), "")
assert.Error(t, err) assert.Error(t, err)
err = mt.DropPrivilegeGroup("pg1") err = mt.DropPrivilegeGroup(context.TODO(), "pg1")
assert.NoError(t, err) 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) 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) assert.Error(t, err)
_, err = mt.GetPrivilegeGroupRoles("") _, err = mt.GetPrivilegeGroupRoles(context.TODO(), "")
assert.Error(t, err) assert.Error(t, err)
_, err = mt.ListPrivilegeGroups() _, err = mt.ListPrivilegeGroups(context.TODO())
assert.NoError(t, err) 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 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 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 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) DescribeAliasFunc func(ctx context.Context, dbName, alias string, ts Timestamp) (string, error)
ListAliasesFunc func(ctx context.Context, dbName, collectionName 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) GetCollectionIDByNameFunc func(name string) (UniqueID, error)
GetPartitionByNameFunc func(collID UniqueID, partitionName string, ts Timestamp) (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 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 RenameCollectionFunc func(ctx context.Context, oldName string, newName string, ts Timestamp) error
AddCredentialFunc func(credInfo *internalpb.CredentialInfo) error AddCredentialFunc func(ctx context.Context, credInfo *internalpb.CredentialInfo) error
GetCredentialFunc func(username string) (*internalpb.CredentialInfo, error) GetCredentialFunc func(ctx context.Context, username string) (*internalpb.CredentialInfo, error)
DeleteCredentialFunc func(username string) error DeleteCredentialFunc func(ctx context.Context, username string) error
AlterCredentialFunc func(credInfo *internalpb.CredentialInfo) error AlterCredentialFunc func(ctx context.Context, credInfo *internalpb.CredentialInfo) error
ListCredentialUsernamesFunc func() (*milvuspb.ListCredUsersResponse, error) ListCredentialUsernamesFunc func(ctx context.Context) (*milvuspb.ListCredUsersResponse, error)
CreateRoleFunc func(tenant string, entity *milvuspb.RoleEntity) error CreateRoleFunc func(ctx context.Context, tenant string, entity *milvuspb.RoleEntity) error
DropRoleFunc func(tenant string, roleName string) error DropRoleFunc func(ctx context.Context, tenant string, roleName string) error
OperateUserRoleFunc func(tenant string, userEntity *milvuspb.UserEntity, roleEntity *milvuspb.RoleEntity, operateType milvuspb.OperateUserRoleType) error OperateUserRoleFunc func(ctx context.Context, tenant string, userEntity *milvuspb.UserEntity, roleEntity *milvuspb.RoleEntity, operateType milvuspb.OperateUserRoleType) error
SelectRoleFunc func(tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) SelectRoleFunc func(ctx context.Context, tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error)
SelectUserFunc func(tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error) SelectUserFunc func(ctx context.Context, tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error)
OperatePrivilegeFunc func(tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error OperatePrivilegeFunc func(ctx context.Context, tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error
SelectGrantFunc func(tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error) SelectGrantFunc func(ctx context.Context, tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error)
DropGrantFunc func(tenant string, role *milvuspb.RoleEntity) error DropGrantFunc func(ctx context.Context, tenant string, role *milvuspb.RoleEntity) error
ListPolicyFunc func(tenant string) ([]string, error) ListPolicyFunc func(ctx context.Context, tenant string) ([]string, error)
ListUserRoleFunc func(tenant string) ([]string, error) ListUserRoleFunc func(ctx context.Context, tenant string) ([]string, error)
DescribeDatabaseFunc func(ctx context.Context, dbName string) (*model.Database, error) DescribeDatabaseFunc func(ctx context.Context, dbName string) (*model.Database, error)
CreatePrivilegeGroupFunc func(groupName string) error CreatePrivilegeGroupFunc func(ctx context.Context, groupName string) error
DropPrivilegeGroupFunc func(groupName string) error DropPrivilegeGroupFunc func(ctx context.Context, groupName string) error
ListPrivilegeGroupsFunc func() ([]*milvuspb.PrivilegeGroupInfo, error) ListPrivilegeGroupsFunc func(ctx context.Context) ([]*milvuspb.PrivilegeGroupInfo, error)
OperatePrivilegeGroupFunc func(groupName string, privileges []*milvuspb.PrivilegeEntity, operateType milvuspb.OperatePrivilegeGroupType) error OperatePrivilegeGroupFunc func(ctx context.Context, groupName string, privileges []*milvuspb.PrivilegeEntity, operateType milvuspb.OperatePrivilegeGroupType) error
GetPrivilegeGroupRolesFunc func(groupName string) ([]*milvuspb.RoleEntity, 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) { 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) return m.DropAliasFunc(ctx, dbName, alias, ts)
} }
func (m mockMetaTable) IsAlias(dbName, name string) bool { func (m mockMetaTable) IsAlias(ctx context.Context, dbName, name string) bool {
return m.IsAliasFunc(dbName, name) return m.IsAliasFunc(ctx, dbName, name)
} }
func (m mockMetaTable) DescribeAlias(ctx context.Context, dbName, alias string, ts Timestamp) (string, error) { 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) return m.ListAliasesFunc(ctx, dbName, collectionName, ts)
} }
func (m mockMetaTable) ListAliasesByID(collID UniqueID) []string { func (m mockMetaTable) ListAliasesByID(ctx context.Context, collID UniqueID) []string {
return m.ListAliasesByIDFunc(collID) return m.ListAliasesByIDFunc(ctx, collID)
} }
func (m mockMetaTable) AlterCollection(ctx context.Context, oldColl *model.Collection, newColl *model.Collection, ts Timestamp) error { 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) return m.GetPartitionByNameFunc(collID, partitionName, ts)
} }
func (m mockMetaTable) GetCollectionVirtualChannels(colID int64) []string { func (m mockMetaTable) GetCollectionVirtualChannels(ctx context.Context, colID int64) []string {
return m.GetCollectionVirtualChannelsFunc(colID) return m.GetCollectionVirtualChannelsFunc(ctx, colID)
} }
func (m mockMetaTable) AddCredential(credInfo *internalpb.CredentialInfo) error { func (m mockMetaTable) AddCredential(ctx context.Context, credInfo *internalpb.CredentialInfo) error {
return m.AddCredentialFunc(credInfo) return m.AddCredentialFunc(ctx, credInfo)
} }
func (m mockMetaTable) GetCredential(username string) (*internalpb.CredentialInfo, error) { func (m mockMetaTable) GetCredential(ctx context.Context, username string) (*internalpb.CredentialInfo, error) {
return m.GetCredentialFunc(username) return m.GetCredentialFunc(ctx, username)
} }
func (m mockMetaTable) DeleteCredential(username string) error { func (m mockMetaTable) DeleteCredential(ctx context.Context, username string) error {
return m.DeleteCredentialFunc(username) return m.DeleteCredentialFunc(ctx, username)
} }
func (m mockMetaTable) AlterCredential(credInfo *internalpb.CredentialInfo) error { func (m mockMetaTable) AlterCredential(ctx context.Context, credInfo *internalpb.CredentialInfo) error {
return m.AlterCredentialFunc(credInfo) return m.AlterCredentialFunc(ctx, credInfo)
} }
func (m mockMetaTable) ListCredentialUsernames() (*milvuspb.ListCredUsersResponse, error) { func (m mockMetaTable) ListCredentialUsernames(ctx context.Context) (*milvuspb.ListCredUsersResponse, error) {
return m.ListCredentialUsernamesFunc() return m.ListCredentialUsernamesFunc(ctx)
} }
func (m mockMetaTable) CreateRole(tenant string, entity *milvuspb.RoleEntity) error { func (m mockMetaTable) CreateRole(ctx context.Context, tenant string, entity *milvuspb.RoleEntity) error {
return m.CreateRoleFunc(tenant, entity) return m.CreateRoleFunc(ctx, tenant, entity)
} }
func (m mockMetaTable) DropRole(tenant string, roleName string) error { func (m mockMetaTable) DropRole(ctx context.Context, tenant string, roleName string) error {
return m.DropRoleFunc(tenant, roleName) return m.DropRoleFunc(ctx, tenant, roleName)
} }
func (m mockMetaTable) OperateUserRole(tenant string, userEntity *milvuspb.UserEntity, roleEntity *milvuspb.RoleEntity, operateType milvuspb.OperateUserRoleType) error { func (m mockMetaTable) OperateUserRole(ctx context.Context, tenant string, userEntity *milvuspb.UserEntity, roleEntity *milvuspb.RoleEntity, operateType milvuspb.OperateUserRoleType) error {
return m.OperateUserRoleFunc(tenant, userEntity, roleEntity, operateType) return m.OperateUserRoleFunc(ctx, tenant, userEntity, roleEntity, operateType)
} }
func (m mockMetaTable) SelectRole(tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) { func (m mockMetaTable) SelectRole(ctx context.Context, tenant string, entity *milvuspb.RoleEntity, includeUserInfo bool) ([]*milvuspb.RoleResult, error) {
return m.SelectRoleFunc(tenant, entity, includeUserInfo) return m.SelectRoleFunc(ctx, tenant, entity, includeUserInfo)
} }
func (m mockMetaTable) SelectUser(tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error) { func (m mockMetaTable) SelectUser(ctx context.Context, tenant string, entity *milvuspb.UserEntity, includeRoleInfo bool) ([]*milvuspb.UserResult, error) {
return m.SelectUserFunc(tenant, entity, includeRoleInfo) return m.SelectUserFunc(ctx, tenant, entity, includeRoleInfo)
} }
func (m mockMetaTable) OperatePrivilege(tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error { func (m mockMetaTable) OperatePrivilege(ctx context.Context, tenant string, entity *milvuspb.GrantEntity, operateType milvuspb.OperatePrivilegeType) error {
return m.OperatePrivilegeFunc(tenant, entity, operateType) return m.OperatePrivilegeFunc(ctx, tenant, entity, operateType)
} }
func (m mockMetaTable) SelectGrant(tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error) { func (m mockMetaTable) SelectGrant(ctx context.Context, tenant string, entity *milvuspb.GrantEntity) ([]*milvuspb.GrantEntity, error) {
return m.SelectGrantFunc(tenant, entity) return m.SelectGrantFunc(ctx, tenant, entity)
} }
func (m mockMetaTable) DropGrant(tenant string, role *milvuspb.RoleEntity) error { func (m mockMetaTable) DropGrant(ctx context.Context, tenant string, role *milvuspb.RoleEntity) error {
return m.DropGrantFunc(tenant, role) return m.DropGrantFunc(ctx, tenant, role)
} }
func (m mockMetaTable) ListPolicy(tenant string) ([]string, error) { func (m mockMetaTable) ListPolicy(ctx context.Context, tenant string) ([]string, error) {
return m.ListPolicyFunc(tenant) return m.ListPolicyFunc(ctx, tenant)
} }
func (m mockMetaTable) ListUserRole(tenant string) ([]string, error) { func (m mockMetaTable) ListUserRole(ctx context.Context, tenant string) ([]string, error) {
return m.ListUserRoleFunc(tenant) return m.ListUserRoleFunc(ctx, tenant)
} }
func (m mockMetaTable) CreatePrivilegeGroup(groupName string) error { func (m mockMetaTable) CreatePrivilegeGroup(ctx context.Context, groupName string) error {
return m.CreatePrivilegeGroupFunc(groupName) return m.CreatePrivilegeGroupFunc(ctx, groupName)
} }
func (m mockMetaTable) DropPrivilegeGroup(groupName string) error { func (m mockMetaTable) DropPrivilegeGroup(ctx context.Context, groupName string) error {
return m.DropPrivilegeGroupFunc(groupName) return m.DropPrivilegeGroupFunc(ctx, groupName)
} }
func (m mockMetaTable) ListPrivilegeGroups() ([]*milvuspb.PrivilegeGroupInfo, error) { func (m mockMetaTable) ListPrivilegeGroups(ctx context.Context) ([]*milvuspb.PrivilegeGroupInfo, error) {
return m.ListPrivilegeGroupsFunc() return m.ListPrivilegeGroupsFunc(ctx)
} }
func (m mockMetaTable) OperatePrivilegeGroup(groupName string, privileges []*milvuspb.PrivilegeEntity, operateType milvuspb.OperatePrivilegeGroupType) error { func (m mockMetaTable) OperatePrivilegeGroup(ctx context.Context, groupName string, privileges []*milvuspb.PrivilegeEntity, operateType milvuspb.OperatePrivilegeGroupType) error {
return m.OperatePrivilegeGroupFunc(groupName, privileges, operateType) return m.OperatePrivilegeGroupFunc(ctx, groupName, privileges, operateType)
} }
func (m mockMetaTable) GetPrivilegeGroupRoles(groupName string) ([]*milvuspb.RoleEntity, error) { func (m mockMetaTable) GetPrivilegeGroupRoles(ctx context.Context, groupName string) ([]*milvuspb.RoleEntity, error) {
return m.GetPrivilegeGroupRolesFunc(groupName) return m.GetPrivilegeGroupRolesFunc(ctx, groupName)
} }
func newMockMetaTable() *mockMetaTable { func newMockMetaTable() *mockMetaTable {
@ -498,49 +498,49 @@ func withInvalidMeta() Opt {
meta.DropAliasFunc = func(ctx context.Context, dbName string, alias string, ts Timestamp) error { meta.DropAliasFunc = func(ctx context.Context, dbName string, alias string, ts Timestamp) error {
return errors.New("error mock DropAlias") 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") 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") 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") 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") 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") 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") 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") 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") 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") 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") 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") 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") 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") 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") 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") return nil, errors.New("error mock ListUserRole")
} }
meta.DescribeAliasFunc = func(ctx context.Context, dbName, alias string, ts Timestamp) (string, error) { 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) { meta.DescribeDatabaseFunc = func(ctx context.Context, dbName string) (*model.Database, error) {
return nil, errors.New("error mock DescribeDatabase") 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") 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") 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") 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") 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 nil, errors.New("error mock GetPrivilegeGroupRoles")
} }
return withMeta(meta) 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.scheduler = newScheduler(c.ctx, c.idAllocator, c.tsoAllocator)
c.factory.Init(Params) 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) c.chanTimeTick = newTimeTickSync(c.ctx, c.session.ServerID, c.factory, chanMap)
log.Info("create TimeTick sync done") log.Info("create TimeTick sync done")
@ -549,11 +549,11 @@ func (c *Core) Init() error {
} }
func (c *Core) initCredentials() error { func (c *Core) initCredentials() error {
credInfo, _ := c.meta.GetCredential(util.UserRoot) credInfo, _ := c.meta.GetCredential(c.ctx, util.UserRoot)
if credInfo == nil { if credInfo == nil {
log.Debug("RootCoord init user root") log.Debug("RootCoord init user root")
encryptedRootPassword, _ := crypto.PasswordEncrypt(Params.CommonCfg.DefaultRootPassword.GetValue()) 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 err
} }
return nil return nil
@ -563,7 +563,7 @@ func (c *Core) initRbac() error {
var err error var err error
// create default roles, including admin, public // create default roles, including admin, public
for _, role := range util.DefaultRoles { 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) { if err != nil && !common.IsIgnorableError(err) {
return errors.Wrap(err, "failed to create role") return errors.Wrap(err, "failed to create role")
} }
@ -593,7 +593,7 @@ func (c *Core) initPublicRolePrivilege() error {
var err error var err error
for _, globalPrivilege := range globalPrivileges { 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}, Role: &milvuspb.RoleEntity{Name: util.RolePublic},
Object: &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Global.String()}, Object: &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Global.String()},
ObjectName: util.AnyWord, ObjectName: util.AnyWord,
@ -608,7 +608,7 @@ func (c *Core) initPublicRolePrivilege() error {
} }
} }
for _, collectionPrivilege := range collectionPrivileges { 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}, Role: &milvuspb.RoleEntity{Name: util.RolePublic},
Object: &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Collection.String()}, Object: &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Collection.String()},
ObjectName: util.AnyWord, ObjectName: util.AnyWord,
@ -672,7 +672,7 @@ func (c *Core) initBuiltinPrivilegeGroups() []*milvuspb.PrivilegeGroupInfo {
func (c *Core) initBuiltinRoles() error { func (c *Core) initBuiltinRoles() error {
rolePrivilegesMap := Params.RoleCfg.Roles.GetAsRoleDetails() rolePrivilegesMap := Params.RoleCfg.Roles.GetAsRoleDetails()
for role, privilegesJSON := range rolePrivilegesMap { 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) { if err != nil && !common.IsIgnorableError(err) {
log.Error("create a builtin role fail", zap.String("roleName", role), zap.Error(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) 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] { for _, privilege := range privilegesJSON[util.RoleConfigPrivileges] {
privilegeName := privilege[util.RoleConfigPrivilege] privilegeName := privilege[util.RoleConfigPrivilege]
if !util.IsAnyWord(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 { if err != nil {
return errors.Wrapf(err, "failed to get metastore privilege name for: %s", privilege[util.RoleConfigPrivilege]) return errors.Wrapf(err, "failed to get metastore privilege name for: %s", privilege[util.RoleConfigPrivilege])
} }
privilegeName = dbPrivName 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}, Role: &milvuspb.RoleEntity{Name: role},
Object: &milvuspb.ObjectEntity{Name: privilege[util.RoleConfigObjectType]}, Object: &milvuspb.ObjectEntity{Name: privilege[util.RoleConfigObjectType]},
ObjectName: privilege[util.RoleConfigObjectName], ObjectName: privilege[util.RoleConfigObjectName],
@ -1666,7 +1666,7 @@ func (c *Core) GetPChannelInfo(ctx context.Context, in *rootcoordpb.GetPChannelI
Status: merr.Status(err), Status: merr.Status(err),
}, nil }, nil
} }
return c.meta.GetPChannelInfo(in.GetPchannel()), nil return c.meta.GetPChannelInfo(ctx, in.GetPchannel()), nil
} }
// AllocTimestamp alloc timestamp // AllocTimestamp alloc timestamp
@ -2082,7 +2082,7 @@ func (c *Core) CreateCredential(ctx context.Context, credInfo *internalpb.Creden
} }
// insert to db // insert to db
err := c.meta.AddCredential(credInfo) err := c.meta.AddCredential(ctx, credInfo)
if err != nil { if err != nil {
ctxLog.Warn("CreateCredential save credential failed", zap.Error(err)) ctxLog.Warn("CreateCredential save credential failed", zap.Error(err))
metrics.RootCoordDDLReqCounter.WithLabelValues(method, metrics.FailLabel).Inc() 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 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 { if err != nil {
ctxLog.Warn("GetCredential query credential failed", zap.Error(err)) ctxLog.Warn("GetCredential query credential failed", zap.Error(err))
metrics.RootCoordDDLReqCounter.WithLabelValues(method, metrics.FailLabel).Inc() 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 return merr.Status(err), nil
} }
// update data on storage // update data on storage
err := c.meta.AlterCredential(credInfo) err := c.meta.AlterCredential(ctx, credInfo)
if err != nil { if err != nil {
ctxLog.Warn("UpdateCredential save credential failed", zap.Error(err)) ctxLog.Warn("UpdateCredential save credential failed", zap.Error(err))
metrics.RootCoordDDLReqCounter.WithLabelValues(method, metrics.FailLabel).Inc() 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 := newBaseRedoTask(c.stepExecutor)
redoTask.AddSyncStep(NewSimpleStep("delete credential meta data", func(ctx context.Context) ([]nestedStep, error) { 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 { if err != nil {
ctxLog.Warn("delete credential meta data failed", zap.Error(err)) 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 return &milvuspb.ListCredUsersResponse{Status: merr.Status(err)}, nil
} }
credInfo, err := c.meta.ListCredentialUsernames() credInfo, err := c.meta.ListCredentialUsernames(ctx)
if err != nil { if err != nil {
ctxLog.Warn("ListCredUsers query usernames failed", zap.Error(err)) ctxLog.Warn("ListCredUsers query usernames failed", zap.Error(err))
metrics.RootCoordDDLReqCounter.WithLabelValues(method, metrics.FailLabel).Inc() metrics.RootCoordDDLReqCounter.WithLabelValues(method, metrics.FailLabel).Inc()
@ -2269,7 +2269,7 @@ func (c *Core) CreateRole(ctx context.Context, in *milvuspb.CreateRoleRequest) (
} }
entity := in.Entity 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 { if err != nil {
errMsg := "fail to create role" errMsg := "fail to create role"
ctxLog.Warn(errMsg, zap.Error(err)) 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()) err := merr.WrapErrPrivilegeNotPermitted("the role[%s] is a builtin role, which can't be dropped", in.GetRoleName())
return merr.Status(err), nil 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" errMsg := "not found the role, maybe the role isn't existed or internal system error"
ctxLog.Warn(errMsg, zap.Error(err)) ctxLog.Warn(errMsg, zap.Error(err))
return merr.StatusWithErrorCode(errors.New(errMsg), commonpb.ErrorCode_DropRoleFailure), nil return merr.StatusWithErrorCode(errors.New(errMsg), commonpb.ErrorCode_DropRoleFailure), nil
} }
if !in.ForceDrop { 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}, Role: &milvuspb.RoleEntity{Name: in.RoleName},
}) })
if len(grantEntities) != 0 { if len(grantEntities) != 0 {
@ -2323,7 +2323,7 @@ func (c *Core) DropRole(ctx context.Context, in *milvuspb.DropRoleRequest) (*com
} }
redoTask := newBaseRedoTask(c.stepExecutor) redoTask := newBaseRedoTask(c.stepExecutor)
redoTask.AddSyncStep(NewSimpleStep("drop role meta data", func(ctx context.Context) ([]nestedStep, error) { 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 { if err != nil {
ctxLog.Warn("drop role mata data failed", zap.Error(err)) 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 { if !in.ForceDrop {
return nil, nil 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 { if err != nil {
ctxLog.Warn("drop the privilege list failed for the role", zap.Error(err)) 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 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" errMsg := "not found the role, maybe the role isn't existed or internal system error"
ctxLog.Warn(errMsg, zap.Error(err)) ctxLog.Warn(errMsg, zap.Error(err))
return merr.StatusWithErrorCode(errors.New(errMsg), commonpb.ErrorCode_OperateUserRoleFailure), nil return merr.StatusWithErrorCode(errors.New(errMsg), commonpb.ErrorCode_OperateUserRoleFailure), nil
} }
if in.Type != milvuspb.OperateUserRoleType_RemoveUserFromRole { 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" errMsg := "not found the user, maybe the user isn't existed or internal system error"
ctxLog.Warn(errMsg, zap.Error(err)) ctxLog.Warn(errMsg, zap.Error(err))
return merr.StatusWithErrorCode(errors.New(errMsg), commonpb.ErrorCode_OperateUserRoleFailure), nil 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 := newBaseRedoTask(c.stepExecutor)
redoTask.AddSyncStep(NewSimpleStep("operate user role meta data", func(ctx context.Context) ([]nestedStep, error) { 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) { if err != nil && !common.IsIgnorableError(err) {
log.Warn("operate user role mata data failed", zap.Error(err)) log.Warn("operate user role mata data failed", zap.Error(err))
return nil, err return nil, err
@ -2452,7 +2452,7 @@ func (c *Core) SelectRole(ctx context.Context, in *milvuspb.SelectRoleRequest) (
} }
if in.Role != nil { 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) { if errors.Is(err, merr.ErrIoKeyNotFound) {
return &milvuspb.SelectRoleResponse{ return &milvuspb.SelectRoleResponse{
Status: merr.Success(), Status: merr.Success(),
@ -2465,7 +2465,7 @@ func (c *Core) SelectRole(ctx context.Context, in *milvuspb.SelectRoleRequest) (
}, nil }, 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 { if err != nil {
errMsg := "fail to select the role" errMsg := "fail to select the role"
ctxLog.Warn(errMsg, zap.Error(err)) 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 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) { if errors.Is(err, merr.ErrIoKeyNotFound) {
return &milvuspb.SelectUserResponse{ return &milvuspb.SelectUserResponse{
Status: merr.Success(), Status: merr.Success(),
@ -2512,7 +2512,7 @@ func (c *Core) SelectUser(ctx context.Context, in *milvuspb.SelectUserRequest) (
}, nil }, 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 { if err != nil {
errMsg := "fail to select the user" errMsg := "fail to select the user"
ctxLog.Warn(errMsg, zap.Error(err)) ctxLog.Warn(errMsg, zap.Error(err))
@ -2537,7 +2537,7 @@ func (c *Core) isValidRole(entity *milvuspb.RoleEntity) error {
if entity.Name == "" { if entity.Name == "" {
return errors.New("the name in the role entity is empty") 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)) 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") 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 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 { if entity == nil {
return errors.New("the grantor entity is nil") return errors.New("the grantor entity is nil")
} }
if entity.User == nil || entity.User.Name == "" { if entity.User == nil || entity.User.Name == "" {
return errors.New("the user entity in the grantor entity is nil or empty") 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)) 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") 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 // 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 { if err != nil {
return err return err
} }
@ -2629,7 +2629,7 @@ func (c *Core) OperatePrivilege(ctx context.Context, in *milvuspb.OperatePrivile
ctxLog.Warn("", zap.Error(err)) ctxLog.Warn("", zap.Error(err))
return merr.StatusWithErrorCode(err, commonpb.ErrorCode_OperatePrivilegeFailure), nil 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)) ctxLog.Error("", zap.Error(err))
return merr.StatusWithErrorCode(err, commonpb.ErrorCode_OperatePrivilegeFailure), nil 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) { redoTask.AddSyncStep(NewSimpleStep("operate privilege meta data", func(ctx context.Context) ([]nestedStep, error) {
if !util.IsAnyWord(privName) { if !util.IsAnyWord(privName) {
// set up privilege name for metastore // set up privilege name for metastore
dbPrivName, err := c.getMetastorePrivilegeName(privName) dbPrivName, err := c.getMetastorePrivilegeName(ctx, privName)
if err != nil { if err != nil {
return nil, err return nil, err
} }
in.Entity.Grantor.Privilege.Name = dbPrivName 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) { if err != nil && !common.IsIgnorableError(err) {
log.Warn("fail to operate the privilege", zap.Any("in", in), zap.Error(err)) log.Warn("fail to operate the privilege", zap.Any("in", in), zap.Error(err))
return nil, err return nil, err
@ -2674,7 +2674,7 @@ func (c *Core) OperatePrivilege(ctx context.Context, in *milvuspb.OperatePrivile
} }
grants := []*milvuspb.GrantEntity{in.Entity} grants := []*milvuspb.GrantEntity{in.Entity}
allGroups, err := c.meta.ListPrivilegeGroups() allGroups, err := c.meta.ListPrivilegeGroups(ctx)
allGroups = append(allGroups, c.initBuiltinPrivilegeGroups()...) allGroups = append(allGroups, c.initBuiltinPrivilegeGroups()...)
if err != nil { if err != nil {
return nil, err 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) { groups := lo.SliceToMap(allGroups, func(group *milvuspb.PrivilegeGroupInfo) (string, []*milvuspb.PrivilegeEntity) {
return group.GroupName, group.Privileges return group.GroupName, group.Privileges
}) })
expandGrants, err := c.expandPrivilegeGroups(grants, groups) expandGrants, err := c.expandPrivilegeGroups(ctx, grants, groups)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -2709,13 +2709,13 @@ func (c *Core) OperatePrivilege(ctx context.Context, in *milvuspb.OperatePrivile
return merr.Success(), nil 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 it is built-in privilege, return the privilege name directly
if util.IsPrivilegeNameDefined(privName) { if util.IsPrivilegeNameDefined(privName) {
return util.PrivilegeNameForMetastore(privName), nil return util.PrivilegeNameForMetastore(privName), nil
} }
// return the privilege group name if it is a custom privilege group // 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 { if err != nil {
return "", err 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) { if errors.Is(err, merr.ErrIoKeyNotFound) {
return &milvuspb.SelectGrantResponse{ return &milvuspb.SelectGrantResponse{
Status: merr.Success(), Status: merr.Success(),
@ -2801,7 +2801,7 @@ func (c *Core) ListPolicy(ctx context.Context, in *internalpb.ListPolicyRequest)
}, nil }, nil
} }
policies, err := c.meta.ListPolicy(util.DefaultTenant) policies, err := c.meta.ListPolicy(ctx, util.DefaultTenant)
if err != nil { if err != nil {
errMsg := "fail to list policy" errMsg := "fail to list policy"
ctxLog.Warn(errMsg, zap.Error(err)) 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), Status: merr.StatusWithErrorCode(errors.New(errMsg), commonpb.ErrorCode_ListPolicyFailure),
}, nil }, nil
} }
userRoles, err := c.meta.ListUserRole(util.DefaultTenant) userRoles, err := c.meta.ListUserRole(ctx, util.DefaultTenant)
if err != nil { if err != nil {
errMsg := "fail to list user-role" errMsg := "fail to list user-role"
ctxLog.Warn(errMsg, zap.Any("in", in), zap.Error(err)) 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), Status: merr.StatusWithErrorCode(errors.New(errMsg), commonpb.ErrorCode_ListPolicyFailure),
}, nil }, nil
} }
privGroups, err := c.meta.ListPrivilegeGroups() privGroups, err := c.meta.ListPrivilegeGroups(ctx)
if err != nil { if err != nil {
errMsg := "fail to list privilege groups" errMsg := "fail to list privilege groups"
ctxLog.Warn(errMsg, zap.Error(err)) 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 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)) ctxLog.Warn("fail to create privilege group", zap.Error(err))
return merr.Status(err), nil return merr.Status(err), nil
} }
@ -3069,7 +3069,7 @@ func (c *Core) DropPrivilegeGroup(ctx context.Context, in *milvuspb.DropPrivileg
return merr.Status(err), nil 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)) ctxLog.Warn("fail to drop privilege group", zap.Error(err))
return merr.Status(err), nil return merr.Status(err), nil
} }
@ -3094,7 +3094,7 @@ func (c *Core) ListPrivilegeGroups(ctx context.Context, in *milvuspb.ListPrivile
}, nil }, nil
} }
privGroups, err := c.meta.ListPrivilegeGroups() privGroups, err := c.meta.ListPrivilegeGroups(ctx)
if err != nil { if err != nil {
ctxLog.Warn("fail to list privilege group", zap.Error(err)) ctxLog.Warn("fail to list privilege group", zap.Error(err))
return &milvuspb.ListPrivilegeGroupsResponse{ return &milvuspb.ListPrivilegeGroupsResponse{
@ -3124,7 +3124,7 @@ func (c *Core) OperatePrivilegeGroup(ctx context.Context, in *milvuspb.OperatePr
redoTask := newBaseRedoTask(c.stepExecutor) redoTask := newBaseRedoTask(c.stepExecutor)
redoTask.AddSyncStep(NewSimpleStep("operate privilege group", func(ctx context.Context) ([]nestedStep, error) { 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) { if err != nil && !common.IsIgnorableError(err) {
log.Warn("fail to list privilege groups", zap.Error(err)) log.Warn("fail to list privilege groups", zap.Error(err))
return nil, err return nil, err
@ -3134,7 +3134,7 @@ func (c *Core) OperatePrivilegeGroup(ctx context.Context, in *milvuspb.OperatePr
}) })
// get roles granted to the group // get roles granted to the group
roles, err := c.meta.GetPrivilegeGroupRoles(in.GroupName) roles, err := c.meta.GetPrivilegeGroupRoles(ctx, in.GroupName)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -3170,18 +3170,18 @@ func (c *Core) OperatePrivilegeGroup(ctx context.Context, in *milvuspb.OperatePr
a.DbName == b.DbName a.DbName == b.DbName
} }
for _, role := range roles { 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, Role: role,
DbName: util.AnyWord, DbName: util.AnyWord,
}) })
if err != nil { if err != nil {
return nil, err return nil, err
} }
currGrants, err := c.expandPrivilegeGroups(grants, currGroups) currGrants, err := c.expandPrivilegeGroups(ctx, grants, currGroups)
if err != nil { if err != nil {
return nil, err return nil, err
} }
newGrants, err := c.expandPrivilegeGroups(grants, newGroups) newGrants, err := c.expandPrivilegeGroups(ctx, grants, newGroups)
if err != nil { if err != nil {
return nil, err 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) { 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) { if err != nil && !common.IsIgnorableError(err) {
log.Warn("fail to operate privilege group", zap.Error(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 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{} newGrants := []*milvuspb.GrantEntity{}
for _, grant := range grants { for _, grant := range grants {
privName := grant.Grantor.Privilege.Name privName := grant.Grantor.Privilege.Name
if privGroup, exists := groups[privName]; !exists { if privGroup, exists := groups[privName]; !exists {
metaName, err := c.getMetastorePrivilegeName(privName) metaName, err := c.getMetastorePrivilegeName(ctx, privName)
if err != nil { if err != nil {
return nil, err return nil, err
} }
@ -3271,7 +3271,7 @@ func (c *Core) expandPrivilegeGroups(grants []*milvuspb.GrantEntity, groups map[
}) })
} else { } else {
for _, priv := range privGroup { for _, priv := range privGroup {
metaName, err := c.getMetastorePrivilegeName(priv.Name) metaName, err := c.getMetastorePrivilegeName(ctx, priv.Name)
if err != nil { if err != nil {
return nil, err 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) { t.Run("operate user role failed", func(t *testing.T) {
mockMeta := c.meta.(*mockMetaTable) 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 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 return nil, nil
} }
resp, err := c.OperateUserRole(ctx, &milvuspb.OperateUserRoleRequest{RoleName: "foo", Username: "bar", Type: milvuspb.OperateUserRoleType_AddUserToRole}) resp, err := c.OperateUserRole(ctx, &milvuspb.OperateUserRoleRequest{RoleName: "foo", Username: "bar", Type: milvuspb.OperateUserRoleType_AddUserToRole})
assert.NoError(t, err) assert.NoError(t, err)
assert.NotEqual(t, commonpb.ErrorCode_Success, resp.ErrorCode) 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") 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") return nil, errors.New("mock error")
} }
}) })
@ -1745,10 +1745,10 @@ func TestRootCoord_RBACError(t *testing.T) {
} }
mockMeta := c.meta.(*mockMetaTable) 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 return nil, nil
} }
mockMeta.ListPrivilegeGroupsFunc = func() ([]*milvuspb.PrivilegeGroupInfo, error) { mockMeta.ListPrivilegeGroupsFunc = func(ctx context.Context) ([]*milvuspb.PrivilegeGroupInfo, error) {
return nil, nil return nil, nil
} }
{ {
@ -1765,7 +1765,7 @@ func TestRootCoord_RBACError(t *testing.T) {
assert.NotEqual(t, commonpb.ErrorCode_Success, resp.ErrorCode) 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 return nil, nil
} }
resp, err := c.OperatePrivilege(ctx, &milvuspb.OperatePrivilegeRequest{Entity: &milvuspb.GrantEntity{ resp, err := c.OperatePrivilege(ctx, &milvuspb.OperatePrivilegeRequest{Entity: &milvuspb.GrantEntity{
@ -1779,23 +1779,23 @@ func TestRootCoord_RBACError(t *testing.T) {
}, Type: milvuspb.OperatePrivilegeType_Grant}) }, Type: milvuspb.OperatePrivilegeType_Grant})
assert.NoError(t, err) assert.NoError(t, err)
assert.NotEqual(t, commonpb.ErrorCode_Success, resp.ErrorCode) 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") 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") return nil, errors.New("mock error")
} }
}) })
t.Run("operate privilege group failed", func(t *testing.T) { t.Run("operate privilege group failed", func(t *testing.T) {
mockMeta := c.meta.(*mockMetaTable) 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") 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") 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") 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()) assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
} }
mockMeta := c.meta.(*mockMetaTable) 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 return nil, nil
} }
{ {
@ -1845,21 +1845,21 @@ func TestRootCoord_RBACError(t *testing.T) {
assert.NoError(t, err) assert.NoError(t, err)
assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode()) 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") return nil, errors.New("mock error")
} }
}) })
t.Run("select grant success", func(t *testing.T) { t.Run("select grant success", func(t *testing.T) {
mockMeta := c.meta.(*mockMetaTable) 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{ return []*milvuspb.RoleResult{
{ {
Role: &milvuspb.RoleEntity{Name: "foo"}, Role: &milvuspb.RoleEntity{Name: "foo"},
}, },
}, nil }, 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{ return []*milvuspb.GrantEntity{
{ {
Role: &milvuspb.RoleEntity{Name: "foo"}, Role: &milvuspb.RoleEntity{Name: "foo"},
@ -1874,11 +1874,11 @@ func TestRootCoord_RBACError(t *testing.T) {
assert.Equal(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode()) 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") 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") 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()) assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode())
mockMeta := c.meta.(*mockMetaTable) mockMeta := c.meta.(*mockMetaTable)
mockMeta.ListPolicyFunc = func(tenant string) ([]string, error) { mockMeta.ListPolicyFunc = func(ctx context.Context, tenant string) ([]string, error) {
return []string{}, nil return []string{}, nil
} }
resp, err = c.ListPolicy(ctx, &internalpb.ListPolicyRequest{}) resp, err = c.ListPolicy(ctx, &internalpb.ListPolicyRequest{})
assert.NoError(t, err) assert.NoError(t, err)
assert.NotEqual(t, commonpb.ErrorCode_Success, resp.GetStatus().GetErrorCode()) 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") 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) { t.Run("init builtin roles success", func(t *testing.T) {
c := newTestCore(withHealthyCode(), withInvalidMeta()) c := newTestCore(withHealthyCode(), withInvalidMeta())
mockMeta := c.meta.(*mockMetaTable) 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 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 return nil
} }
mockMeta.ListPrivilegeGroupsFunc = func() ([]*milvuspb.PrivilegeGroupInfo, error) { mockMeta.ListPrivilegeGroupsFunc = func(ctx context.Context) ([]*milvuspb.PrivilegeGroupInfo, error) {
return nil, nil return nil, nil
} }
err := c.initBuiltinRoles() 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) { t.Run("init builtin roles fail to create role", func(t *testing.T) {
c := newTestCore(withHealthyCode(), withInvalidMeta()) c := newTestCore(withHealthyCode(), withInvalidMeta())
mockMeta := c.meta.(*mockMetaTable) 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 return merr.ErrPrivilegeNotPermitted
} }
err := c.initBuiltinRoles() 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) { t.Run("init builtin roles fail to operate privileg", func(t *testing.T) {
c := newTestCore(withHealthyCode(), withInvalidMeta()) c := newTestCore(withHealthyCode(), withInvalidMeta())
mockMeta := c.meta.(*mockMetaTable) 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 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 return merr.ErrPrivilegeNotPermitted
} }
err := c.initBuiltinRoles() 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) { t.Run("init default role and public role privilege", func(t *testing.T) {
meta := mockrootcoord.NewIMetaTable(t) meta := mockrootcoord.NewIMetaTable(t)
c := newTestCore(withHealthyCode(), withMeta(meta)) c := newTestCore(withHealthyCode(), withMeta(meta))
meta.EXPECT().CreateRole(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).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.RoleCfg.Enabled.Key, "false")
Params.Save(Params.ProxyCfg.EnablePublicPrivilege.Key, "true") 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": "*"}]}}` builtinRoles := `{"db_admin": {"privileges": [{"object_type": "Global", "object_name": "*", "privilege": "CreateCollection", "db_name": "*"}]}}`
meta := mockrootcoord.NewIMetaTable(t) meta := mockrootcoord.NewIMetaTable(t)
c := newTestCore(withHealthyCode(), withMeta(meta)) c := newTestCore(withHealthyCode(), withMeta(meta))
meta.EXPECT().CreateRole(mock.Anything, mock.Anything).Return(nil).Times(3) meta.EXPECT().CreateRole(mock.Anything, mock.Anything, mock.Anything).Return(nil).Times(3)
meta.EXPECT().OperatePrivilege(mock.Anything, mock.Anything, mock.Anything).Return(nil).Once() 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.Enabled.Key, "true")
Params.Save(Params.RoleCfg.Roles.Key, builtinRoles) 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) privilegeColls.Insert(util.AnyWord)
return privilegeColls, nil return privilegeColls, nil
} }
userRoles, err := t.core.meta.SelectUser("", &milvuspb.UserEntity{ userRoles, err := t.core.meta.SelectUser(ctx, "", &milvuspb.UserEntity{
Name: curUser, Name: curUser,
}, true) }, true)
if err != nil { if err != nil {
@ -79,7 +79,7 @@ func (t *showCollectionTask) Execute(ctx context.Context) error {
privilegeColls.Insert(util.AnyWord) privilegeColls.Insert(util.AnyWord)
return privilegeColls, nil return privilegeColls, nil
} }
entities, err := t.core.meta.SelectGrant("", &milvuspb.GrantEntity{ entities, err := t.core.meta.SelectGrant(ctx, "", &milvuspb.GrantEntity{
Role: role, Role: role,
DbName: t.Req.GetDbName(), 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) meta := mockrootcoord.NewIMetaTable(t)
core := newTestCore(withMeta(meta)) 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() Return(nil, errors.New("mock error: select user")).Once()
task := &showCollectionTask{ task := &showCollectionTask{
@ -189,7 +189,7 @@ func TestShowCollectionsAuth(t *testing.T) {
meta := mockrootcoord.NewIMetaTable(t) meta := mockrootcoord.NewIMetaTable(t)
core := newTestCore(withMeta(meta)) 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() Return([]*milvuspb.UserResult{}, nil).Once()
task := &showCollectionTask{ task := &showCollectionTask{
@ -210,7 +210,7 @@ func TestShowCollectionsAuth(t *testing.T) {
meta := mockrootcoord.NewIMetaTable(t) meta := mockrootcoord.NewIMetaTable(t)
core := newTestCore(withMeta(meta)) 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{ Return([]*milvuspb.UserResult{
{ {
User: &milvuspb.UserEntity{ User: &milvuspb.UserEntity{
@ -250,7 +250,7 @@ func TestShowCollectionsAuth(t *testing.T) {
meta := mockrootcoord.NewIMetaTable(t) meta := mockrootcoord.NewIMetaTable(t)
core := newTestCore(withMeta(meta)) 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{ Return([]*milvuspb.UserResult{
{ {
User: &milvuspb.UserEntity{ User: &milvuspb.UserEntity{
@ -263,7 +263,7 @@ func TestShowCollectionsAuth(t *testing.T) {
}, },
}, },
}, nil).Once() }, 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{ task := &showCollectionTask{
baseTask: newBaseTask(context.Background(), core), baseTask: newBaseTask(context.Background(), core),
@ -281,7 +281,7 @@ func TestShowCollectionsAuth(t *testing.T) {
meta := mockrootcoord.NewIMetaTable(t) meta := mockrootcoord.NewIMetaTable(t)
core := newTestCore(withMeta(meta)) 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{ Return([]*milvuspb.UserResult{
{ {
User: &milvuspb.UserEntity{ User: &milvuspb.UserEntity{
@ -294,7 +294,7 @@ func TestShowCollectionsAuth(t *testing.T) {
}, },
}, },
}, nil).Once() }, 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()}, Object: &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Global.String()},
Grantor: &milvuspb.GrantorEntity{ Grantor: &milvuspb.GrantorEntity{
@ -331,7 +331,7 @@ func TestShowCollectionsAuth(t *testing.T) {
meta := mockrootcoord.NewIMetaTable(t) meta := mockrootcoord.NewIMetaTable(t)
core := newTestCore(withMeta(meta)) 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{ Return([]*milvuspb.UserResult{
{ {
User: &milvuspb.UserEntity{ User: &milvuspb.UserEntity{
@ -344,7 +344,7 @@ func TestShowCollectionsAuth(t *testing.T) {
}, },
}, },
}, nil).Once() }, 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()}, Object: &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Collection.String()},
ObjectName: util.AnyWord, ObjectName: util.AnyWord,
@ -376,7 +376,7 @@ func TestShowCollectionsAuth(t *testing.T) {
meta := mockrootcoord.NewIMetaTable(t) meta := mockrootcoord.NewIMetaTable(t)
core := newTestCore(withMeta(meta)) 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{ Return([]*milvuspb.UserResult{
{ {
User: &milvuspb.UserEntity{ User: &milvuspb.UserEntity{
@ -389,7 +389,7 @@ func TestShowCollectionsAuth(t *testing.T) {
}, },
}, },
}, nil).Once() }, 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()}, Object: &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Collection.String()},
ObjectName: "a", ObjectName: "a",