fix: add privilege group orm (milvus-io#2356)

issue: milvus-io/milvus#37031 issue: milvus-io#2353 --------- Signed-off-by: shaoting-huang <[email protected]>
weiliu1031 · Nov 18, 2024 · 24cba21 · 24cba21
1 parent 016ff55
commit 24cba21
Show file tree

Hide file tree

Showing 4 changed files with 129 additions and 27 deletions.
diff --git a/pymilvus/client/check.py b/pymilvus/client/check.py
@@ -276,6 +276,25 @@ def is_legal_operate_privilege_type(operate_privilege_type: Any) -> bool:
     )
 
 
+def is_legal_privilege_group(privilege_group: Any) -> bool:
+    return privilege_group and isinstance(privilege_group, str)
+
+
+def is_legal_privileges(privileges: Any) -> bool:
+    return (
+        privileges
+        and isinstance(privileges, list)
+        and all(is_legal_privilege(p) for p in privileges)
+    )
+
+
+def is_legal_operate_privilege_group_type(operate_privilege_group_type: Any) -> bool:
+    return operate_privilege_group_type in (
+        milvus_types.OperatePrivilegeGroupType.AddPrivilegesToGroup,
+        milvus_types.OperatePrivilegeGroupType.RemovePrivilegesFromGroup,
+    )
+
+
 class ParamChecker(metaclass=Singleton):
     def __init__(self) -> None:
         self.check_dict = {
@@ -320,6 +339,9 @@ def __init__(self) -> None:
             "timeout": is_legal_timeout,
             "drop_ratio_build": is_legal_drop_ratio,
             "drop_ratio_search": is_legal_drop_ratio,
+            "privilege_group": is_legal_privilege_group,
+            "privileges": is_legal_privileges,
+            "operate_privilege_group_type": is_legal_operate_privilege_group_type,
         }
 
     def check(self, key: str, value: Callable):

diff --git a/pymilvus/client/grpc_handler.py b/pymilvus/client/grpc_handler.py
@@ -2023,14 +2023,16 @@ def alloc_timestamp(self, timeout: Optional[float] = None) -> int:
         return response.timestamp
 
     @retry_on_rpc_failure()
-    def create_privilege_group(self, group_name: str, timeout: Optional[float] = None, **kwargs):
-        req = Prepare.create_privilege_group_req(group_name)
+    def create_privilege_group(
+        self, privilege_group: str, timeout: Optional[float] = None, **kwargs
+    ):
+        req = Prepare.create_privilege_group_req(privilege_group)
         resp = self._stub.CreatePrivilegeGroup(req, wait_for_ready=True, timeout=timeout)
         check_status(resp)
 
     @retry_on_rpc_failure()
-    def drop_privilege_group(self, group_name: str, timeout: Optional[float] = None, **kwargs):
-        req = Prepare.drop_privilege_group_req(group_name)
+    def drop_privilege_group(self, privilege_group: str, timeout: Optional[float] = None, **kwargs):
+        req = Prepare.drop_privilege_group_req(privilege_group)
         resp = self._stub.DropPrivilegeGroup(req, wait_for_ready=True, timeout=timeout)
         check_status(resp)
 
@@ -2043,20 +2045,22 @@ def list_privilege_groups(self, timeout: Optional[float] = None, **kwargs):
 
     @retry_on_rpc_failure()
     def add_privileges_to_group(
-        self, group_name: str, privileges: List[str], timeout: Optional[float] = None, **kwargs
+        self, privilege_group: str, privileges: List[str], timeout: Optional[float] = None, **kwargs
     ):
         req = Prepare.operate_privilege_group_req(
-            group_name, privileges, milvus_types.OperatePrivilegeGroupType.AddPrivilegesToGroup
+            privilege_group, privileges, milvus_types.OperatePrivilegeGroupType.AddPrivilegesToGroup
         )
         resp = self._stub.OperatePrivilegeGroup(req, wait_for_ready=True, timeout=timeout)
         check_status(resp)
 
     @retry_on_rpc_failure()
     def remove_privileges_from_group(
-        self, group_name: str, privileges: List[str], timeout: Optional[float] = None, **kwargs
+        self, privilege_group: str, privileges: List[str], timeout: Optional[float] = None, **kwargs
     ):
         req = Prepare.operate_privilege_group_req(
-            group_name, privileges, milvus_types.OperatePrivilegeGroupType.RemovePrivilegesFromGroup
+            privilege_group,
+            privileges,
+            milvus_types.OperatePrivilegeGroupType.RemovePrivilegesFromGroup,
         )
         resp = self._stub.OperatePrivilegeGroup(req, wait_for_ready=True, timeout=timeout)
         check_status(resp)
diff --git a/pymilvus/client/prepare.py b/pymilvus/client/prepare.py
@@ -1590,33 +1590,28 @@ def describe_database_req(cls, db_name: str):
         return milvus_types.DescribeDatabaseRequest(db_name=db_name)
 
     @classmethod
