enhance: [2.5] avoid frequent LoadWithPrefix etcd calls in ShowCollections and DescribeCollections (#43903)

pr: #43902
related: https://github.com/milvus-io/milvus/issues/43901

Signed-off-by: shaoting-huang <shaoting.huang@zilliz.com>
This commit is contained in:
sthuang 2025-08-19 12:29:46 +08:00 committed by GitHub
parent b57d104742
commit c6e6bcece4
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 36 additions and 11 deletions

View File

@ -262,6 +262,7 @@ func TestDescribeCollectionsAuth(t *testing.T) {
},
},
}, nil).Once()
meta.EXPECT().ListPrivilegeGroups(mock.Anything).Return(nil, nil).Once()
task := getTask(core)
@ -341,6 +342,7 @@ func TestDescribeCollectionsAuth(t *testing.T) {
ID: 1,
Name: "test db",
}, nil).Once()
meta.EXPECT().ListPrivilegeGroups(mock.Anything).Return(nil, nil).Once()
task := getTask(core)
ctx := GetContext(context.Background(), "foo:root")
@ -378,6 +380,7 @@ func TestDescribeCollectionsAuth(t *testing.T) {
},
}, nil).Once()
meta.EXPECT().SelectGrant(mock.Anything, mock.Anything, mock.Anything).Return(nil, errors.New("mock error: select grant")).Once()
meta.EXPECT().ListPrivilegeGroups(mock.Anything).Return(nil, nil).Once()
task := getTask(core)
ctx := GetContext(context.Background(), "foo:root")
@ -424,6 +427,7 @@ func TestDescribeCollectionsAuth(t *testing.T) {
ID: 1,
Name: "test db",
}, nil).Once()
meta.EXPECT().ListPrivilegeGroups(mock.Anything).Return(nil, nil).Once()
task := getTask(core)
ctx := GetContext(context.Background(), "foo:root")
@ -476,7 +480,7 @@ func TestDescribeCollectionsAuth(t *testing.T) {
ID: 1,
Name: "test db",
}, nil).Once()
meta.EXPECT().IsCustomPrivilegeGroup(mock.Anything, util.PrivilegeNameForAPI(commonpb.ObjectPrivilege_PrivilegeGroupCollectionReadOnly.String())).Return(false, nil).Once()
meta.EXPECT().ListPrivilegeGroups(mock.Anything).Return(nil, nil).Once()
task := getTask(core)
ctx := GetContext(context.Background(), "foo:root")
@ -524,6 +528,7 @@ func TestDescribeCollectionsAuth(t *testing.T) {
ID: 1,
Name: "test db",
}, nil).Once()
meta.EXPECT().ListPrivilegeGroups(mock.Anything).Return(nil, nil).Once()
task := getTask(core)
ctx := GetContext(context.Background(), "foo:root")
@ -577,7 +582,7 @@ func TestDescribeCollectionsAuth(t *testing.T) {
ID: 1,
Name: "test db",
}, nil).Once()
meta.EXPECT().IsCustomPrivilegeGroup(mock.Anything, mock.Anything).Return(false, nil).Once()
meta.EXPECT().ListPrivilegeGroups(mock.Anything).Return(nil, nil).Once()
task := getTask(core)
ctx := GetContext(context.Background(), "foo:root")
@ -620,7 +625,11 @@ func TestDescribeCollectionsAuth(t *testing.T) {
ObjectName: "test coll",
},
}, nil).Once()
meta.EXPECT().IsCustomPrivilegeGroup(mock.Anything, "privilege_group").Return(true, nil).Once()
meta.EXPECT().ListPrivilegeGroups(mock.Anything).Return([]*milvuspb.PrivilegeGroupInfo{
{
GroupName: "privilege_group",
},
}, nil).Once()
meta.EXPECT().GetCollectionByName(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(&model.Collection{
CollectionID: 1,
Name: "test coll",

View File

@ -3346,6 +3346,16 @@ func (c *Core) getCurrentUserVisibleCollections(ctx context.Context, databaseNam
if len(userRoles) == 0 {
return privilegeColls, nil
}
// Get all custom privilege groups
customPrivilegeGroups, err := c.meta.ListPrivilegeGroups(ctx)
if err != nil {
return nil, err
}
customPrivilegeGroupMap := lo.SliceToMap(customPrivilegeGroups, func(group *milvuspb.PrivilegeGroupInfo) (string, struct{}) {
return group.GroupName, struct{}{}
})
for _, role := range userRoles[0].Roles {
if role.GetName() == util.RoleAdmin {
privilegeColls.Insert(util.AnyWord)
@ -3371,11 +3381,8 @@ func (c *Core) getCurrentUserVisibleCollections(ctx context.Context, databaseNam
}
// should list collection level built-in privilege group or custom privilege group objects
if objectType != commonpb.ObjectType_Collection.String() {
customGroup, err := c.meta.IsCustomPrivilegeGroup(ctx, priv)
if err != nil {
return nil, err
}
if !customGroup && !Params.RbacConfig.IsCollectionPrivilegeGroup(priv) {
_, isCustomGroup := customPrivilegeGroupMap[priv]
if !isCustomGroup && !Params.RbacConfig.IsCollectionPrivilegeGroup(priv) {
continue
}
}

View File

@ -212,6 +212,7 @@ func TestShowCollectionsAuth(t *testing.T) {
},
},
}, nil).Once()
meta.EXPECT().ListPrivilegeGroups(mock.Anything).Return(nil, nil).Once()
task := &showCollectionTask{
baseTask: newBaseTask(context.Background(), core),
@ -293,6 +294,7 @@ func TestShowCollectionsAuth(t *testing.T) {
CreateTime: tsoutil.GetCurrentTime(),
},
}, nil).Once()
meta.EXPECT().ListPrivilegeGroups(mock.Anything).Return(nil, nil).Once()
task := &showCollectionTask{
baseTask: newBaseTask(context.Background(), core),
@ -326,6 +328,7 @@ func TestShowCollectionsAuth(t *testing.T) {
},
}, nil).Once()
meta.EXPECT().SelectGrant(mock.Anything, mock.Anything, mock.Anything).Return(nil, errors.New("mock error: select grant")).Once()
meta.EXPECT().ListPrivilegeGroups(mock.Anything).Return(nil, nil).Once()
task := &showCollectionTask{
baseTask: newBaseTask(context.Background(), core),
@ -374,6 +377,7 @@ func TestShowCollectionsAuth(t *testing.T) {
CreateTime: tsoutil.GetCurrentTime(),
},
}, nil).Once()
meta.EXPECT().ListPrivilegeGroups(mock.Anything).Return(nil, nil).Once()
task := &showCollectionTask{
baseTask: newBaseTask(context.Background(), core),
@ -425,7 +429,7 @@ func TestShowCollectionsAuth(t *testing.T) {
CreateTime: tsoutil.GetCurrentTime(),
},
}, nil).Once()
meta.EXPECT().IsCustomPrivilegeGroup(mock.Anything, util.PrivilegeNameForAPI(commonpb.ObjectPrivilege_PrivilegeGroupCollectionReadOnly.String())).Return(false, nil).Once()
meta.EXPECT().ListPrivilegeGroups(mock.Anything).Return(nil, nil).Once()
task := &showCollectionTask{
baseTask: newBaseTask(context.Background(), core),
@ -472,6 +476,7 @@ func TestShowCollectionsAuth(t *testing.T) {
CreateTime: tsoutil.GetCurrentTime(),
},
}, nil).Once()
meta.EXPECT().ListPrivilegeGroups(mock.Anything).Return(nil, nil).Once()
task := &showCollectionTask{
baseTask: newBaseTask(context.Background(), core),
@ -530,7 +535,7 @@ func TestShowCollectionsAuth(t *testing.T) {
CreateTime: tsoutil.GetCurrentTime(),
},
}, nil).Once()
meta.EXPECT().IsCustomPrivilegeGroup(mock.Anything, mock.Anything).Return(false, nil).Once()
meta.EXPECT().ListPrivilegeGroups(mock.Anything).Return(nil, nil).Once()
task := &showCollectionTask{
baseTask: newBaseTask(context.Background(), core),
@ -574,7 +579,11 @@ func TestShowCollectionsAuth(t *testing.T) {
ObjectName: "test_collection",
},
}, nil).Once()
meta.EXPECT().IsCustomPrivilegeGroup(mock.Anything, "privilege_group").Return(true, nil).Once()
meta.EXPECT().ListPrivilegeGroups(mock.Anything).Return([]*milvuspb.PrivilegeGroupInfo{
{
GroupName: "privilege_group",
},
}, nil).Once()
meta.EXPECT().ListCollections(mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return([]*model.Collection{
{
DBID: 1,