From 463765922e22116cdc0f6704438ad14c329af299 Mon Sep 17 00:00:00 2001 From: SimFG Date: Tue, 23 Jan 2024 15:56:54 +0800 Subject: [PATCH] enhance: support related privilege for grant api (#30153) /kind improvement Signed-off-by: SimFG --- go.sum | 14 ---------- internal/proxy/impl.go | 16 ++++++++++++ internal/proxy/proxy_test.go | 50 ++++++++++++++++++++++++++++++++++++ pkg/util/constant.go | 10 ++++++++ 4 files changed, 76 insertions(+), 14 deletions(-) diff --git a/go.sum b/go.sum index d1f7f972df..20ad5cad69 100644 --- a/go.sum +++ b/go.sum @@ -590,20 +590,6 @@ github.com/milvus-io/gorocksdb v0.0.0-20220624081344-8c5f4212846b h1:TfeY0NxYxZz github.com/milvus-io/gorocksdb v0.0.0-20220624081344-8c5f4212846b/go.mod h1:iwW+9cWfIzzDseEBCCeDSN5SD16Tidvy8cwQ7ZY8Qj4= github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240109020841-d367b5a59df1 h1:oNpMivd94JAMhdSVsFw8t1b+olXz8pbzd5PES21sth8= github.com/milvus-io/milvus-proto/go-api/v2 v2.3.4-0.20240109020841-d367b5a59df1/go.mod h1:1OIl0v5PQeNxIJhCvY+K55CBUOYDZevw9g9380u1Wek= -github.com/milvus-io/milvus-storage/go v0.0.0-20231109072809-1cd7b0866092 h1:UYJ7JB+QlMOoFHNdd8mUa3/lV63t9dnBX7ILXmEEWPY= -github.com/milvus-io/milvus-storage/go v0.0.0-20231109072809-1cd7b0866092/go.mod h1:GPETMcTZq1gLY1WA6Na5kiNAKnq8SEMMiVKUZrM3sho= -github.com/milvus-io/milvus-storage/go v0.0.0-20231213080429-ed6b9bd5c9d2 h1:2epYWKCSY6Rq/aJ/6UyUS1d3+Yts0UK8HNiWGjVN4Pc= -github.com/milvus-io/milvus-storage/go v0.0.0-20231213080429-ed6b9bd5c9d2/go.mod h1:GPETMcTZq1gLY1WA6Na5kiNAKnq8SEMMiVKUZrM3sho= -github.com/milvus-io/milvus-storage/go v0.0.0-20231226033437-76e506e3ae48 h1:EXDWA9yjmLLjIlIFjTdwtA3p1G0FDJdT07QdgCAWFWU= -github.com/milvus-io/milvus-storage/go v0.0.0-20231226033437-76e506e3ae48/go.mod h1:GPETMcTZq1gLY1WA6Na5kiNAKnq8SEMMiVKUZrM3sho= -github.com/milvus-io/milvus-storage/go v0.0.0-20231226075239-137cb5c55a5f h1:l43tW6aahbKcatIsX2X1guQktWSv/wgCBcGhmMWJgTg= -github.com/milvus-io/milvus-storage/go v0.0.0-20231226075239-137cb5c55a5f/go.mod h1:GPETMcTZq1gLY1WA6Na5kiNAKnq8SEMMiVKUZrM3sho= -github.com/milvus-io/milvus-storage/go v0.0.0-20231226081638-4a9a35e739b6 h1:v8WP0xJoOFno/YKdTrVfjWNn/VBmRX4IirK3/dhtH+8= -github.com/milvus-io/milvus-storage/go v0.0.0-20231226081638-4a9a35e739b6/go.mod h1:GPETMcTZq1gLY1WA6Na5kiNAKnq8SEMMiVKUZrM3sho= -github.com/milvus-io/milvus-storage/go v0.0.0-20231226083239-422d03dd1e1c h1:Xnc1Jt4joXVu2OsZp3xNZYQ/rKptRfRzYIHNaZkCpF8= -github.com/milvus-io/milvus-storage/go v0.0.0-20231226083239-422d03dd1e1c/go.mod h1:GPETMcTZq1gLY1WA6Na5kiNAKnq8SEMMiVKUZrM3sho= -github.com/milvus-io/milvus-storage/go v0.0.0-20231226085237-57519406e94f h1:4qnOXYGDVXdbIWUp9tk+JYtQ58QKf5d8q+XVk9+UVXo= -github.com/milvus-io/milvus-storage/go v0.0.0-20231226085237-57519406e94f/go.mod h1:GPETMcTZq1gLY1WA6Na5kiNAKnq8SEMMiVKUZrM3sho= github.com/milvus-io/milvus-storage/go v0.0.0-20231227072638-ebd0b8e56d70 h1:Z+sp64fmAOxAG7mU0dfVOXvAXlwRB0c8a96rIM5HevI= github.com/milvus-io/milvus-storage/go v0.0.0-20231227072638-ebd0b8e56d70/go.mod h1:GPETMcTZq1gLY1WA6Na5kiNAKnq8SEMMiVKUZrM3sho= github.com/milvus-io/pulsar-client-go v0.6.10 h1:eqpJjU+/QX0iIhEo3nhOqMNXL+TyInAs1IAHZCrCM/A= diff --git a/internal/proxy/impl.go b/internal/proxy/impl.go index 597f7cd455..2ac93fa311 100644 --- a/internal/proxy/impl.go +++ b/internal/proxy/impl.go @@ -4700,6 +4700,22 @@ func (node *Proxy) OperatePrivilege(ctx context.Context, req *milvuspb.OperatePr log.Warn("fail to operate privilege", zap.Error(err)) return merr.Status(err), nil } + relatedPrivileges := util.RelatedPrivileges[util.PrivilegeNameForMetastore(req.Entity.Grantor.Privilege.Name)] + if len(relatedPrivileges) != 0 { + for _, relatedPrivilege := range relatedPrivileges { + relatedReq := proto.Clone(req).(*milvuspb.OperatePrivilegeRequest) + relatedReq.Entity.Grantor.Privilege.Name = util.PrivilegeNameForAPI(relatedPrivilege) + result, err = node.rootCoord.OperatePrivilege(ctx, relatedReq) + if err != nil { + log.Warn("fail to operate related privilege", zap.String("related_privilege", relatedPrivilege), zap.Error(err)) + return merr.Status(err), nil + } + if !merr.Ok(result) { + log.Warn("fail to operate related privilege", zap.String("related_privilege", relatedPrivilege), zap.Any("result", result)) + return result, nil + } + } + } return result, nil } diff --git a/internal/proxy/proxy_test.go b/internal/proxy/proxy_test.go index 98435d0568..438a0663ff 100644 --- a/internal/proxy/proxy_test.go +++ b/internal/proxy/proxy_test.go @@ -53,6 +53,7 @@ import ( grpcquerynode "github.com/milvus-io/milvus/internal/distributed/querynode" grpcrootcoord "github.com/milvus-io/milvus/internal/distributed/rootcoord" rcc "github.com/milvus-io/milvus/internal/distributed/rootcoord/client" + "github.com/milvus-io/milvus/internal/mocks" "github.com/milvus-io/milvus/internal/proto/internalpb" "github.com/milvus-io/milvus/internal/proto/proxypb" "github.com/milvus-io/milvus/internal/proto/querypb" @@ -4616,6 +4617,55 @@ func TestProxy_ListImportTasks(t *testing.T) { }) } +func TestProxy_RelatedPrivilege(t *testing.T) { + req := &milvuspb.OperatePrivilegeRequest{ + Entity: &milvuspb.GrantEntity{ + Role: &milvuspb.RoleEntity{Name: "public"}, + ObjectName: "col1", + Object: &milvuspb.ObjectEntity{Name: commonpb.ObjectType_Collection.String()}, + Grantor: &milvuspb.GrantorEntity{Privilege: &milvuspb.PrivilegeEntity{Name: util.MetaStore2API(commonpb.ObjectPrivilege_PrivilegeLoad.String())}}, + }, + } + ctx := GetContext(context.Background(), "root:123456") + + t.Run("related privilege grpc error", func(t *testing.T) { + rootCoord := mocks.NewMockRootCoordClient(t) + proxy := &Proxy{rootCoord: rootCoord} + proxy.UpdateStateCode(commonpb.StateCode_Healthy) + + rootCoord.EXPECT().OperatePrivilege(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, request *milvuspb.OperatePrivilegeRequest, option ...grpc.CallOption) (*commonpb.Status, error) { + privilegeName := request.Entity.Grantor.Privilege.Name + if privilegeName == util.MetaStore2API(commonpb.ObjectPrivilege_PrivilegeLoad.String()) { + return merr.Success(), nil + } + return nil, errors.New("mock grpc error") + }) + + resp, err := proxy.OperatePrivilege(ctx, req) + assert.NoError(t, err) + assert.False(t, merr.Ok(resp)) + }) + + t.Run("related privilege status error", func(t *testing.T) { + rootCoord := mocks.NewMockRootCoordClient(t) + proxy := &Proxy{rootCoord: rootCoord} + proxy.UpdateStateCode(commonpb.StateCode_Healthy) + + rootCoord.EXPECT().OperatePrivilege(mock.Anything, mock.Anything).RunAndReturn(func(ctx context.Context, request *milvuspb.OperatePrivilegeRequest, option ...grpc.CallOption) (*commonpb.Status, error) { + privilegeName := request.Entity.Grantor.Privilege.Name + if privilegeName == util.MetaStore2API(commonpb.ObjectPrivilege_PrivilegeLoad.String()) || + privilegeName == util.MetaStore2API(commonpb.ObjectPrivilege_PrivilegeGetLoadState.String()) { + return merr.Success(), nil + } + return merr.Status(errors.New("mock status error")), nil + }) + + resp, err := proxy.OperatePrivilege(ctx, req) + assert.NoError(t, err) + assert.False(t, merr.Ok(resp)) + }) +} + func TestProxy_GetStatistics(t *testing.T) { } diff --git a/pkg/util/constant.go b/pkg/util/constant.go index 95011d8aca..50a6ac3182 100644 --- a/pkg/util/constant.go +++ b/pkg/util/constant.go @@ -142,6 +142,16 @@ var ( MetaStore2API(commonpb.ObjectPrivilege_PrivilegeSelectUser.String()), }, } + + RelatedPrivileges = map[string][]string{ + commonpb.ObjectPrivilege_PrivilegeLoad.String(): { + commonpb.ObjectPrivilege_PrivilegeGetLoadState.String(), + commonpb.ObjectPrivilege_PrivilegeGetLoadingProgress.String(), + }, + commonpb.ObjectPrivilege_PrivilegeFlush.String(): { + commonpb.ObjectPrivilege_PrivilegeGetFlushState.String(), + }, + } ) // StringSet convert array to map for conveniently check if the array contains an element