enhance: [2.4] RBAC Custom Privilege Group API (#2344)

issue: milvus-io/milvus#37031 --------- Signed-off-by: shaoting-huang <[email protected]>
milvus-io · Nov 13, 2024 · b490a0a · b490a0a
1 parent 7576627
commit b490a0a
Show file tree

Hide file tree

Showing 11 changed files with 606 additions and 177 deletions.
diff --git a/pymilvus/client/grpc_handler.py b/pymilvus/client/grpc_handler.py
@@ -2013,3 +2013,42 @@ def alloc_timestamp(self, timeout: Optional[float] = None) -> int:
         response = self._stub.AllocTimestamp(request, timeout=timeout)
         check_status(response.status)
         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)
+        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)
+        resp = self._stub.DropPrivilegeGroup(req, wait_for_ready=True, timeout=timeout)
+        check_status(resp)
+
+    @retry_on_rpc_failure()
+    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
+
+    @retry_on_rpc_failure()
+    def add_privileges_to_group(
+        self, group_name: str, privileges: List[str], timeout: Optional[float] = None, **kwargs
+    ):
+        req = Prepare.operate_privilege_group_req(
+            group_name, 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
+    ):
+        req = Prepare.operate_privilege_group_req(
+            group_name, 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
@@ -1460,3 +1460,35 @@ def alter_database_req(cls, db_name: str, properties: Dict):
     def describe_database_req(cls, db_name: str):
         check_pass_param(db_name=db_name)
         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)
