enhance: RBAC grant v2 api (#2361)

issue: milvus-io/milvus#37031

Signed-off-by: shaoting-huang <[email protected]>

pymilvus/client/grpc_handler.py

@@ -1872,6 +1872,46 @@ def revoke_privilege(
         resp = self._stub.OperatePrivilege(req, wait_for_ready=True, timeout=timeout)
         check_status(resp)
 
+    @retry_on_rpc_failure()
+    def grant_privilege_v2(
+        self,
+        role_name: str,
+        privilege: str,
+        db_name: str,
+        collection_name: str,
+        timeout: Optional[float] = None,
+        **kwargs,
+    ):
+        req = Prepare.operate_privilege_v2_request(
+            role_name,
+            privilege,
+            milvus_types.OperatePrivilegeType.Grant,
+            db_name,
+            collection_name,
+        )
+        resp = self._stub.OperatePrivilegeV2(req, wait_for_ready=True, timeout=timeout)
+        check_status(resp)
+
+    @retry_on_rpc_failure()
+    def revoke_privilege_v2(
+        self,
+        role_name: str,
+        privilege: str,
+        db_name: str,
+        collection_name: str,
+        timeout: Optional[float] = None,
+        **kwargs,
+    ):
+        req = Prepare.operate_privilege_v2_request(
+            role_name,
+            privilege,
+            milvus_types.OperatePrivilegeType.Revoke,
+            db_name,
+            collection_name,
+        )
+        resp = self._stub.OperatePrivilegeV2(req, wait_for_ready=True, timeout=timeout)
+        check_status(resp)
+
     @retry_on_rpc_failure()
     def select_grant_for_one_role(
         self, role_name: str, db_name: str, timeout: Optional[float] = None, **kwargs

pymilvus/client/prepare.py

@@ -1461,6 +1461,32 @@ def operate_privilege_request(
             type=operate_privilege_type,
         )
 
+    @classmethod
+    def operate_privilege_v2_request(
+        cls,
+        role_name: str,
+        privilege: str,
+        operate_privilege_type: Any,
+        db_name: str,
+        collection_name: str,
+    ):
+        check_pass_param(
+            role_name=role_name,
+            privilege=privilege,
+            operate_privilege_type=operate_privilege_type,
+            db_name=db_name,
+            collection_name=collection_name,
+        )
+        return milvus_types.OperatePrivilegeV2Request(
+            role=milvus_types.RoleEntity(name=role_name),
+            grantor=milvus_types.GrantorEntity(
+                privilege=milvus_types.PrivilegeEntity(name=privilege)
+            ),
+            type=operate_privilege_type,
+            db_name=db_name,
+            collection_name=collection_name,
+        )
+
     @classmethod
     def select_grant_request(cls, role_name: str, object: str, object_name: str, db_name: str):
         check_pass_param(role_name=role_name)

pymilvus/grpc_gen/common_pb2.pyi

@@ -322,6 +322,15 @@ class ObjectPrivilege(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     PrivilegeDropPrivilegeGroup: _ClassVar[ObjectPrivilege]
     PrivilegeListPrivilegeGroups: _ClassVar[ObjectPrivilege]
     PrivilegeOperatePrivilegeGroup: _ClassVar[ObjectPrivilege]
+    PrivilegeGroupClusterReadOnly: _ClassVar[ObjectPrivilege]
+    PrivilegeGroupClusterReadWrite: _ClassVar[ObjectPrivilege]
+    PrivilegeGroupClusterAdmin: _ClassVar[ObjectPrivilege]
+    PrivilegeGroupDatabaseReadOnly: _ClassVar[ObjectPrivilege]
+    PrivilegeGroupDatabaseReadWrite: _ClassVar[ObjectPrivilege]
+    PrivilegeGroupDatabaseAdmin: _ClassVar[ObjectPrivilege]
+    PrivilegeGroupCollectionReadOnly: _ClassVar[ObjectPrivilege]
+    PrivilegeGroupCollectionReadWrite: _ClassVar[ObjectPrivilege]
+    PrivilegeGroupCollectionAdmin: _ClassVar[ObjectPrivilege]
 
 class StateCode(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     __slots__ = ()
@@ -617,6 +626,15 @@ PrivilegeCreatePrivilegeGroup: ObjectPrivilege
 PrivilegeDropPrivilegeGroup: ObjectPrivilege
 PrivilegeListPrivilegeGroups: ObjectPrivilege
 PrivilegeOperatePrivilegeGroup: ObjectPrivilege
+PrivilegeGroupClusterReadOnly: ObjectPrivilege
+PrivilegeGroupClusterReadWrite: ObjectPrivilege
+PrivilegeGroupClusterAdmin: ObjectPrivilege
+PrivilegeGroupDatabaseReadOnly: ObjectPrivilege
+PrivilegeGroupDatabaseReadWrite: ObjectPrivilege
+PrivilegeGroupDatabaseAdmin: ObjectPrivilege
+PrivilegeGroupCollectionReadOnly: ObjectPrivilege
+PrivilegeGroupCollectionReadWrite: ObjectPrivilege
+PrivilegeGroupCollectionAdmin: ObjectPrivilege
 Initializing: StateCode
 Healthy: StateCode
 Abnormal: StateCode

pymilvus/grpc_gen/milvus_pb2.pyi

@@ -27,6 +27,12 @@ class OperateUserRoleType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     AddUserToRole: _ClassVar[OperateUserRoleType]
     RemoveUserFromRole: _ClassVar[OperateUserRoleType]
 
+class PrivilegeLevel(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
+    __slots__ = ()
+    Cluster: _ClassVar[PrivilegeLevel]
+    Database: _ClassVar[PrivilegeLevel]
+    Collection: _ClassVar[PrivilegeLevel]
+
 class OperatePrivilegeType(int, metaclass=_enum_type_wrapper.EnumTypeWrapper):
     __slots__ = ()
     Grant: _ClassVar[OperatePrivilegeType]
@@ -45,6 +51,9 @@ AddPrivilegesToGroup: OperatePrivilegeGroupType
 RemovePrivilegesFromGroup: OperatePrivilegeGroupType
 AddUserToRole: OperateUserRoleType
 RemoveUserFromRole: OperateUserRoleType
+Cluster: PrivilegeLevel
+Database: PrivilegeLevel
+Collection: PrivilegeLevel
 Grant: OperatePrivilegeType
 Revoke: OperatePrivilegeType
 Unknown: QuotaState
@@ -1710,14 +1719,32 @@ class SelectGrantResponse(_message.Message):
     def __init__(self, status: _Optional[_Union[_common_pb2.Status, _Mapping]] = ..., entities: _Optional[_Iterable[_Union[GrantEntity, _Mapping]]] = ...) -> None: ...
 
 class OperatePrivilegeRequest(_message.Message):
-    __slots__ = ("base", "entity", "type")
+    __slots__ = ("base", "entity", "type", "version")
     BASE_FIELD_NUMBER: _ClassVar[int]
     ENTITY_FIELD_NUMBER: _ClassVar[int]
     TYPE_FIELD_NUMBER: _ClassVar[int]
+    VERSION_FIELD_NUMBER: _ClassVar[int]
     base: _common_pb2.MsgBase
     entity: GrantEntity
     type: OperatePrivilegeType
-    def __init__(self, base: _Optional[_Union[_common_pb2.MsgBase, _Mapping]] = ..., entity: _Optional[_Union[GrantEntity, _Mapping]] = ..., type: _Optional[_Union[OperatePrivilegeType, str]] = ...) -> None: ...
+    version: str
+    def __init__(self, base: _Optional[_Union[_common_pb2.MsgBase, _Mapping]] = ..., entity: _Optional[_Union[GrantEntity, _Mapping]] = ..., type: _Optional[_Union[OperatePrivilegeType, str]] = ..., version: _Optional[str] = ...) -> None: ...
+
+class OperatePrivilegeV2Request(_message.Message):
+    __slots__ = ("base", "role", "grantor", "type", "db_name", "collection_name")
+    BASE_FIELD_NUMBER: _ClassVar[int]
+    ROLE_FIELD_NUMBER: _ClassVar[int]
+    GRANTOR_FIELD_NUMBER: _ClassVar[int]
+    TYPE_FIELD_NUMBER: _ClassVar[int]
+    DB_NAME_FIELD_NUMBER: _ClassVar[int]
+    COLLECTION_NAME_FIELD_NUMBER: _ClassVar[int]
+    base: _common_pb2.MsgBase
+    role: RoleEntity
+    grantor: GrantorEntity
+    type: OperatePrivilegeType
+    db_name: str
+    collection_name: str
+    def __init__(self, base: _Optional[_Union[_common_pb2.MsgBase, _Mapping]] = ..., role: _Optional[_Union[RoleEntity, _Mapping]] = ..., grantor: _Optional[_Union[GrantorEntity, _Mapping]] = ..., type: _Optional[_Union[OperatePrivilegeType, str]] = ..., db_name: _Optional[str] = ..., collection_name: _Optional[str] = ...) -> None: ...
 
 class UserInfo(_message.Message):
     __slots__ = ("user", "password", "roles")

pymilvus/grpc_gen/milvus_pb2_grpc.py

@@ -346,6 +346,11 @@ def __init__(self, channel):
                 request_serializer=milvus__pb2.OperatePrivilegeRequest.SerializeToString,
                 response_deserializer=common__pb2.Status.FromString,
                 )
+        self.OperatePrivilegeV2 = channel.unary_unary(
+                '/milvus.proto.milvus.MilvusService/OperatePrivilegeV2',
+                request_serializer=milvus__pb2.OperatePrivilegeV2Request.SerializeToString,
+                response_deserializer=common__pb2.Status.FromString,
+                )
         self.SelectGrant = channel.unary_unary(
                 '/milvus.proto.milvus.MilvusService/SelectGrant',
                 request_serializer=milvus__pb2.SelectGrantRequest.SerializeToString,
@@ -889,6 +894,12 @@ def OperatePrivilege(self, request, context):
         context.set_details('Method not implemented!')
         raise NotImplementedError('Method not implemented!')
 
+    def OperatePrivilegeV2(self, request, context):
+        """Missing associated documentation comment in .proto file."""
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
     def SelectGrant(self, request, context):
         """Missing associated documentation comment in .proto file."""
         context.set_code(grpc.StatusCode.UNIMPLEMENTED)
@@ -1384,6 +1395,11 @@ def add_MilvusServiceServicer_to_server(servicer, server):
                     request_deserializer=milvus__pb2.OperatePrivilegeRequest.FromString,
                     response_serializer=common__pb2.Status.SerializeToString,
             ),
+            'OperatePrivilegeV2': grpc.unary_unary_rpc_method_handler(
+                    servicer.OperatePrivilegeV2,
+                    request_deserializer=milvus__pb2.OperatePrivilegeV2Request.FromString,
+                    response_serializer=common__pb2.Status.SerializeToString,
+            ),
             'SelectGrant': grpc.unary_unary_rpc_method_handler(
                     servicer.SelectGrant,
                     request_deserializer=milvus__pb2.SelectGrantRequest.FromString,
@@ -2651,6 +2667,23 @@ def OperatePrivilege(request,
             options, channel_credentials,
             insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
 
+    @staticmethod
+    def OperatePrivilegeV2(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(request, target, '/milvus.proto.milvus.MilvusService/OperatePrivilegeV2',
+            milvus__pb2.OperatePrivilegeV2Request.SerializeToString,
+            common__pb2.Status.FromString,
+            options, channel_credentials,
+            insecure, call_credentials, compression, wait_for_ready, timeout, metadata)
+
     @staticmethod
     def SelectGrant(request,
             target,

pymilvus/milvus_client/milvus_client.py

@@ -1003,6 +1003,34 @@ def revoke_privilege(
             role_name, object_type, object_name, privilege, db_name, timeout=timeout, **kwargs
         )
 
+    def grant_privilege_v2(
+        self,
+        role_name: str,
+        privilege: str,
+        db_name: str,
+        collection_name: str,
+        timeout: Optional[float] = None,
+        **kwargs,
+    ):
+        conn = self._get_connection()
+        conn.grant_privilege_v2(
+            role_name, privilege, db_name, collection_name, timeout=timeout, **kwargs
+        )
+
+    def revoke_privilege_v2(
+        self,
+        role_name: str,
+        privilege: str,
+        db_name: str,
+        collection_name: str,
+        timeout: Optional[float] = None,
+        **kwargs,
+    ):
+        conn = self._get_connection()
+        conn.revoke_privilege_v2(
+            role_name, privilege, db_name, collection_name, timeout=timeout, **kwargs
+        )
+
     def create_alias(
         self, collection_name: str, alias: str, timeout: Optional[float] = None, **kwargs
     ):

pymilvus/orm/role.py

@@ -176,6 +176,46 @@ 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):
+        """Grant a privilege for the role
+            :param privilege: privilege name.
+            :type  privilege: str
+            :param db_name: db name.
+            :type  db_name: str
+            :param collection_name: collection name.
+            :type  collection_name: str
+
+        :example:
+            >>> from pymilvus import connections
+            >>> from pymilvus.orm.role import Role
+            >>> connections.connect()
+            >>> role = Role(role_name)
+            >>> role.grant_v2("Insert", db_name, collection_name)
+        """
+        return self._get_connection().grant_privilege_v2(
+            self._name, privilege, db_name, collection_name
+        )
+
+    def revoke_v2(self, privilege: str, db_name: str, collection_name: str):
+        """Revoke a privilege for the role
+            :param privilege: privilege name.
+            :type  privilege: str
+            :param db_name: db name.
+            :type  db_name: str
+            :param collection_name: collection name.
+            :type  collection_name: str
+
+        :example:
+            >>> from pymilvus import connections
+            >>> from pymilvus.orm.role import Role
+            >>> connections.connect()
+            >>> role = Role(role_name)
+            >>> role.revoke_v2("Insert", db_name, collection_name)
+        """
+        return self._get_connection().revoke_privilege_v2(
+            self._name, privilege, db_name, collection_name
+        )
+
     def list_grant(self, object: str, object_name: str, db_name: str = ""):
         """List a grant info for the role and the specific object
             :param object: object type.