-    def create_privilege_group_req(cls, group_name: str):
-        check_pass_param(group_name=group_name)
-        return milvus_types.CreatePrivilegeGroupRequest(group_name=group_name)
+    def create_privilege_group_req(cls, privilege_group: str):
+        check_pass_param(privilege_group=privilege_group)
+        return milvus_types.CreatePrivilegeGroupRequest(group_name=privilege_group)
 
     @classmethod
-    def drop_privilege_group_req(cls, group_name: str):
-        check_pass_param(group_name=group_name)
-        return milvus_types.DropPrivilegeGroupRequest(group_name=group_name)
+    def drop_privilege_group_req(cls, privilege_group: str):
+        check_pass_param(privilege_group=privilege_group)
+        return milvus_types.DropPrivilegeGroupRequest(group_name=privilege_group)
 
     @classmethod
     def list_privilege_groups_req(cls):
         return milvus_types.ListPrivilegeGroupsRequest()
 
     @classmethod
-    def operate_privilege_group_req(cls, group_name: str, privileges: List[str], operate_type: Any):
-        check_pass_param(group_name=group_name)
-        check_pass_param(operate_type=operate_type)
-        if not isinstance(
-            privileges,
-            (list),
-        ):
-            msg = f"Privileges {privileges} is not a list"
-            raise ParamError(message=msg)
-        for p in privileges:
-            check_pass_param(privilege=p)
+    def operate_privilege_group_req(
+        cls, privilege_group: str, privileges: List[str], operate_privilege_group_type: Any
+    ):
+        check_pass_param(privilege_group=privilege_group)
+        check_pass_param(privileges=privileges)
+        check_pass_param(operate_privilege_group_type=operate_privilege_group_type)
         return milvus_types.OperatePrivilegeGroupRequest(
-            group_name=group_name,
+            group_name=privilege_group,
             privileges=[milvus_types.PrivilegeEntity(name=p) for p in privileges],
-            type=operate_type,
+            type=operate_privilege_group_type,
         )
diff --git a/pymilvus/orm/role.py b/pymilvus/orm/role.py
@@ -221,3 +221,84 @@ def list_grants(self, db_name: str = ""):
             >>> role.list_grants()
         """
         return self._get_connection().select_grant_for_one_role(self._name, db_name)
+
+    def create_privilege_group(self, privilege_group: str):
+        """Create a privilege group for the role
+            :param privilege_group: privilege group name.
+            :type  privilege_group: str
+
+        :example:
+            >>> from pymilvus import connections
+            >>> from pymilvus.orm.role import Role
+            >>> connections.connect()
+            >>> role = Role(role_name)
+            >>> role.create_privilege_group(privilege_group)
+        """
+        return self._get_connection().create_privilege_group(self._name, privilege_group)
+
+    def drop_privilege_group(self, privilege_group: str):
+        """Drop a privilege group for the role
+            :param privilege_group: privilege group name.
+            :type  privilege_group: str
+
+        :example:
+            >>> from pymilvus import connections
+            >>> from pymilvus.orm.role import Role
+            >>> connections.connect()
+            >>> role = Role(role_name)
+            >>> role.drop_privilege_group(privilege_group)
+        """
+        return self._get_connection().drop_privilege_group(self._name, privilege_group)
+
+    def list_privilege_groups(self):
+        """List all privilege groups for the role
+            :return a PrivilegeGroupInfo object
+            :rtype PrivilegeGroupInfo
+
+            PrivilegeGroupInfo groups:
+            - PrivilegeGroupItem: <group_name:group1>, <privileges:['Insert', 'Select']>
+
+        :example:
+            >>> from pymilvus import connections
+            >>> from pymilvus.orm.role import Role
+            >>> connections.connect()
+            >>> role = Role(role_name)
+            >>> role.list_privilege_groups()
+        """
+        return self._get_connection().list_privilege_groups(self._name)
+
+    def add_privileges_to_group(self, privilege_group: str, privileges: list):
+        """Add privileges to a privilege group for the role
+            :param privilege_group: privilege group name.
+            :type  privilege_group: str
+            :param privileges: a list of privilege names.
+            :type  privileges: list
+
+        :example:
+            >>> from pymilvus import connections
+            >>> from pymilvus.orm.role import Role
+            >>> connections.connect()
+            >>> role = Role(role_name)
+            >>> role.add_privileges_to_group(privilege_group, ["Insert", "Select"])
+        """
+        return self._get_connection().add_privileges_to_group(
+            self._name, privilege_group, privileges
+        )
+
+    def remove_privileges_from_group(self, privilege_group: str, privileges: list):
+        """Remove privileges from a privilege group for the role
+            :param privilege_group: privilege group name.
+            :type  privilege_group: str
+            :param privileges: a list of privilege names.
+            :type  privileges: list
+
+        :example:
+            >>> from pymilvus import connections
+            >>> from pymilvus.orm.role import Role
+            >>> connections.connect()
+            >>> role = Role(role_name)
+            >>> role.remove_privileges_from_group(privilege_group, ["Insert", "Select"])
+        """
+        return self._get_connection().remove_privileges_from_group(
+            self._name, privilege_group, privileges
+        )