fix: [2.4] optional db for grant/revoke v2 (#856)

cherry-pick from master: #852 issue: milvus-io/milvus#37031 Signed-off-by: shaoting-huang <shaoting.huang@zilliz.com>
milvus-io · Dec 10, 2024 · d724dc8 · d724dc8
1 parent 48cb4a2
commit d724dc8
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 8 deletions.
diff --git a/client/client.go b/client/client.go
@@ -226,10 +226,10 @@ type Client interface {
 	Grant(ctx context.Context, role string, objectType entity.PriviledgeObjectType, object string, privilege string, options ...entity.OperatePrivilegeOption) error
 	// 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
-	// RevokeV2 removes privilege from role.
-	RevokeV2(ctx context.Context, role string, privilege string, dbName string, colName string) error
+	// GrantV2 adds privilege for role. It will use default database if the option is not provided.
+	GrantV2(ctx context.Context, role string, colName string, privilege string, options ...entity.OperatePrivilegeOption) error
+	// RevokeV2 removes privilege from role. It will use default database if the option is not provided.
+	RevokeV2(ctx context.Context, role string, colName 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,11 +394,15 @@ 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, colName 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{
 			Name: role,
@@ -409,7 +413,7 @@ func (c *GrpcClient) GrantV2(ctx context.Context, role string, privilege string,
 			},
 		},
 		Type:           milvuspb.OperatePrivilegeType_Grant,
-		DbName:         dbName,
+		DbName:         grantOpt.Database,
 		CollectionName: colName,
 	}
 
@@ -422,11 +426,15 @@ 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, colName 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{
 			Name: role,
@@ -437,7 +445,7 @@ func (c *GrpcClient) RevokeV2(ctx context.Context, role string, privilege string
 			},
 		},
 		Type:           milvuspb.OperatePrivilegeType_Revoke,
-		DbName:         dbName,
+		DbName:         revokeOpt.Database,
 		CollectionName: colName,
 	}