From 47898c48c1c3c46b6600d4a5ea530519a95e848d Mon Sep 17 00:00:00 2001 From: sthuang <167743503+shaoting-huang@users.noreply.github.com> Date: Wed, 11 Dec 2024 15:14:36 +0800 Subject: [PATCH] enhance: rbac privilege group and grant v2 examples (#2425) issue: https://github.com/milvus-io/milvus/issues/37031 Signed-off-by: shaoting-huang --- examples/orm/rbac.py | 114 ++++++++++++++++++++++++ examples/rbac.py | 52 ++++++++++- pymilvus/milvus_client/milvus_client.py | 30 +++++++ 3 files changed, 195 insertions(+), 1 deletion(-) create mode 100644 examples/orm/rbac.py diff --git a/examples/orm/rbac.py b/examples/orm/rbac.py new file mode 100644 index 000000000..6ecfb5c18 --- /dev/null +++ b/examples/orm/rbac.py @@ -0,0 +1,114 @@ +from pymilvus import ( + connections, + FieldSchema, CollectionSchema, DataType, + Collection, + utility, + Role +) + +# This example shows how to: +# 1. connect to Milvus server +# 2. create a role +# 3. add user to the role +# 4. assign the role to the user +# 5. create privilege group +# 6. add privileges to the privilege group +# 7. remove privileges from the privilege group +# 8. list privilege groups of the role +# 9. grant the privilege group to the role +# 10. grant a built-in privilege group to the role +# 11. revoke the privilege group from the role +# 12. revoke a built-in privilege group from the role +# 13. drop the role + +_HOST = 'localhost' +_USER = 'root' +_PASSWORD = 'Milvus' + +# Create a Milvus connection +def create_connection(): + print(f"\nCreate connection...") + connections.connect(host=_HOST, user=_USER, password=_PASSWORD) + print(f"\nList connections:") + print(connections.list_connections()) + +def create_privilege_group(role, group_name): + role.create_privilege_group(group_name) + +def add_privileges_to_group(role, group_name, privileges): + role.add_privileges_to_group(group_name, privileges) + +def list_privilege_groups(role): + return role.list_privilege_groups() + +def remove_privileges_from_group(role, group_name, privileges): + role.remove_privileges_from_group(group_name, privileges) + +def drop_privilege_group(role, group_name): + role.drop_privilege_group(group_name) + +def grant_v2(role, privilege, collection_name, db_name): + role.grant_v2(privilege, collection_name, db_name=db_name) + +def revoke_v2(role, privilege, collection_name, db_name): + role.revoke_v2(privilege, collection_name, db_name=db_name) + +def list_grants(role): + return role.list_grants() + +def main(): + # Connect to Milvus server + create_connection() + + # create role + role = Role("privilege_group_role") + role.create() + + # list roles + utility.list_roles(True) + + # create user + utility.create_user(user="user1", password="Milvus") + + # add user + role.add_user("user1") + + # create privilege group + privilege_group = "search_query" + create_privilege_group(role, privilege_group) + + # add privileges to group + add_privileges_to_group(role, privilege_group, ["Search", "Query"]) + + # list privilege groups + groups = list_privilege_groups(role) + print("List privilege groups: ", groups) + + # remove privileges from group + remove_privileges_from_group(role, privilege_group, ["Query"]) + + # grant custom privielge group to role + grant_v2(role, privilege_group, "*", db_name="*") + + # grant built-in privielge group to role + grant_v2(role, "ClusterReadOnly", "*", db_name="*") + + # list grants + grants =list_grants(role) + print("List grants: ", grants) + + # revoke custom privielge group from role + revoke_v2(role, privilege_group, "*", db_name="*") + + # revoke built-in privielge group from role + revoke_v2(role, "ClusterReadOnly", "*", db_name="*") + + # drop privilege group + drop_privilege_group(role, privilege_group) + + # drop role + role.drop() + + +if __name__ == '__main__': + main() diff --git a/examples/rbac.py b/examples/rbac.py index 63ca5d005..cdc69ebdf 100644 --- a/examples/rbac.py +++ b/examples/rbac.py @@ -53,6 +53,10 @@ role_db_rw = "db_rw" role_db_ro = "db_ro" +role_custom = "custom_role" +role_cluster_admin = "cluster_admin" +role_database_readonly = "database_readonly" +role_collection_readwrite = "collection_readwrite" current_roles = milvus_client.list_roles() print("current roles:", current_roles) @@ -64,7 +68,31 @@ milvus_client.revoke_privilege(role, item["object_type"], item["privilege"], item["object_name"]) milvus_client.drop_role(role) - + +# manage custom privilege group and grant it to custom role +privilege_group_name = "custom_privilege_group" +milvus_client.create_privilege_group(privilege_group_name) +milvus_client.add_privileges_to_group(privilege_group_name, ["Search", "Query"]) +milvus_client.list_privilege_groups() +milvus_client.remove_privileges_from_group(privilege_group_name, ["Search"]) +milvus_client.list_privilege_groups() +milvus_client.create_role(role_custom) +milvus_client.grant_privilege_v2(role_custom, privilege_group_name, "*") + +# grant cluster level built-in privilege group +milvus_client.create_role(role_cluster_admin) +milvus_client.grant_privilege_v2(role_cluster_admin, "ClusterAdmin", "*", "*") + +# grant database level built-in privilege group +milvus_client.create_role(role_database_readonly) +milvus_client.grant_privilege_v2(role_database_readonly, "DatabaseReadOnly", "*", "db1") + +# grant collection level built-in privilege group +milvus_client.create_role(role_collection_readwrite) +milvus_client.grant_privilege_v2(role_collection_readwrite, "CollectionReadWrite", "col1", "db1") + +roles = milvus_client.list_roles() +print("roles:", roles) milvus_client.create_role(role_db_rw) for item in db_rw_privileges: @@ -88,9 +116,12 @@ print(f"grant {role_db_rw} to user1") milvus_client.grant_role("user1", role_db_rw) print("user info for user1:", user1_info) +milvus_client.grant_role("user1", role_collection_readwrite) milvus_client.grant_role("user2", role_db_ro) milvus_client.grant_role("user2", role_db_rw) +milvus_client.grant_role("user2", role_database_readonly) +milvus_client.grant_role("user2", role_cluster_admin) user2_info = milvus_client.describe_user("user2") print("user info for user2:", user2_info) @@ -102,3 +133,22 @@ user3_info = milvus_client.describe_user("user3") print("user info for user3:", user3_info) + +# revoke all privileges before dropping roles and users +milvus_client.revoke_privilege_v2(role_cluster_admin, "ClusterAdmin", "*", "*") +milvus_client.revoke_privilege_v2(role_database_readonly, "DatabaseReadOnly", "*", "db1") +milvus_client.revoke_privilege_v2(role_collection_readwrite, "CollectionReadWrite", "col1", "db1") +milvus_client.revoke_privilege_v2(role_custom, privilege_group_name, "*") + +milvus_client.drop_role(role_cluster_admin) +milvus_client.drop_role(role_database_readonly) +milvus_client.drop_role(role_collection_readwrite) +milvus_client.drop_role(role_custom) +milvus_client.list_roles + +milvus_client.drop_privilege_group(privilege_group_name) + +milvus_client.drop_user("user1") +milvus_client.drop_user("user2") +milvus_client.drop_user("user3") + diff --git a/pymilvus/milvus_client/milvus_client.py b/pymilvus/milvus_client/milvus_client.py index 27b514d9b..89be1e2c2 100644 --- a/pymilvus/milvus_client/milvus_client.py +++ b/pymilvus/milvus_client/milvus_client.py @@ -1013,6 +1013,21 @@ def grant_privilege_v2( timeout: Optional[float] = None, **kwargs, ): + """Grant a privilege or a privilege group to a role. + + Args: + role_name (``str``): The name of the role. + privilege (``str``): The privilege or privilege group to grant. + collection_name (``str``): The name of the collection. + db_name (``str``, optional): The name of the database. It will use default database + if not specified. + 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. + + Raises: + MilvusException: If anything goes wrong. + """ conn = self._get_connection() conn.grant_privilege_v2( role_name, @@ -1032,6 +1047,21 @@ def revoke_privilege_v2( timeout: Optional[float] = None, **kwargs, ): + """Revoke a privilege or a privilege group from a role. + + Args: + role_name (``str``): The name of the role. + privilege (``str``): The privilege or privilege group to revoke. + collection_name (``str``): The name of the collection. + db_name (``str``, optional): The name of the database. It will use default database + if not specified. + 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. + + Raises: + MilvusException: If anything goes wrong. + """ conn = self._get_connection() conn.revoke_privilege_v2( role_name,