From a086a9d95f9534d0c8a4d3e9303d13860a22fd99 Mon Sep 17 00:00:00 2001 From: sthuang <167743503+shaoting-huang@users.noreply.github.com> Date: Tue, 26 Nov 2024 14:52:28 +0800 Subject: [PATCH] enhance: format list privilege group result (#2372) issue: https://github.com/milvus-io/milvus/issues/37031 Signed-off-by: shaoting-huang --- pymilvus/client/grpc_handler.py | 3 +- pymilvus/client/types.py | 45 +++++++++++++++++++++++++ pymilvus/milvus_client/milvus_client.py | 17 ++++++---- pymilvus/orm/role.py | 6 ++-- 4 files changed, 61 insertions(+), 10 deletions(-) diff --git a/pymilvus/client/grpc_handler.py b/pymilvus/client/grpc_handler.py index 9a9791a84..fb6383008 100644 --- a/pymilvus/client/grpc_handler.py +++ b/pymilvus/client/grpc_handler.py @@ -51,6 +51,7 @@ IndexState, LoadState, Plan, + PrivilegeGroupInfo, Replica, ResourceGroupConfig, ResourceGroupInfo, @@ -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( diff --git a/pymilvus/client/types.py b/pymilvus/client/types.py index 509fb6a39..8f0b8c629 100644 --- a/pymilvus/client/types.py +++ b/pymilvus/client/types.py @@ -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: , " + + @property + def privilege_group(self): + return self._privilege_group + + @property + def privileges(self): + return self._privileges + + +class PrivilegeGroupInfo: + """ + PrivilegeGroupInfo groups: + - PrivilegeGroupItem: , + """ + + 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 diff --git a/pymilvus/milvus_client/milvus_client.py b/pymilvus/milvus_client/milvus_client.py index dd2d29c78..2a58d00b4 100644 --- a/pymilvus/milvus_client/milvus_client.py +++ b/pymilvus/milvus_client/milvus_client.py @@ -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: @@ -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( @@ -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. diff --git a/pymilvus/orm/role.py b/pymilvus/orm/role.py index 989ef8c73..4e699d140 100644 --- a/pymilvus/orm/role.py +++ b/pymilvus/orm/role.py @@ -296,7 +296,7 @@ def list_privilege_groups(self): :rtype PrivilegeGroupInfo PrivilegeGroupInfo groups: - - PrivilegeGroupItem: , + - PrivilegeGroupItem: , :example: >>> from pymilvus import connections @@ -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) @@ -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)