Skip to content

Commit

Permalink
enhance: Add Describe User methods
Browse files Browse the repository at this point in the history
Signed-off-by: punkerpunker <[email protected]>
  • Loading branch information
punkerpunker committed Apr 5, 2024
1 parent 6d2bfc9 commit dd101c1
Show file tree
Hide file tree
Showing 4 changed files with 246 additions and 4 deletions.
6 changes: 5 additions & 1 deletion client/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -186,7 +186,11 @@ type Client interface {
ListRoles(ctx context.Context) ([]entity.Role, error)
// ListUsers lists the user objects in system.
ListUsers(ctx context.Context) ([]entity.User, error)
// Grant adds object privileged for role.
// DescribeUser describes specific user attributes in the system
DescribeUser(ctx context.Context, username string) (entity.UserDescription, error)
// DescribeUsers describe all users attributes in the system
DescribeUsers(ctx context.Context) ([]entity.UserDescription, error)
// Grant adds privilege for role.
Grant(ctx context.Context, role string, objectType entity.PriviledgeObjectType, object string) error
// Revoke removes privilege from role.
Revoke(ctx context.Context, role string, objectType entity.PriviledgeObjectType, object string) error
Expand Down
76 changes: 74 additions & 2 deletions client/rbac.go
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@ package client

import (
"context"

Check failure on line 20 in client/rbac.go

View workflow job for this annotation

GitHub Actions / lint

File is not `goimports`-ed (goimports)

"github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
"github.com/milvus-io/milvus-proto/go-api/v2/milvuspb"
"github.com/milvus-io/milvus-sdk-go/v2/entity"
Expand Down Expand Up @@ -97,7 +96,6 @@ func (c *GrpcClient) RemoveUserRole(ctx context.Context, username string, role s
if err != nil {
return err
}

return handleRespStatus(resp)
}

Expand Down Expand Up @@ -149,6 +147,80 @@ func (c *GrpcClient) ListUsers(ctx context.Context) ([]entity.User, error) {
return users, nil
}

// DescribeUser lists the user descriptions in the system (name, roles)
func (c *GrpcClient) DescribeUser(ctx context.Context, username string) (entity.UserDescription, error) {
if c.Service == nil {
return entity.UserDescription{}, ErrClientNotReady
}

req := &milvuspb.SelectUserRequest{
User: &milvuspb.UserEntity{
Name: username,
},
IncludeRoleInfo: true,
}

resp, err := c.Service.SelectUser(ctx, req)

if err != nil {
return entity.UserDescription{}, err
}
if err = handleRespStatus(resp.GetStatus()); err != nil {
return entity.UserDescription{}, err
}
results := resp.GetResults()

if len(results) == 0 {
return entity.UserDescription{}, nil
}

userDescription := entity.UserDescription{
Name: results[0].GetUser().GetName(),
Roles: make([]string, 0, len(results[0].GetRoles())),
}

for _, role := range results[0].GetRoles() {
userDescription.Roles = append(userDescription.Roles, role.GetName())
}
return userDescription, nil
}

// DescribeUsers lists all users with descriptions (names, roles)
func (c *GrpcClient) DescribeUsers(ctx context.Context) ([]entity.UserDescription, error) {
if c.Service == nil {
return nil, ErrClientNotReady
}

req := &milvuspb.SelectUserRequest{
IncludeRoleInfo: true,
}

resp, err := c.Service.SelectUser(ctx, req)

if err != nil {
return nil, err
}
if err = handleRespStatus(resp.GetStatus()); err != nil {
return nil, err
}
results := resp.GetResults()

userDescriptions := make([]entity.UserDescription, 0, len(results))

for _, result := range results {
userDescription := entity.UserDescription{
Name: result.GetUser().GetName(),
Roles: make([]string, 0, len(result.GetRoles())),
}
for _, role := range result.GetRoles() {
userDescription.Roles = append(userDescription.Roles, role.GetName())
}
userDescriptions = append(userDescriptions, userDescription)
}

return userDescriptions, nil
}

// Grant adds object privileged for role.
func (c *GrpcClient) Grant(ctx context.Context, role string, objectType entity.PriviledgeObjectType, object string) error {
if c.Service == nil {
Expand Down
158 changes: 158 additions & 0 deletions client/rbac_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -347,6 +347,164 @@ func (s *RBACSuite) TestListUser() {
})
}

func (s *RBACSuite) TestDescribeUser() {
ctx := context.Background()
userName := "testUser"

s.Run("normal run", func() {
ctx, cancel := context.WithCancel(ctx)
defer cancel()
defer s.resetMock()
s.mock.EXPECT().SelectUser(mock.Anything, mock.Anything).Run(func(ctx context.Context, req *milvuspb.SelectUserRequest) {
s.True(req.GetIncludeRoleInfo())
s.Equal(req.GetUser().GetName(), userName)
}).Return(&milvuspb.SelectUserResponse{
Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
Results: []*milvuspb.UserResult{
{
User: &milvuspb.UserEntity{
Name: userName,
},
Roles: []*milvuspb.RoleEntity{
{Name: "role1"},
{Name: "role2"},
},
},
},
}, nil)

userDesc, err := s.client.DescribeUser(ctx, userName)

s.NoError(err)
s.Equal(userDesc.Name, userName)
s.ElementsMatch(userDesc.Roles, []string{"role1", "role2"})
})

s.Run("rpc error", func() {
ctx, cancel := context.WithCancel(ctx)
defer cancel()
defer s.resetMock()
s.mock.EXPECT().SelectUser(mock.Anything, mock.Anything).Return(nil, errors.New("mock error"))

_, err := s.client.DescribeUser(ctx, userName)
s.Error(err)
})

s.Run("status error", func() {
ctx, cancel := context.WithCancel(ctx)
defer cancel()
defer s.resetMock()
s.mock.EXPECT().SelectUser(mock.Anything, mock.Anything).Return(&milvuspb.SelectUserResponse{
Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_UnexpectedError},
}, nil)

_, err := s.client.DescribeUser(ctx, userName)
s.Error(err)
})

s.Run("service not ready", func() {
ctx, cancel := context.WithCancel(ctx)
defer cancel()

c := &GrpcClient{}
_, err := c.DescribeUser(ctx, userName)
s.Error(err)
s.ErrorIs(err, ErrClientNotReady)
})
}

func (s *RBACSuite) TestDescribeUsers() {
ctx := context.Background()

s.Run("normal run", func() {
ctx, cancel := context.WithCancel(ctx)
defer cancel()
defer s.resetMock()

mockResults := []*milvuspb.UserResult{
{
User: &milvuspb.UserEntity{
Name: "user1",
},
Roles: []*milvuspb.RoleEntity{
{Name: "role1"},
{Name: "role2"},
},
},
{
User: &milvuspb.UserEntity{
Name: "user2",
},
Roles: []*milvuspb.RoleEntity{
{Name: "role3"},
},
},
}

s.mock.EXPECT().SelectUser(mock.Anything, mock.Anything).Run(func(ctx context.Context, req *milvuspb.SelectUserRequest) {
s.True(req.GetIncludeRoleInfo())
}).Return(&milvuspb.SelectUserResponse{
Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_Success},
Results: mockResults,
}, nil)

userDescs, err := s.client.DescribeUsers(ctx)

s.NoError(err)
s.Len(userDescs, 2)

expectedDescs := []entity.UserDescription{
{
Name: "user1",
Roles: []string{"role1", "role2"},
},
{
Name: "user2",
Roles: []string{"role3"},
},
}

for i, desc := range userDescs {
s.Equal(expectedDescs[i].Name, desc.Name)
s.ElementsMatch(expectedDescs[i].Roles, desc.Roles)
}
})

s.Run("rpc error", func() {
ctx, cancel := context.WithCancel(ctx)
defer cancel()
defer s.resetMock()

s.mock.EXPECT().SelectUser(mock.Anything, mock.Anything).Return(nil, errors.New("mock error"))

_, err := s.client.DescribeUsers(ctx)
s.Error(err)
})

s.Run("status error", func() {
ctx, cancel := context.WithCancel(ctx)
defer cancel()
defer s.resetMock()

s.mock.EXPECT().SelectUser(mock.Anything, mock.Anything).Return(&milvuspb.SelectUserResponse{
Status: &commonpb.Status{ErrorCode: commonpb.ErrorCode_UnexpectedError},
}, nil)

_, err := s.client.DescribeUsers(ctx)
s.Error(err)
})

s.Run("service not ready", func() {
ctx, cancel := context.WithCancel(ctx)
defer cancel()

c := &GrpcClient{}
_, err := c.DescribeUsers(ctx)
s.Error(err)
s.ErrorIs(err, ErrClientNotReady)
})
}

func (s *RBACSuite) TestGrant() {
ctx := context.Background()

Expand Down
10 changes: 9 additions & 1 deletion entity/rbac.go
Original file line number Diff line number Diff line change
@@ -1,12 +1,20 @@
package entity

import common "github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
import (
common "github.com/milvus-io/milvus-proto/go-api/v2/commonpb"
)

// User is the model for RBAC user object.
type User struct {
Name string
}

// UserDescription is the model for RBAC user description object.
type UserDescription struct {
Name string
Roles []string
}

// Role is the model for RBAC role object.
type Role struct {
Name string
Expand Down

0 comments on commit dd101c1

Please sign in to comment.