From 266edabe009286517d2901353673ff844bb7e5ea Mon Sep 17 00:00:00 2001 From: nameczz Date: Tue, 22 Aug 2023 16:12:33 +0800 Subject: [PATCH 1/2] support alias operations Signed-off-by: nameczz --- milvus_cli/Alias.py | 41 +++++++++++++ milvus_cli/Cli.py | 2 + milvus_cli/Collection.py | 1 + milvus_cli/scripts/alias_cli.py | 100 +++++++++++++++++++++++++++++++ milvus_cli/scripts/milvus_cli.py | 1 + milvus_cli/test/test_alias.py | 72 ++++++++++++++++++++++ milvus_cli/utils.py | 9 ++- 7 files changed, 225 insertions(+), 1 deletion(-) create mode 100644 milvus_cli/Alias.py create mode 100644 milvus_cli/scripts/alias_cli.py create mode 100644 milvus_cli/test/test_alias.py diff --git a/milvus_cli/Alias.py b/milvus_cli/Alias.py new file mode 100644 index 0000000..f639ca1 --- /dev/null +++ b/milvus_cli/Alias.py @@ -0,0 +1,41 @@ +from pymilvus import utility + + +class MilvusAlias(object): + alias = "default" + + def create_alias(self, collectionName, aliasName, alias=None): + tempAlias = alias if alias else self.alias + try: + utility.create_alias( + collection_name=collectionName, alias=aliasName, using=tempAlias + ) + return f"Create alias {aliasName} successfully!" + except Exception as e: + raise Exception(f"Create alias error!{str(e)}") + + def list_aliases(self, collectionName, alias=None): + tempAlias = alias if alias else self.alias + try: + res = utility.list_aliases(collection_name=collectionName, using=tempAlias) + return res + except Exception as e: + raise Exception(f"List alias error!{str(e)}") + + def drop_alias(self, aliasName, alias=None): + tempAlias = alias if alias else self.alias + try: + utility.drop_alias(alias=aliasName, using=tempAlias) + return f"Drop alias {aliasName} successfully!" + except Exception as e: + raise Exception(f"Drop alias error!{str(e)}") + + def alter_alias(self, aliasName, collectionName, alias=None): + tempAlias = alias if alias else self.alias + try: + utility.alter_alias( + alias=aliasName, collection_name=collectionName, using=tempAlias + ) + return f"Alter alias {aliasName} successfully!" + except Exception as e: + raise Exception(f"Alter alias error!{str(e)}") diff --git a/milvus_cli/Cli.py b/milvus_cli/Cli.py index 539e120..33cf688 100644 --- a/milvus_cli/Cli.py +++ b/milvus_cli/Cli.py @@ -3,6 +3,7 @@ from Collection import MilvusCollection from Index import MilvusIndex from Data import MilvusData +from Alias import MilvusAlias from pymilvus import __version__ from Types import ParameterException @@ -25,3 +26,4 @@ class MilvusCli(object): collection = MilvusCollection() index = MilvusIndex() data = MilvusData() + alias = MilvusAlias() diff --git a/milvus_cli/Collection.py b/milvus_cli/Collection.py index f1e2b86..7dffa91 100644 --- a/milvus_cli/Collection.py +++ b/milvus_cli/Collection.py @@ -81,6 +81,7 @@ def create_collection( def list_collections(self, alias=None): tempAlias = alias if alias else self.alias + print("--- alias ----", tempAlias) try: res = list_collections(using=tempAlias) except Exception as e: diff --git a/milvus_cli/scripts/alias_cli.py b/milvus_cli/scripts/alias_cli.py new file mode 100644 index 0000000..3b0f22e --- /dev/null +++ b/milvus_cli/scripts/alias_cli.py @@ -0,0 +1,100 @@ +from helper_cli import create, getList, delete +import click + + +@create.command("alias") +@click.option( + "-c", + "--collection-name", + "collectionName", + help="Collection name to be specified alias.", + type=str, +) +@click.option( + "-a", + "--alias-name", + "aliasName", + help="The alias of the collection.", + type=str, +) +@click.option( + "-A", + "--alter", + "alter", + help="[Optional, Flag] - Change an existing alias to current collection.", + default=False, + is_flag=True, +) +@click.pass_obj +def create_alias( + obj, + collectionName, + aliasName, + alter, +): + """ + Specify alias for a collection. + Alias cannot be duplicated, you can't assign same alias to different collections. + But you can specify multiple aliases for a collection, for example: + + create alias -c car -a carAlias1 + + You can also change alias of a collection to another collection. + If the alias doesn't exist, it will return error. + Use "-A" option to change alias owner collection, for example: + + create alias -c car2 -A -a carAlias1 + """ + try: + if alter: + result = obj.alias.alter_alias(collectionName, aliasName) + else: + result = obj.alias.create_alias(collectionName, aliasName) + except Exception as ce: + click.echo("Error!\n{}".format(str(ce))) + else: + click.echo(result) + + +@getList.command("aliases") +@click.option( + "-c", + "--collection-name", + "collectionName", + help="Collection name to be specified alias.", + type=str, +) +@click.pass_obj +def list_aliases(obj, collectionName): + """ + List all aliases of a collection. + Example: + + list aliases -c car + """ + try: + click.echo(obj.alias.list_aliases(collectionName)) + except Exception as e: + click.echo(message=e, err=True) + + +@delete.command("alias") +@click.option( + "-a", + "--alias-name", + "aliasName", + help="The alias of the collection.", + type=str, +) +@click.pass_obj +def delete_alias(obj, aliasName): + """ + Delete alias of a collection. + Example: + + delete alias -a carAlias1 + """ + try: + click.echo(obj.alias.drop_alias(aliasName)) + except Exception as e: + click.echo(message=e, err=True) diff --git a/milvus_cli/scripts/milvus_cli.py b/milvus_cli/scripts/milvus_cli.py index a86828a..bf7f4f1 100644 --- a/milvus_cli/scripts/milvus_cli.py +++ b/milvus_cli/scripts/milvus_cli.py @@ -5,6 +5,7 @@ from collection_cli import * from index_cli import * from data_cli import * +from alias_cli import * if __name__ == "__main__": diff --git a/milvus_cli/test/test_alias.py b/milvus_cli/test/test_alias.py new file mode 100644 index 0000000..217f318 --- /dev/null +++ b/milvus_cli/test/test_alias.py @@ -0,0 +1,72 @@ +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 Collection import MilvusCollection +from Alias import MilvusAlias + +uri = "http://localhost:19530" +tempAlias = "zilliz2" +collectionName = "test_collection" +collectionName2 = "test_collection2" +colAlias = "collection_alias" + +milvusConnection = MilvusConnection() +collection = MilvusCollection() +alias = MilvusAlias() + + +class TestAlias(unittest.TestCase): + @classmethod + def setUpClass(cls): + milvusConnection.connect(uri=uri, alias=tempAlias) + fields = [ + "id:VARCHAR:128", + "title:VARCHAR:512", + "title_vector:FLOAT_VECTOR:768", + ] + collection.create_collection( + collectionName=collectionName, + fields=fields, + alias=tempAlias, + autoId=False, + description="this is a test collection", + primaryField="id", + isDynamic=True, + consistencyLevel="Strong", + ) + + collection.create_collection( + collectionName=collectionName2, + fields=fields, + alias=tempAlias, + autoId=False, + description="this is a test collection", + primaryField="id", + isDynamic=True, + consistencyLevel="Strong", + ) + + @classmethod + def tearDownClass(cls): + collection.drop_collection(collectionName=collectionName, alias=tempAlias) + collection.drop_collection(collectionName=collectionName2, alias=tempAlias) + + milvusConnection.disconnect(alias=tempAlias) + + def test_create_alias(self): + result = alias.create_alias( + collectionName=collectionName, aliasName=colAlias, alias=tempAlias + ) + print(result) + aliasList = alias.list_aliases(collectionName, tempAlias) + self.assertIn(colAlias, aliasList) + + def test_drop_alias(self): + alias.drop_alias(aliasName=colAlias, alias=tempAlias) + aliasList = alias.list_aliases(collectionName, tempAlias) + self.assertNotIn(colAlias, aliasList) diff --git a/milvus_cli/utils.py b/milvus_cli/utils.py index f516bfa..91acf41 100644 --- a/milvus_cli/utils.py +++ b/milvus_cli/utils.py @@ -680,7 +680,14 @@ class Completer(object): "help": [], "import": [], "insert": [], - "list": ["collections", "databases", "partitions", "indexes", "users"], + "list": [ + "collections", + "databases", + "partitions", + "indexes", + "users", + "aliases", + ], "load_balance": [], "load": ["collection", "partition"], "query": [], From 6c4ec1838ae42d290140d940fd5985867438d85f Mon Sep 17 00:00:00 2001 From: nameczz Date: Tue, 22 Aug 2023 16:13:39 +0800 Subject: [PATCH 2/2] remove print Signed-off-by: nameczz --- milvus_cli/Collection.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/milvus_cli/Collection.py b/milvus_cli/Collection.py index 7dffa91..da580ea 100644 --- a/milvus_cli/Collection.py +++ b/milvus_cli/Collection.py @@ -81,7 +81,7 @@ def create_collection( def list_collections(self, alias=None): tempAlias = alias if alias else self.alias - print("--- alias ----", tempAlias) + try: res = list_collections(using=tempAlias) except Exception as e: