Merge pull request #91 from zilliztech/issue-90

feat: support role apis and update pymilvus

Author: shanghaikid

README.md

@@ -30,8 +30,8 @@ Python >= 3.8.5
 
 #### Install from PyPI (Recommended)
 
-Run `pip install pymilvus==2.4.3`
-Run `pip install milvus-cli==0.4.3`
+Run `pip install pymilvus>=2.4.3`
+Run `pip install milvus-cli==1.0.1`
 
 #### Install from a tarball
 

milvus_cli/Cli.py

@@ -6,6 +6,7 @@
 from User import MilvusUser
 from Alias import MilvusAlias
 from Partition import MilvusPartition
+from Role import MilvusRole
 from pymilvus import __version__
 
 
@@ -16,5 +17,6 @@ class MilvusCli(object):
     index = MilvusIndex()
     data = MilvusData()
     user = MilvusUser()
+    role = MilvusRole()
     alias = MilvusAlias()
     partition = MilvusPartition()

milvus_cli/Collection.py

@@ -56,13 +56,13 @@ def create_collection(
                     FieldSchema(
                         name=fieldName,
                         dtype=DataType[upperFieldType],
-                        max_length=restData[0],
+                        max_length=int(restData[0]),
                     )
                 )
             elif upperFieldType == "ARRAY":
                 upperElementType = restData[1].upper()
