From c6b391ccd56431640276093d4cef7631de019c9e Mon Sep 17 00:00:00 2001
From: Alex Giurgiu <alex@giurgiu.io>
Date: Tue, 9 Jul 2024 18:09:39 +0300
Subject: [PATCH] fix grant creation by adding the actual privilege

---
 client/client.go    | 2 +-
 client/rbac.go      | 7 ++++++-
 client/rbac_test.go | 9 +++++----
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/client/client.go b/client/client.go
index ec3765d58..1728926e8 100644
--- a/client/client.go
+++ b/client/client.go
@@ -202,7 +202,7 @@ type Client interface {
 	// ListGrants lists all assigned privileges and objects for the role.
 	ListGrants(ctx context.Context, role string, dbName string) ([]entity.RoleGrants, error)
 	// Grant adds privilege for role.
-	Grant(ctx context.Context, role string, objectType entity.PriviledgeObjectType, object string) error
+	Grant(ctx context.Context, role string, objectType entity.PriviledgeObjectType, object string, privilege string) error
 	// Revoke removes privilege from role.
 	Revoke(ctx context.Context, role string, objectType entity.PriviledgeObjectType, object string) error
 
diff --git a/client/rbac.go b/client/rbac.go
index 2e96bede6..a94bfc05f 100644
--- a/client/rbac.go
+++ b/client/rbac.go
@@ -320,7 +320,7 @@ func (c *GrpcClient) ListGrant(ctx context.Context, role string, object string,
 }
 
 // Grant adds object privileged for role.
-func (c *GrpcClient) Grant(ctx context.Context, role string, objectType entity.PriviledgeObjectType, object string) error {
+func (c *GrpcClient) Grant(ctx context.Context, role string, objectType entity.PriviledgeObjectType, object string, privilege string) error {
 	if c.Service == nil {
 		return ErrClientNotReady
 	}
@@ -333,6 +333,11 @@ func (c *GrpcClient) Grant(ctx context.Context, role string, objectType entity.P
 			Object: &milvuspb.ObjectEntity{
 				Name: commonpb.ObjectType_name[int32(objectType)],
 			},
+			Grantor: &milvuspb.GrantorEntity{
+				Privilege: &milvuspb.PrivilegeEntity{
+					Name: privilege,
+				},
+			},
 			ObjectName: object,
 		},
 		Type: milvuspb.OperatePrivilegeType_Grant,
diff --git a/client/rbac_test.go b/client/rbac_test.go
index 51f230d15..2e6673d9d 100644
--- a/client/rbac_test.go
+++ b/client/rbac_test.go
@@ -629,6 +629,7 @@ func (s *RBACSuite) TestGrant() {
 	roleName := "testRole"
 	objectName := testCollectionName
 	objectType := entity.PriviledegeObjectTypeCollection
+	privilegeName := "testPrivilege"
 
 	s.Run("normal run", func() {
 		ctx, cancel := context.WithCancel(ctx)
@@ -641,7 +642,7 @@ func (s *RBACSuite) TestGrant() {
 			s.Equal(milvuspb.OperatePrivilegeType_Grant, req.GetType())
 		}).Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_Success}, nil)
 
-		err := s.client.Grant(ctx, roleName, objectType, objectName)
+		err := s.client.Grant(ctx, roleName, objectType, objectName, privilegeName)
 
 		s.NoError(err)
 	})
@@ -652,7 +653,7 @@ func (s *RBACSuite) TestGrant() {
 		defer s.resetMock()
 		s.mock.EXPECT().OperatePrivilege(mock.Anything, mock.Anything).Return(nil, errors.New("mock error"))
 
-		err := s.client.Grant(ctx, roleName, objectType, objectName)
+		err := s.client.Grant(ctx, roleName, objectType, objectName, privilegeName)
 		s.Error(err)
 	})
 
@@ -662,7 +663,7 @@ func (s *RBACSuite) TestGrant() {
 		defer s.resetMock()
 		s.mock.EXPECT().OperatePrivilege(mock.Anything, mock.Anything).Return(&commonpb.Status{ErrorCode: commonpb.ErrorCode_UnexpectedError}, nil)
 
-		err := s.client.Grant(ctx, roleName, objectType, objectName)
+		err := s.client.Grant(ctx, roleName, objectType, objectName, privilegeName)
 		s.Error(err)
 	})
 
@@ -671,7 +672,7 @@ func (s *RBACSuite) TestGrant() {
 		defer cancel()
 
 		c := &GrpcClient{}
-		err := c.Grant(ctx, roleName, objectType, objectName)
+		err := c.Grant(ctx, roleName, objectType, objectName, privilegeName)
 		s.Error(err)
 		s.ErrorIs(err, ErrClientNotReady)
 	})