+
+    @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)
+
+    @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)
+        return milvus_types.OperatePrivilegeGroupRequest(
+            group_name=group_name,
+            privileges=[milvus_types.PrivilegeEntity(name=p) for p in privileges],
+            type=operate_type,
+        )
diff --git a/pymilvus/grpc_gen/common_pb2.py b/pymilvus/grpc_gen/common_pb2.py
diff --git a/pymilvus/grpc_gen/common_pb2.pyi b/pymilvus/grpc_gen/common_pb2.pyi
@@ -66,6 +66,10 @@ class ErrorCode(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     TimeTickLongDelay: _ClassVar[ErrorCode]
     NotReadyServe: _ClassVar[ErrorCode]
     NotReadyCoordActivating: _ClassVar[ErrorCode]
+    CreatePrivilegeGroupFailure: _ClassVar[ErrorCode]
+    DropPrivilegeGroupFailure: _ClassVar[ErrorCode]
+    ListPrivilegeGroupsFailure: _ClassVar[ErrorCode]
+    OperatePrivilegeGroupFailure: _ClassVar[ErrorCode]
     DataCoordNA: _ClassVar[ErrorCode]
     DDRequestRace: _ClassVar[ErrorCode]
 
@@ -204,6 +208,10 @@ class MsgType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     SelectGrant: _ClassVar[MsgType]
     RefreshPolicyInfoCache: _ClassVar[MsgType]
     ListPolicy: _ClassVar[MsgType]
+    CreatePrivilegeGroup: _ClassVar[MsgType]
+    DropPrivilegeGroup: _ClassVar[MsgType]
+    ListPrivilegeGroups: _ClassVar[MsgType]
+    OperatePrivilegeGroup: _ClassVar[MsgType]
     CreateResourceGroup: _ClassVar[MsgType]
     DropResourceGroup: _ClassVar[MsgType]
     ListResourceGroups: _ClassVar[MsgType]
@@ -310,6 +318,10 @@ class ObjectPrivilege(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     PrivilegeGroupReadOnly: _ClassVar[ObjectPrivilege]
     PrivilegeGroupReadWrite: _ClassVar[ObjectPrivilege]
     PrivilegeGroupAdmin: _ClassVar[ObjectPrivilege]
+    PrivilegeCreatePrivilegeGroup: _ClassVar[ObjectPrivilege]
+    PrivilegeDropPrivilegeGroup: _ClassVar[ObjectPrivilege]
+    PrivilegeListPrivilegeGroups: _ClassVar[ObjectPrivilege]
+    PrivilegeOperatePrivilegeGroup: _ClassVar[ObjectPrivilege]
 
 class StateCode(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     __slots__ = ()
@@ -382,6 +394,10 @@ DiskQuotaExhausted: ErrorCode
 TimeTickLongDelay: ErrorCode
 NotReadyServe: ErrorCode
 NotReadyCoordActivating: ErrorCode
+CreatePrivilegeGroupFailure: ErrorCode
+DropPrivilegeGroupFailure: ErrorCode
+ListPrivilegeGroupsFailure: ErrorCode
+OperatePrivilegeGroupFailure: ErrorCode
 DataCoordNA: ErrorCode
 DDRequestRace: ErrorCode
 IndexStateNone: IndexState
@@ -505,6 +521,10 @@ OperatePrivilege: MsgType
 SelectGrant: MsgType
 RefreshPolicyInfoCache: MsgType
 ListPolicy: MsgType
+CreatePrivilegeGroup: MsgType
+DropPrivilegeGroup: MsgType
+ListPrivilegeGroups: MsgType
+OperatePrivilegeGroup: MsgType
 CreateResourceGroup: MsgType
 DropResourceGroup: MsgType
 ListResourceGroups: MsgType
@@ -593,6 +613,10 @@ PrivilegeRestoreRBAC: ObjectPrivilege
 PrivilegeGroupReadOnly: ObjectPrivilege
 PrivilegeGroupReadWrite: ObjectPrivilege
 PrivilegeGroupAdmin: ObjectPrivilege
+PrivilegeCreatePrivilegeGroup: ObjectPrivilege
+PrivilegeDropPrivilegeGroup: ObjectPrivilege
+PrivilegeListPrivilegeGroups: ObjectPrivilege
+PrivilegeOperatePrivilegeGroup: ObjectPrivilege
 Initializing: StateCode
 Healthy: StateCode
 Abnormal: StateCode

diff --git a/pymilvus/grpc_gen/milvus-proto b/pymilvus/grpc_gen/milvus-proto
diff --git a/pymilvus/grpc_gen/milvus_pb2.py b/pymilvus/grpc_gen/milvus_pb2.py
diff --git a/pymilvus/grpc_gen/milvus_pb2.pyi b/pymilvus/grpc_gen/milvus_pb2.pyi
@@ -17,6 +17,11 @@ class ShowType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     All: _ClassVar[ShowType]
     InMemory: _ClassVar[ShowType]
 
+class OperatePrivilegeGroupType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+    __slots__ = ()
+    AddPrivilegesToGroup: _ClassVar[OperatePrivilegeGroupType]
+    RemovePrivilegesFromGroup: _ClassVar[OperatePrivilegeGroupType]
+
 class OperateUserRoleType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     __slots__ = ()
     AddUserToRole: _ClassVar[OperateUserRoleType]
@@ -36,6 +41,8 @@ class QuotaState(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     DenyToWrite: _ClassVar[QuotaState]
 All: ShowType
 InMemory: ShowType
+AddPrivilegesToGroup: OperatePrivilegeGroupType
+RemovePrivilegesFromGroup: OperatePrivilegeGroupType
 AddUserToRole: OperateUserRoleType
 RemoveUserFromRole: OperateUserRoleType
 Grant: OperatePrivilegeType
@@ -1506,6 +1513,48 @@ class DropRoleRequest(_message.Message):
     force_drop: bool
     def __init__(self, base: _Optional[_Union[_common_pb2.MsgBase, _Mapping]] = ..., role_name: _Optional[str] = ..., force_drop: bool = ...) -> None: ...
 
+class CreatePrivilegeGroupRequest(_message.Message):
+    __slots__ = ("base", "group_name")
+    BASE_FIELD_NUMBER: _ClassVar[int]
+    GROUP_NAME_FIELD_NUMBER: _ClassVar[int]
+    base: _common_pb2.MsgBase
+    group_name: str
+    def __init__(self, base: _Optional[_Union[_common_pb2.MsgBase, _Mapping]] = ..., group_name: _Optional[str] = ...) -> None: ...
+
+class DropPrivilegeGroupRequest(_message.Message):
+    __slots__ = ("base", "group_name")
+    BASE_FIELD_NUMBER: _ClassVar[int]
+    GROUP_NAME_FIELD_NUMBER: _ClassVar[int]
+    base: _common_pb2.MsgBase
+    group_name: str
+    def __init__(self, base: _Optional[_Union[_common_pb2.MsgBase, _Mapping]] = ..., group_name: _Optional[str] = ...) -> None: ...
+
+class ListPrivilegeGroupsRequest(_message.Message):
+    __slots__ = ("base",)
+    BASE_FIELD_NUMBER: _ClassVar[int]
+    base: _common_pb2.MsgBase
+    def __init__(self, base: _Optional[_Union[_common_pb2.MsgBase, _Mapping]] = ...) -> None: ...
+
+class ListPrivilegeGroupsResponse(_message.Message):
+    __slots__ = ("status", "privilege_groups")
+    STATUS_FIELD_NUMBER: _ClassVar[int]
+    PRIVILEGE_GROUPS_FIELD_NUMBER: _ClassVar[int]
+    status: _common_pb2.Status
+    privilege_groups: _containers.RepeatedCompositeFieldContainer[PrivilegeGroupInfo]
+    def __init__(self, status: _Optional[_Union[_common_pb2.Status, _Mapping]] = ..., privilege_groups: _Optional[_Iterable[_Union[PrivilegeGroupInfo, _Mapping]]] = ...) -> None: ...
+
+class OperatePrivilegeGroupRequest(_message.Message):
+    __slots__ = ("base", "group_name", "privileges", "type")
+    BASE_FIELD_NUMBER: _ClassVar[int]
+    GROUP_NAME_FIELD_NUMBER: _ClassVar[int]
+    PRIVILEGES_FIELD_NUMBER: _ClassVar[int]
+    TYPE_FIELD_NUMBER: _ClassVar[int]
+    base: _common_pb2.MsgBase
+    group_name: str
+    privileges: _containers.RepeatedCompositeFieldContainer[PrivilegeEntity]
+    type: OperatePrivilegeGroupType
+    def __init__(self, base: _Optional[_Union[_common_pb2.MsgBase, _Mapping]] = ..., group_name: _Optional[str] = ..., privileges: _Optional[_Iterable[_Union[PrivilegeEntity, _Mapping]]] = ..., type: _Optional[_Union[OperatePrivilegeGroupType, str]] = ...) -> None: ...
+
 class OperateUserRoleRequest(_message.Message):
     __slots__ = ("base", "username", "role_name", "type")
     BASE_FIELD_NUMBER: _ClassVar[int]
@@ -1518,6 +1567,14 @@ class OperateUserRoleRequest(_message.Message):
     type: OperateUserRoleType
     def __init__(self, base: _Optional[_Union[_common_pb2.MsgBase, _Mapping]] = ..., username: _Optional[str] = ..., role_name: _Optional[str] = ..., type: _Optional[_Union[OperateUserRoleType, str]] = ...) -> None: ...
 
+class PrivilegeGroupInfo(_message.Message):
+    __slots__ = ("group_name", "privileges")
+    GROUP_NAME_FIELD_NUMBER: _ClassVar[int]
+    PRIVILEGES_FIELD_NUMBER: _ClassVar[int]
+    group_name: str
+    privileges: _containers.RepeatedCompositeFieldContainer[PrivilegeEntity]
+    def __init__(self, group_name: _Optional[str] = ..., privileges: _Optional[_Iterable[_Union[PrivilegeEntity, _Mapping]]] = ...) -> None: ...
+
 class SelectRoleRequest(_message.Message):
     __slots__ = ("base", "role", "include_user_info")
     BASE_FIELD_NUMBER: _ClassVar[int]
@@ -1647,14 +1704,16 @@ class UserInfo(_message.Message):
     def __init__(self, user: _Optional[str] = ..., password: _Optional[str] = ..., roles: _Optional[_Iterable[_Union[RoleEntity, _Mapping]]] = ...) -> None: ...
 
 class RBACMeta(_message.Message):
-    __slots__ = ("users", "roles", "grants")
+    __slots__ = ("users", "roles", "grants", "privilege_groups")
     USERS_FIELD_NUMBER: _ClassVar[int]
     ROLES_FIELD_NUMBER: _ClassVar[int]
     GRANTS_FIELD_NUMBER: _ClassVar[int]
+    PRIVILEGE_GROUPS_FIELD_NUMBER: _ClassVar[int]
     users: _containers.RepeatedCompositeFieldContainer[UserInfo]
     roles: _containers.RepeatedCompositeFieldContainer[RoleEntity]
     grants: _containers.RepeatedCompositeFieldContainer[GrantEntity]
-    def __init__(self, users: _Optional[_Iterable[_Union[UserInfo, _Mapping]]] = ..., roles: _Optional[_Iterable[_Union[RoleEntity, _Mapping]]] = ..., grants: _Optional[_Iterable[_Union[GrantEntity, _Mapping]]] = ...) -> None: ...
+    privilege_groups: _containers.RepeatedCompositeFieldContainer[PrivilegeGroupInfo]
+    def __init__(self, users: _Optional[_Iterable[_Union[UserInfo, _Mapping]]] = ..., roles: _Optional[_Iterable[_Union[RoleEntity, _Mapping]]] = ..., grants: _Optional[_Iterable[_Union[GrantEntity, _Mapping]]] = ..., privilege_groups: _Optional[_Iterable[_Union[PrivilegeGroupInfo, _Mapping]]] = ...) -> None: ...
 
 class BackupRBACMetaRequest(_message.Message):
     __slots__ = ("base",)
+379 −320		go-api/commonpb/common.pb.go
+3,621 −3,026		go-api/milvuspb/milvus.pb.go
+148 −0		go-api/milvuspb/milvus_grpc.pb.go
+134 −51		go-api/rgpb/rg.pb.go
+14 −0		proto/common.proto
+72 −0		proto/milvus.proto
+7 −0		proto/rg.proto