grant/revoke v2 optional db and collection params

Signed-off-by: shaoting-huang <[email protected]>
milvus-io · Nov 28, 2024 · d16b340 · d16b340
1 parent 807a6cc
commit d16b340
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 10 deletions.
diff --git a/client/client.go b/client/client.go
@@ -223,9 +223,9 @@ type Client interface {
 	// Revoke removes privilege from role.
 	Revoke(ctx context.Context, role string, objectType entity.PriviledgeObjectType, object string, privilege string, options ...entity.OperatePrivilegeOption) error
 	// GrantV2 adds privilege for role.
-	GrantV2(ctx context.Context, role string, privilege string, dbName string, colName string) error
+	GrantV2(ctx context.Context, role string, privilege string, options ...entity.OperatePrivilegeOption) error
 	// RevokeV2 removes privilege from role.
-	RevokeV2(ctx context.Context, role string, privilege string, dbName string, colName string) error
+	RevokeV2(ctx context.Context, role string, privilege string, options ...entity.OperatePrivilegeOption) error
 
 	// GetLoadingProgress get the collection or partitions loading progress
 	GetLoadingProgress(ctx context.Context, collectionName string, partitionNames []string) (int64, error)

diff --git a/client/rbac.go b/client/rbac.go
@@ -394,10 +394,14 @@ func (c *GrpcClient) Revoke(ctx context.Context, role string, objectType entity.
 }
 
 // GrantV2 adds object privilege for role without object type
-func (c *GrpcClient) GrantV2(ctx context.Context, role string, privilege string, dbName string, colName string) error {
+func (c *GrpcClient) GrantV2(ctx context.Context, role string, privilege string, options ...entity.OperatePrivilegeOption) error {
 	if c.Service == nil {
 		return ErrClientNotReady
 	}
+	grantOpt := &entity.OperatePrivilegeOpt{}
+	for _, opt := range options {
+		opt(grantOpt)
+	}
 
 	req := &milvuspb.OperatePrivilegeV2Request{
 		Role: &milvuspb.RoleEntity{
@@ -409,8 +413,8 @@ func (c *GrpcClient) GrantV2(ctx context.Context, role string, privilege string,
 			},
 		},
 		Type:           milvuspb.OperatePrivilegeType_Grant,
-		DbName:         dbName,
-		CollectionName: colName,
+		DbName:         grantOpt.Database,
+		CollectionName: grantOpt.Collection,
 	}
 
 	resp, err := c.Service.OperatePrivilegeV2(ctx, req)
@@ -422,10 +426,14 @@ func (c *GrpcClient) GrantV2(ctx context.Context, role string, privilege string,
 }
 
 // Revoke removes privilege from role without object type
-func (c *GrpcClient) RevokeV2(ctx context.Context, role string, privilege string, dbName string, colName string) error {
+func (c *GrpcClient) RevokeV2(ctx context.Context, role string, privilege string, options ...entity.OperatePrivilegeOption) error {
 	if c.Service == nil {
 		return ErrClientNotReady
 	}
+	revokeOpt := &entity.OperatePrivilegeOpt{}
+	for _, opt := range options {
+		opt(revokeOpt)
+	}
 
 	req := &milvuspb.OperatePrivilegeV2Request{
 		Role: &milvuspb.RoleEntity{
@@ -437,8 +445,8 @@ func (c *GrpcClient) RevokeV2(ctx context.Context, role string, privilege string
 			},
 		},
 		Type:           milvuspb.OperatePrivilegeType_Revoke,
-		DbName:         dbName,
-		CollectionName: colName,
+		DbName:         revokeOpt.Database,
+		CollectionName: revokeOpt.Collection,
 	}
 
 	resp, err := c.Service.OperatePrivilegeV2(ctx, req)

diff --git a/entity/rbac.go b/entity/rbac.go
@@ -44,8 +44,9 @@ const (
 )
 
 type OperatePrivilegeOpt struct {
-	Base     *common.MsgBase
-	Database string
+	Base       *common.MsgBase
+	Database   string
+	Collection string
 }
 
 type OperatePrivilegeOption func(o *OperatePrivilegeOpt)