-                max_capacity = restData[0]
-                maxLength = restData[2] if len(restData) == 3 else None
+                max_capacity = int(restData[0])
+                maxLength = int(restData[2]) if len(restData) == 3 else None
                 fieldList.append(
                     FieldSchema(
                         name=fieldName,

milvus_cli/Partition.py

@@ -1,5 +1,4 @@
 from Collection import getTargetCollection
-from tabulate import tabulate
 
 
 class MilvusPartition(object):

milvus_cli/Role.py

@@ -0,0 +1,116 @@
+from pymilvus import connections, Role, utility
+from tabulate import tabulate
+
+
+class MilvusRole(object):
+    def createRole(self, roleName):
+        try:
+            role = Role(name=roleName)
+            role.create()
+        except Exception as e:
+            raise Exception(f"Create role error!{str(e)}")
+        else:
+            return f"Create role successfully!"
+
+    def listRoles(self):
+        try:
+            res = utility.list_roles(include_user_info=True)
+            data = [[role.role_name, role.users] for role in res.groups]
+            print(tabulate(data, headers=["roleName", "users"], tablefmt="pretty"))
+        except Exception as e:
+            raise Exception(f"List role error!{str(e)}")
+        else:
+            return res.groups
+
+    def dropRole(self, roleName):
+        try:
+            role = Role(name=roleName)
+            role.drop()
+        except Exception as e:
+            raise Exception(f"Drop role error!{str(e)}")
+        else:
+            return f"Drop role successfully!"
+
+    def grantRole(self, roleName, username):
+        try:
+            role = Role(name=roleName)
+            role.add_user(username=username)
+            users = role.get_users()
+            print(tabulate([users], headers=["users"], tablefmt="pretty"))
+        except Exception as e:
+            raise Exception(f"Grant role error!{str(e)}")
+        else:
+            return users
+
+    def revokeRole(self, roleName, username):
+        try:
+            role = Role(name=roleName)
+            role.remove_user(username=username)
+            users = role.get_users()
+            print(tabulate([users], headers=["users"], tablefmt="pretty"))
+        except Exception as e:
+            raise Exception(f"Revoke role error!{str(e)}")
+        else:
+            return users
+
+    def grantPrivilege(
+        self, roleName, objectName, objectType, privilege, dbName="default"
+    ):
+        try:
+            role = Role(name=roleName)
+            role.grant(
+                object=objectType,
+                privilege=privilege,
+                object_name=objectName,
+                db_name=dbName,
+            )
+        except Exception as e:
+            raise Exception(f"Grant privilege error!{str(e)}")
+        else:
+            return f"Grant privilege successfully!"
+
+    def revokePrivilege(
+        self, roleName, objectName, objectType, privilege, dbName="default"
+    ):
+        try:
+            role = Role(name=roleName)
+            role.revoke(
+                object=objectType,
+                privilege=privilege,
+                object_name=objectName,
+                db_name=dbName,
+            )
+        except Exception as e:
+            raise Exception(f"Revoke privilege error!{str(e)}")
+        else:
+            return f"Revoke privilege successfully!"
+
+    def listGrants(self, roleName, objectName, objectType):
+        try:
+            role = Role(name=roleName)
+            grants = role.list_grant(object=objectType, object_name=objectName)
+            headers = [
+                "Object",
+                "Object Name",
+                "DB Name",
+                "Role Name",
+                "Grantor Name",
+                "Privilege",
+            ]
+            data = [
+                [
+                    grant.object,
+                    grant.object_name,
+                    grant.db_name,
+                    grant.role_name,
+                    grant.grantor_name,
+                    grant.privilege,
+                ]
+                for grant in grants.groups
+            ]
+
+            print(tabulate(data, headers=headers, tablefmt="pretty"))
+        except Exception as e:
+            raise Exception(f"List grants error!{str(e)}")
+        else:
+            return data

milvus_cli/Types.py

@@ -171,3 +171,55 @@ def __str__(self):
 }
 
 Operators = ["<", "<=", ">", ">=", "==", "!=", "in"]
+
+Privileges = [
+    "CreateIndex",
+    "DropIndex",
+    "IndexDetail",
+    "Load",
+    "GetLoadingProgress",
+    "GetLoadState",
+    "Release",
+    "Insert",
+    "Delete",
+    "Upsert",
+    "Search",
+    "Flush",
+    "GetFlushState",
+    "Query",
+    "GetStatistics",
+    "Compaction",
+    "Import",
+    "LoadBalance",
+    "CreatePartition",
+    "DropPartition",
+    "ShowPartitions",
+    "HasPartition",
+    "All",
+    "CreateCollection",
+    "DropCollection",
+    "DescribeCollection",
+    "ShowCollections",
+    "RenameCollection",
+    "FlushAll",
+    "CreateOwnership",
+    "DropOwnership",
+    "SelectOwnership",
+    "ManageOwnership",
+    "CreateResourceGroup",
+    "DropResourceGroup",
+    "DescribeResourceGroup",
+    "ListResourceGroups",
+    "TransferNode",
+    "TransferReplica",
+    "CreateDatabase",
+    "DropDatabase",
+    "ListDatabases",
+    "CreateAlias",
+    "DropAlias",
+    "DescribeAlias",
+    "ListAliases",
+    "UpdateUser",
+    "SelectUser",
+    "*",
+]

milvus_cli/scripts/helper_cli.py

@@ -44,7 +44,7 @@ def show(obj):
 @cli.group("list", no_args_is_help=False)
 @click.pass_obj
 def getList(obj):
-    """List collections,databases, partitions,users or indexes."""
+    """List collections,databases, partitions, users, grants or indexes."""
     pass
 
 
@@ -58,7 +58,21 @@ def rename(obj):
 @cli.group("create", no_args_is_help=False)
 @click.pass_obj
 def create(obj):
-    """Create collection, database, partition,user or index."""
+    """Create collection, database, partition,user,role or index."""
+    pass
+
+
+@cli.group("grant", no_args_is_help=False)
+@click.pass_obj
+def grant(obj):
+    """Grant role,  privilege."""
+    pass
+
+
+@cli.group("revoke", no_args_is_help=False)
+@click.pass_obj
+def revoke(obj):
+    """Revoke role,  privilege."""
     pass
 
 
@@ -79,7 +93,7 @@ def release(obj):
 @cli.group("delete", no_args_is_help=False)
 @click.pass_obj
 def delete(obj):
-    """Delete collection, database, partition,alias,user or index."""
+    """Delete collection, database, partition,alias,user, role or index."""
     pass
 
 

milvus_cli/scripts/milvus_cli.py

@@ -7,7 +7,7 @@
 from .user_cli import *
 from .alias_cli import *
 from .partition_cli import *
-
+from .role_cli import *
 
 if __name__ == "__main__":
     runCliPrompt()