Skip to content

Commit

Permalink
enhance: format list privilege group result (milvus-io#2372)
Browse files Browse the repository at this point in the history
issue: milvus-io/milvus#37031

Signed-off-by: shaoting-huang <[email protected]>
  • Loading branch information
shaoting-huang authored Nov 26, 2024
1 parent b7c6fd2 commit a086a9d
Show file tree
Hide file tree
Showing 4 changed files with 61 additions and 10 deletions.
3 changes: 2 additions & 1 deletion pymilvus/client/grpc_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -51,6 +51,7 @@
IndexState,
LoadState,
Plan,
PrivilegeGroupInfo,
Replica,
ResourceGroupConfig,
ResourceGroupInfo,
Expand Down Expand Up @@ -2079,7 +2080,7 @@ def list_privilege_groups(self, timeout: Optional[float] = None, **kwargs):
req = Prepare.list_privilege_groups_req()
resp = self._stub.ListPrivilegeGroups(req, wait_for_ready=True, timeout=timeout)
check_status(resp.status)
return resp.privilege_groups
return PrivilegeGroupInfo(resp.privilege_groups)

@retry_on_rpc_failure()
def add_privileges_to_group(
Expand Down
45 changes: 45 additions & 0 deletions pymilvus/client/types.py
Original file line number Diff line number Diff line change
Expand Up @@ -693,6 +693,51 @@ def groups(self):
return self._groups


class PrivilegeGroupItem:
def __init__(self, privilege_group: str, privileges: List[milvus_types.PrivilegeEntity]):
self._privilege_group = privilege_group
privielges = []
for privilege in privileges:
if isinstance(privilege, milvus_types.PrivilegeEntity):
privielges.append(privilege.name)
self._privileges = tuple(privielges)

def __repr__(self) -> str:
return f"PrivilegeGroupItem: <privilege_group:{self.privilege_group}>, <privileges:{self.privileges}>"

@property
def privilege_group(self):
return self._privilege_group

@property
def privileges(self):
return self._privileges


class PrivilegeGroupInfo:
"""
PrivilegeGroupInfo groups:
- PrivilegeGroupItem: <privilge_group:group>, <privileges:('Load', 'CreateCollection')>
"""

def __init__(self, results: List[milvus_types.PrivilegeGroupInfo]) -> None:
groups = []
for result in results:
if isinstance(result, milvus_types.PrivilegeGroupInfo):
groups.append(PrivilegeGroupItem(result.group_name, result.privileges))
self._groups = groups

def __repr__(self) -> str:
s = "PrivilegeGroupInfo groups:"
for g in self.groups:
s += f"\n- {g}"
return s

@property
def groups(self):
return self._groups


class UserItem:
def __init__(self, username: str, entities: List[milvus_types.RoleEntity]) -> None:
self._username = username
Expand Down
17 changes: 11 additions & 6 deletions pymilvus/milvus_client/milvus_client.py
Original file line number Diff line number Diff line change
Expand Up @@ -1210,7 +1210,7 @@ def list_privilege_groups(
self,
timeout: Optional[float] = None,
**kwargs,
) -> Dict[str, List[str]]:
) -> List[Dict[str, str]]:
"""List all privilege groups.
Args:
Expand All @@ -1219,16 +1219,20 @@ def list_privilege_groups(
or error occur.
Returns:
Dict[str, List[str]]: A dictionary of privilege groups and their privileges.
List[Dict[str, str]]: A list of privilege groups.
Raises:
MilvusException: If anything goes wrong.
"""
conn = self._get_connection()
pgs = conn.list_privilege_groups(timeout=timeout, **kwargs)
ret = {}
for pg in pgs:
ret[pg.group_name] = [p.name for p in pg.privileges]
try:
res = conn.list_privilege_groups(timeout=timeout, **kwargs)
except Exception as ex:
logger.exception("Failed to list privilege groups.")
raise ex from ex
ret = []
for g in res.groups:
ret.append({"privilge_group": g.privilege_group, "privileges": g.privileges})
return ret

def add_privileges_to_group(
Expand All @@ -1243,6 +1247,7 @@ def add_privileges_to_group(
Args:
group_name (``str``): The name of the privilege group.
privileges (``List[str]``): A list of privileges to be added to the group.
Privilges should be the same type in a group otherwise it will raise an exception.
timeout (``float``, optional): An optional duration of time in seconds to allow
for the RPC. When timeout is set to None, client waits until server response
or error occur.
Expand Down
6 changes: 3 additions & 3 deletions pymilvus/orm/role.py
Original file line number Diff line number Diff line change
Expand Up @@ -296,7 +296,7 @@ def list_privilege_groups(self):
:rtype PrivilegeGroupInfo
PrivilegeGroupInfo groups:
- PrivilegeGroupItem: <group_name:group1>, <privileges:['Insert', 'Select']>
- PrivilegeGroupItem: <group_name:group1>, <privileges:['Insert', 'Release']>
:example:
>>> from pymilvus import connections
Expand All @@ -320,7 +320,7 @@ def add_privileges_to_group(self, privilege_group: str, privileges: list):
>>> connections.connect()
>>> role = Role(role_name)
>>> role.add_privileges_to_group(privilege_group="privilege_group",
>>> privileges=["Insert","Select"])
>>> privileges=["Insert","Release"])
"""
return self._get_connection().add_privileges_to_group(privilege_group, privileges)

Expand All @@ -337,6 +337,6 @@ def remove_privileges_from_group(self, privilege_group: str, privileges: list):
>>> connections.connect()
>>> role = Role(role_name)
>>> role.remove_privileges_from_group(privilege_group="privilege_group",
>>> privileges=["Insert","Select"])
>>> privileges=["Insert","Release"])
"""
return self._get_connection().remove_privileges_from_group(privilege_group, privileges)

0 comments on commit a086a9d

Please sign in to comment.