enhance: support related privilege for grant api (#30153)

/kind improvement

Signed-off-by: SimFG <bang.fu@zilliz.com>

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=

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
 }
 

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) {
 }
 

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