Skip to content

Commit

Permalink
grant/revoke v2 optional db and collection params
Browse files Browse the repository at this point in the history
Signed-off-by: shaoting-huang <[email protected]>
  • Loading branch information
shaoting-huang committed Nov 28, 2024
1 parent 3449513 commit e010cd0
Show file tree
Hide file tree
Showing 4 changed files with 23 additions and 20 deletions.
8 changes: 4 additions & 4 deletions pymilvus/client/grpc_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -1876,8 +1876,8 @@ def grant_privilege_v2(
self,
role_name: str,
privilege: str,
db_name: str,
collection_name: str,
db_name: Optional[str] = None,
collection_name: Optional[str] = None,
timeout: Optional[float] = None,
**kwargs,
):
Expand All @@ -1896,8 +1896,8 @@ def revoke_privilege_v2(
self,
role_name: str,
privilege: str,
db_name: str,
collection_name: str,
db_name: Optional[str] = None,
collection_name: Optional[str] = None,
timeout: Optional[float] = None,
**kwargs,
):
Expand Down
6 changes: 4 additions & 2 deletions pymilvus/client/prepare.py
Original file line number Diff line number Diff line change
Expand Up @@ -1484,9 +1484,11 @@ def operate_privilege_v2_request(
role_name=role_name,
privilege=privilege,
operate_privilege_type=operate_privilege_type,
db_name=db_name,
collection_name=collection_name,
)
if db_name:
check_pass_param(db_name=db_name)
if collection_name:
check_pass_param(collection_name=collection_name)
return milvus_types.OperatePrivilegeV2Request(
role=milvus_types.RoleEntity(name=role_name),
grantor=milvus_types.GrantorEntity(
Expand Down
12 changes: 6 additions & 6 deletions pymilvus/milvus_client/milvus_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -1008,28 +1008,28 @@ def grant_privilege_v2(
self,
role_name: str,
privilege: str,
db_name: str,
collection_name: str,
db_name: Optional[str] = None,
collection_name: Optional[str] = None,
timeout: Optional[float] = None,
**kwargs,
):
conn = self._get_connection()
conn.grant_privilege_v2(
role_name, privilege, db_name, collection_name, timeout=timeout, **kwargs
role_name, privilege, db_name=db_name, collection_name=collection_name, timeout=timeout, **kwargs
)

def revoke_privilege_v2(
self,
role_name: str,
privilege: str,
db_name: str,
collection_name: str,
db_name: Optional[str] = None,
collection_name: Optional[str] = None,
timeout: Optional[float] = None,
**kwargs,
):
conn = self._get_connection()
conn.revoke_privilege_v2(
role_name, privilege, db_name, collection_name, timeout=timeout, **kwargs
role_name, privilege, db_name=db_name, collection_name=collection_name, timeout=timeout, **kwargs
)

def create_alias(
Expand Down
17 changes: 9 additions & 8 deletions pymilvus/orm/role.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
# the License.

from .connections import connections
from typing import Optional

INCLUDE_USER_INFO, NOT_INCLUDE_USER_INFO = True, False

Expand Down Expand Up @@ -176,13 +177,13 @@ def revoke(self, object: str, object_name: str, privilege: str, db_name: str = "
self._name, object, object_name, privilege, db_name
)

def grant_v2(self, privilege: str, db_name: str, collection_name: str):
def grant_v2(self, privilege: str, db_name: Optional[str] = None, collection_name: Optional[str] = None):
"""Grant a privilege for the role
:param privilege: privilege name.
:type privilege: str
:param db_name: db name.
:param db_name: db name. Optional
:type db_name: str
:param collection_name: collection name.
:param collection_name: collection name. Optional
:type collection_name: str
:example:
Expand All @@ -193,16 +194,16 @@ def grant_v2(self, privilege: str, db_name: str, collection_name: str):
>>> role.grant_v2("Insert", db_name, collection_name)
"""
return self._get_connection().grant_privilege_v2(
self._name, privilege, db_name, collection_name
self._name, privilege, db_name=db_name, collection_name=collection_name
)

def revoke_v2(self, privilege: str, db_name: str, collection_name: str):
def revoke_v2(self, privilege: str, db_name: Optional[str] = None, collection_name: Optional[str] = None):
"""Revoke a privilege for the role
:param privilege: privilege name.
:type privilege: str
:param db_name: db name.
:param db_name: db name. Optional
:type db_name: str
:param collection_name: collection name.
:param collection_name: collection name. Optional
:type collection_name: str
:example:
Expand All @@ -213,7 +214,7 @@ def revoke_v2(self, privilege: str, db_name: str, collection_name: str):
>>> role.revoke_v2("Insert", db_name, collection_name)
"""
return self._get_connection().revoke_privilege_v2(
self._name, privilege, db_name, collection_name
self._name, privilege, db_name=db_name, collection_name=collection_name
)

def list_grant(self, object: str, object_name: str, db_name: str = ""):
Expand Down

0 comments on commit e010cd0

Please sign in to comment.