Skip to content

Commit

Permalink
fix: [2.4] add privilege group orm (#2357)
Browse files Browse the repository at this point in the history
issue: milvus-io/milvus#37031
issue: #2353

Signed-off-by: shaoting-huang <shaoting.huang@zilliz.com>
  • Loading branch information
shaoting-huang authored Nov 18, 2024
1 parent 0be8e7f commit 416a8b4
Show file tree
Hide file tree
Showing 4 changed files with 129 additions and 27 deletions.
22 changes: 22 additions & 0 deletions pymilvus/client/check.py
Original file line number Diff line number Diff line change
Expand Up @@ -271,6 +271,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 = {
Expand Down Expand Up @@ -315,6 +334,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):
Expand Down
20 changes: 12 additions & 8 deletions pymilvus/client/grpc_handler.py
Original file line number Diff line number Diff line change
Expand Up @@ -2015,14 +2015,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)

Expand All @@ -2035,20 +2037,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)
33 changes: 14 additions & 19 deletions pymilvus/client/prepare.py
Original file line number Diff line number Diff line change
Expand Up @@ -1462,33 +1462,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,
)
81 changes: 81 additions & 0 deletions pymilvus/orm/role.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
)

0 comments on commit 416a8b4

Please sign in to comment.