From fe021c147428cbe0899510995d119cf000558bc0 Mon Sep 17 00:00:00 2001 From: nameczz Date: Mon, 16 Dec 2024 17:06:10 +0800 Subject: [PATCH] feat: support role apis and update pymilvus Signed-off-by: nameczz --- README.md | 4 +- milvus_cli/Cli.py | 2 + milvus_cli/Collection.py | 6 +- milvus_cli/Partition.py | 1 - milvus_cli/Role.py | 116 ++++++++++++++++++++ milvus_cli/Types.py | 52 +++++++++ milvus_cli/scripts/helper_cli.py | 20 +++- milvus_cli/scripts/milvus_cli.py | 2 +- milvus_cli/scripts/role_cli.py | 183 +++++++++++++++++++++++++++++++ milvus_cli/test/test_index.py | 16 +-- milvus_cli/test/test_role.py | 83 ++++++++++++++ milvus_cli/utils.py | 24 ++-- setup.py | 4 +- 13 files changed, 482 insertions(+), 31 deletions(-) create mode 100644 milvus_cli/Role.py create mode 100644 milvus_cli/scripts/role_cli.py create mode 100644 milvus_cli/test/test_role.py diff --git a/README.md b/README.md index 1d4886b..bd5257a 100644 --- a/README.md +++ b/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 diff --git a/milvus_cli/Cli.py b/milvus_cli/Cli.py index 2c5986a..fe1d19f 100644 --- a/milvus_cli/Cli.py +++ b/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() diff --git a/milvus_cli/Collection.py b/milvus_cli/Collection.py index 69af5b5..001ec96 100644 --- a/milvus_cli/Collection.py +++ b/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, diff --git a/milvus_cli/Partition.py b/milvus_cli/Partition.py index d9fd5b3..1129836 100644 --- a/milvus_cli/Partition.py +++ b/milvus_cli/Partition.py @@ -1,5 +1,4 @@ from Collection import getTargetCollection -from tabulate import tabulate class MilvusPartition(object): diff --git a/milvus_cli/Role.py b/milvus_cli/Role.py new file mode 100644 index 0000000..44abeec --- /dev/null +++ b/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 diff --git a/milvus_cli/Types.py b/milvus_cli/Types.py index 19198b8..ee79a13 100644 --- a/milvus_cli/Types.py +++ b/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", + "*", +] diff --git a/milvus_cli/scripts/helper_cli.py b/milvus_cli/scripts/helper_cli.py index f6f7f46..38b2ba6 100644 --- a/milvus_cli/scripts/helper_cli.py +++ b/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 diff --git a/milvus_cli/scripts/milvus_cli.py b/milvus_cli/scripts/milvus_cli.py index 9e995ad..ae8bbd7 100644 --- a/milvus_cli/scripts/milvus_cli.py +++ b/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() diff --git a/milvus_cli/scripts/role_cli.py b/milvus_cli/scripts/role_cli.py new file mode 100644 index 0000000..feb0d5d --- /dev/null +++ b/milvus_cli/scripts/role_cli.py @@ -0,0 +1,183 @@ +from .helper_cli import create, getList, delete, grant, revoke +import click +from ..Types import Privileges + + +@create.command("role") +@click.option("-r", "--roleName", "roleName", help="The role name of milvus role.") +@click.pass_obj +def create_role(obj, roleName): + """ + Create a new role. + + Example: + + milvus_cli > create role -r role1 + """ + try: + click.echo(obj.role.createRole(roleName)) + obj.role.listRoles() + except Exception as e: + click.echo(message=e, err=True) + + +@getList.command("roles") +@click.pass_obj +def list_roles(obj): + """List all roles in Milvus + Example: + + milvus_cli >list roles + """ + try: + obj.role.listRoles() + except Exception as e: + click.echo(message=e, err=True) + + +@delete.command("role") +@click.option("-r", "--roleName", "roleName", help="The role name of milvus role.") +@click.pass_obj +def drop_role(obj, roleName): + """ + Drop role. + + Example: + + milvus_cli > drop role -r role1 + """ + try: + obj.role.dropRole(roleName) + obj.role.listRoles() + except Exception as e: + click.echo(message=e, err=True) + + +@grant.command("role") +@click.option("-r", "--roleName", "roleName", help="The role name of milvus role.") +@click.option("-u", "--username", "username", help="The username of milvus user.") +@click.pass_obj +def add_role_users(obj, roleName, username): + """ + Add user to role. + + Example: + + milvus_cli > grant role -r role1 -u user1 + """ + try: + obj.role.grantRole(roleName, username) + except Exception as e: + click.echo(message=e, err=True) + + +@revoke.command("role") +@click.option("-r", "--roleName", "roleName", help="The role name of milvus role.") +@click.option("-u", "--username", "username", help="The username of milvus user.") +@click.pass_obj +def drop_role_users(obj, roleName, username): + """ + Drop user from role. + + Example: + + milvus_cli > revoke role -r role1 -u user1 + """ + try: + obj.role.revokeRole(roleName, username) + except Exception as e: + click.echo(message=e, err=True) + + +@grant.command("privilege") +@click.pass_obj +def grant_privilege(obj): + """ + Grant privilege to role. + + Example: + + milvus_cli > grant privilege + """ + try: + roleName = click.prompt("Role name") + print(roleName) + objectType = click.prompt( + "The type of object for which the privilege is to be assigned.", + type=click.Choice(["Global", "Collection", "User"]), + ) + objectName = click.prompt( + "The name of the object to control access for", type=str + ) + privilege = click.prompt( + "The name of the privilege to assign.", + type=click.Choice( + Privileges, + case_sensitive=True, + ), + ) + dbName = click.prompt( + "The name of the database to which the object belongs.", default="default" + ) + obj.role.grantPrivilege(roleName, objectName, objectType, privilege, dbName) + obj.role.listGrants(roleName, objectName, objectType) + except Exception as e: + click.echo(message=e, err=True) + + +@revoke.command("privilege") +@click.pass_obj +def revoke_privilege(obj): + """ + Revoke privilege from role. + + Example: + + milvus_cli > revoke privilege + """ + try: + roleName = click.prompt("Role name") + objectType = click.prompt( + "The type of object for which the privilege is to be assigned.", + type=click.Choice(["Global", "Collection", "User"]), + ) + objectName = click.prompt( + "The name of the object to control access for", type=str + ) + privilege = click.prompt( + "The name of the privilege to assign.", + type=click.Choice( + Privileges, + case_sensitive=True, + ), + ) + dbName = click.prompt( + "The name of the database to which the object belongs.", default="default" + ) + obj.role.revokePrivilege(roleName, objectName, objectType, privilege, dbName) + obj.role.listGrants(roleName, objectName, objectType) + except Exception as e: + click.echo(message=e, err=True) + + +@getList.command("grants") +@click.option("-r", "--roleName", "roleName", help="The role name of milvus role.") +@click.option( + "-o", "--objectName", "objectName", help="The object name of milvus object." +) +@click.option( + "-t", "--objectType", "objectType", help="The object type of milvus object." +) +@click.pass_obj +def list_grants(obj, roleName, objectName, objectType): + """ + List all grants in Milvus + + Example: + + milvus_cli > list grants -r role1 -o object1 -t Collection + """ + try: + obj.role.listGrants(roleName, objectName, objectType) + except Exception as e: + click.echo(message=e, err=True) diff --git a/milvus_cli/test/test_index.py b/milvus_cli/test/test_index.py index f1b8a96..1459b4d 100644 --- a/milvus_cli/test/test_index.py +++ b/milvus_cli/test/test_index.py @@ -10,7 +10,6 @@ from Index import MilvusIndex uri = "http://localhost:19530" -tempAlias = "zilliz2" collectionName = "test_collection" vectorName = "title_vector" indexName = "vec_index" @@ -42,8 +41,8 @@ def setUpClass(cls): @classmethod def tearDownClass(cls): - collection.drop_collection(tempAlias, collectionName) - milvusConnection.disconnect(alias=tempAlias) + collection.drop_collection(collectionName) + milvusConnection.disconnect() def test_create_index(self): params = [f"nlist:128"] @@ -54,23 +53,20 @@ def test_create_index(self): metricType="IP", indexType="IVF_SQ8", params=params, - alias=tempAlias, ) self.assertEqual(res.code, 0) - print(milvusIndex.list_indexes(collectionName, tempAlias)) + print(milvusIndex.list_indexes(collectionName)) def test_describe_index(self): - indexDetail = milvusIndex.get_index_details( - collectionName, indexName, tempAlias - ) + indexDetail = milvusIndex.get_index_details(collectionName, indexName) self.assertIn(indexName, indexDetail) def test_has_index(self): - res = milvusIndex.has_index(collectionName, indexName, tempAlias) + res = milvusIndex.has_index(collectionName, indexName) self.assertTrue(res) def test_drop_index(self): - res = milvusIndex.drop_index(collectionName, vectorName, tempAlias) + res = milvusIndex.drop_index(collectionName, vectorName) self.assertIsInstance(res, str) diff --git a/milvus_cli/test/test_role.py b/milvus_cli/test/test_role.py new file mode 100644 index 0000000..37bce83 --- /dev/null +++ b/milvus_cli/test/test_role.py @@ -0,0 +1,83 @@ +import unittest +import sys +import os + +current_dir = os.path.dirname(os.path.realpath(__file__)) +parent_dir = os.path.dirname(current_dir) +sys.path.append(parent_dir) +from Connection import MilvusConnection +from Role import MilvusRole +from User import MilvusUser + +uri = "http://localhost:19530" +roleName = "test_role_zilliz" +username = "test_zilliz" +objectName = "a" + + +milvusConnection = MilvusConnection() +role = MilvusRole() +user = MilvusUser() + + +class TestRole(unittest.TestCase): + @classmethod + def setUpClass(cls): + print("tear up") + milvusConnection.connect(uri=uri, token="root:Milvus") + user.create_user(username=username, password="123456") + role.createRole(roleName=roleName) + + @classmethod + def tearDownClass(cls): + role.revokeRole(roleName=roleName, username=username) + user.delete_user(username=username) + role.dropRole(roleName=roleName) + milvusConnection.disconnect() + + def test_create_role(self): + data = role.listRoles() + self.assertIn(roleName, [role.role_name for role in data]) + + def test_grant_role(self): + users = role.grantRole(roleName=roleName, username=username) + self.assertIn(username, users) + + def test_revoke_role(self): + users = role.revokeRole(roleName=roleName, username=username) + self.assertNotIn(username, users) + + def test_grant_privilege(self): + role.grantPrivilege( + roleName=roleName, + objectName=objectName, + objectType="Collection", + privilege="CreateIndex", + ) + data = role.listGrants( + roleName=roleName, objectName=objectName, objectType="Collection" + ) + self.assertIn("CreateIndex", [priv[len(priv) - 1] for priv in data]) + + def test_revoke_privilege(self): + role.revokePrivilege( + roleName=roleName, + objectName=objectName, + objectType="Collection", + privilege="CreateIndex", + ) + data = role.listGrants( + roleName=roleName, objectName=objectName, objectType="Collection" + ) + self.assertNotIn("CreateIndex", [priv[len(priv) - 1] for priv in data]) + + def test_drop_role(self): + tempName = "a" + role.createRole(roleName=tempName) + role.dropRole(roleName=tempName) + data = role.listRoles() + self.assertNotIn(tempName, [role.role_name for role in data]) + + +if __name__ == "__main__": + unittest.main() diff --git a/milvus_cli/utils.py b/milvus_cli/utils.py index 0615277..747749d 100644 --- a/milvus_cli/utils.py +++ b/milvus_cli/utils.py @@ -3,17 +3,10 @@ import os from string import Template from pymilvus import __version__ -from Types import ParameterException def getPackageVersion(): - import pkg_resources - - try: - version = pkg_resources.require("milvus_cli")[0].version - except Exception as e: - raise ParameterException(e) - return version + return __version__ class Completer(object): @@ -21,9 +14,19 @@ class Completer(object): # 'list', 'load', 'query', 'release', 'search', 'show', 'version' ] RE_SPACE = re.compile(".*\s+$", re.M) CMDS_DICT = { + "grant": ["privilege", "role"], + "revoke": ["privilege", "role"], "clear": [], "connect": [], - "create": ["alias", "database", "collection", "partition", "index", "user"], + "create": [ + "alias", + "database", + "collection", + "partition", + "index", + "user", + "role", + ], "delete": [ "alias", "database", @@ -32,6 +35,7 @@ class Completer(object): "partition", "index", "user", + "role", ], "describe": ["collection", "partition", "index"], "exit": [], @@ -45,6 +49,8 @@ class Completer(object): "indexes", "users", "aliases", + "roles", + "grants", ], "load": ["collection", "partition"], "query": [], diff --git a/setup.py b/setup.py index 830773f..b47c07b 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ setup( name="milvus_cli", - version="v1.0.0", + version="v1.0.1", author="Milvus Team", author_email="milvus-team@zilliz.com", url="https://github.com/zilliztech/milvus_cli", @@ -17,7 +17,7 @@ include_package_data=True, install_requires=[ "Click==8.0.1", - "pymilvus==2.4.3", + "pymilvus>=2.4.3", "tabulate==0.8.9", "requests==2.31.0", ],