From c6e6bcece484c618c2ce7bd6ac09720ed38fa0b7 Mon Sep 17 00:00:00 2001 From: sthuang <167743503+shaoting-huang@users.noreply.github.com> Date: Tue, 19 Aug 2025 12:29:46 +0800 Subject: [PATCH] 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 --- .../rootcoord/describe_collection_task_test.go | 15 ++++++++++++--- internal/rootcoord/root_coord.go | 17 ++++++++++++----- internal/rootcoord/show_collection_task_test.go | 15 ++++++++++++--- 3 files changed, 36 insertions(+), 11 deletions(-) diff --git a/internal/rootcoord/describe_collection_task_test.go b/internal/rootcoord/describe_collection_task_test.go index 38b0687543..99f729e8d2 100644 --- a/internal/rootcoord/describe_collection_task_test.go +++ b/internal/rootcoord/describe_collection_task_test.go @@ -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", diff --git a/internal/rootcoord/root_coord.go b/internal/rootcoord/root_coord.go index b9c71d6a8b..0de5ca1884 100644 --- a/internal/rootcoord/root_coord.go +++ b/internal/rootcoord/root_coord.go @@ -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 } } diff --git a/internal/rootcoord/show_collection_task_test.go b/internal/rootcoord/show_collection_task_test.go index e133c94ecf..9701ad94ff 100644 --- a/internal/rootcoord/show_collection_task_test.go +++ b/internal/rootcoord/show_collection_task_test.go @@ -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,