From d16b3407f3624a7fca9bab9eb30fdff434afd53e Mon Sep 17 00:00:00 2001 From: shaoting-huang Date: Thu, 28 Nov 2024 16:14:56 +0800 Subject: [PATCH] grant/revoke v2 optional db and collection params Signed-off-by: shaoting-huang --- client/client.go | 4 ++-- client/rbac.go | 20 ++++++++++++++------ entity/rbac.go | 5 +++-- 3 files changed, 19 insertions(+), 10 deletions(-) diff --git a/client/client.go b/client/client.go index 23b4ce1e..29e1061d 100644 --- 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 index 462d7012..35f7a634 100644 --- 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 index c16d6370..3e7a30aa 100644 --- 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)