enhance: rbac privilege group and grant v2 examples (#2425)

issue: milvus-io/milvus#37031

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

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()

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")    
+

pymilvus/milvus_client/milvus_client.py

@@ -1078,6 +1078,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,
@@ -1097,6 +1112,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,