RBAC privielge group API

Signed-off-by: shaoting-huang <[email protected]>
milvus-io · Nov 13, 2024 · 88d525b · 88d525b
1 parent d472ccd
commit 88d525b
Show file tree

Hide file tree

Showing 3 changed files with 50 additions and 9 deletions.
diff --git a/client/client.go b/client/client.go
@@ -92,8 +92,10 @@ type Client interface {
 	DropPrivilegeGroup(ctx context.Context, groupName string) error
 	// ListPrivilegeGroups lists all privilege groups
 	ListPrivilegeGroups(ctx context.Context) ([]*entity.PrivilegeGroup, error)
-	// OperatePrivilegeGroup adds privileges to a privilege group or remove privileges from a privilege group
-	OperatePrivilegeGroup(ctx context.Context, groupName string, privileges []string) error
+	// AddPrivilegeToGroup adds privileges to a privilege group
+	AddPrivilegesToGroup(ctx context.Context, groupName string, privileges []string) error
+	// RemovePrivilegesFromGroup removes privileges from a privilege group
+	RemovePrivilegesFromGroup(ctx context.Context, groupName string, privileges []string) error
 
 	// -- authentication --
 

diff --git a/client/client_mock_test.go b/client/client_mock_test.go
@@ -256,6 +256,7 @@ const (
 	MAlterCollection         ServiceMethod = 109
 	MGetLoadingProgress      ServiceMethod = 110
 	MGetLoadState            ServiceMethod = 111
+	MAlterCollectionField    ServiceMethod = 112
 
 	MCreatePartition   ServiceMethod = 201
 	MDropPartition     ServiceMethod = 202
@@ -1085,6 +1086,15 @@ func (m *MockServer) RestoreRBAC(ctx context.Context, req *milvuspb.RestoreRBACM
 	return SuccessStatus()
 }
 
+func (m *MockServer) AlterCollectionField(ctx context.Context, req *milvuspb.AlterCollectionFieldRequest) (*commonpb.Status, error) {
+	f := m.GetInjection(MAlterCollectionField)
+	if f != nil {
+		r, err := f(ctx, req)
+		return r.(*commonpb.Status), err
+	}
+	return SuccessStatus()
+}
+
 func (m *MockServer) CreatePrivilegeGroup(ctx context.Context, req *milvuspb.CreatePrivilegeGroupRequest) (*commonpb.Status, error) {
 	f := m.GetInjection(MCreatePrivilegeGroup)
 	if f != nil {
@@ -1113,7 +1123,7 @@ func (m *MockServer) ListPrivilegeGroups(ctx context.Context, req *milvuspb.List
 	return &milvuspb.ListPrivilegeGroupsResponse{Status: s}, err
 }
 
-func (m *MockServer) OperatePrivilegeGroupRequest(ctx context.Context, req *milvuspb.OperatePrivilegeGroupRequest) (*commonpb.Status, error) {
+func (m *MockServer) OperatePrivilegeGroup(ctx context.Context, req *milvuspb.OperatePrivilegeGroupRequest) (*commonpb.Status, error) {
 	f := m.GetInjection(MOperatePrivilegeGroup)
 	if f != nil {
 		r, err := f(ctx, req)

diff --git a/client/rbac.go b/client/rbac.go
@@ -557,8 +557,8 @@ func (c *GrpcClient) DropPrivilegeGroup(ctx context.Context, groupName string) e
 	return handleRespStatus(resp)
 }
 
-func (c *GrpcClient) ListPrivilegeGroups(ctx context.Context) ([]entity.PrivilegeGroup, error) {
-	PrivilegeGroupList := make([]entity.PrivilegeGroup, 0)
+func (c *GrpcClient) ListPrivilegeGroups(ctx context.Context) ([]*entity.PrivilegeGroup, error) {
+	PrivilegeGroupList := make([]*entity.PrivilegeGroup, 0)
 	if c.Service == nil {
 		return PrivilegeGroupList, ErrClientNotReady
 	}
@@ -581,7 +581,7 @@ func (c *GrpcClient) ListPrivilegeGroups(ctx context.Context) ([]entity.Privileg
 	}
 
 	for _, pg := range results {
-		PrivilegeGroup := entity.PrivilegeGroup{
+		PrivilegeGroup := &entity.PrivilegeGroup{
 			GroupName: pg.GroupName,
 			Privileges: lo.Map(pg.Privileges, func(p *milvuspb.PrivilegeEntity, _ int) string {
 				return p.Name
@@ -593,16 +593,45 @@ func (c *GrpcClient) ListPrivilegeGroups(ctx context.Context) ([]entity.Privileg
 	return PrivilegeGroupList, nil
 }
 
-func (c *GrpcClient) OperatePrivilegeGroup(ctx context.Context, groupName string, privileges []string) error {
+func (c *GrpcClient) AddPrivilegesToGroup(ctx context.Context, groupName string, privileges []string) error {
 	if c.Service == nil {
 		return ErrClientNotReady
 	}
 
-	req := &milvuspb.DropPrivilegeGroupRequest{
+	req := &milvuspb.OperatePrivilegeGroupRequest{
 		GroupName: groupName,
+		Privileges: lo.Map(privileges, func(p string, _ int) *milvuspb.PrivilegeEntity {
+			return &milvuspb.PrivilegeEntity{
+				Name: p,
+			}
+		}),
+		Type: milvuspb.OperatePrivilegeGroupType_AddPrivilegesToGroup,
 	}
 
-	resp, err := c.Service.DropPrivilegeGroup(ctx, req)
+	resp, err := c.Service.OperatePrivilegeGroup(ctx, req)
+	if err != nil {
+		return err
+	}
+
+	return handleRespStatus(resp)
+}
+
+func (c *GrpcClient) RemovePrivilegesFromGroup(ctx context.Context, groupName string, privileges []string) error {
+	if c.Service == nil {
+		return ErrClientNotReady
+	}
+
+	req := &milvuspb.OperatePrivilegeGroupRequest{
+		GroupName: groupName,
+		Privileges: lo.Map(privileges, func(p string, _ int) *milvuspb.PrivilegeEntity {
+			return &milvuspb.PrivilegeEntity{
+				Name: p,
+			}
+		}),
+		Type: milvuspb.OperatePrivilegeGroupType_RemovePrivilegesFromGroup,
+	}
+
+	resp, err := c.Service.OperatePrivilegeGroup(ctx, req)
 	if err != nil {
 		return err
 	}