Merge pull request #38 from zilliztech/collection-cli

Support collection cli

milvus_cli/Cli.py

@@ -1,19 +1,23 @@
 from Connection import MilvusConnection
 from Database import MilvusDatabase
+from Collection import MilvusCollection
 from pymilvus import __version__
 from Types import ParameterException
 
+
 def getPackageVersion():
     import pkg_resources  # part of setuptools
+
     try:
         version = pkg_resources.require("milvus_cli")[0].version
     except Exception as e:
         raise ParameterException(
-            "Could not get version under single executable file mode.")
+            "Could not get version under single executable file mode."
+        )
     return version
 
-class MilvusCli(object):
-  connection = MilvusConnection()
-  database = MilvusDatabase()
-
 
+class MilvusCli(object):
+    connection = MilvusConnection()
+    database = MilvusDatabase()
+    collection = MilvusCollection()

milvus_cli/Collection.py

@@ -8,7 +8,7 @@
     CollectionSchema,
 )
 from tabulate import tabulate
-from milvus_cli.Types import DataTypeByNum
+from Types import DataTypeByNum
 
 
 class MilvusCollection(object):
@@ -22,8 +22,8 @@ def create_collection(
         autoId=None,
         description=None,
         isDynamic=None,
-        consistencyLevel=None,
-        shardsNum=2,
+        consistencyLevel="Bounded",
+        shardsNum=1,
         alias=None,
     ):
         fieldList = []

milvus_cli/scripts/collection_cli.py

@@ -0,0 +1,234 @@
+from tabulate import tabulate
+from helper_cli import create, getList, delete, rename, show
+import click
+import os
+import sys
+
+current_dir = os.path.dirname(os.path.realpath(__file__))
+parent_dir = os.path.dirname(current_dir)
+sys.path.append(parent_dir)
+
+from Validation import validateCollectionParameter
+from Types import ParameterException, ConnectException
+
+
+@create.command("collection")
+@click.option(
+    "-c",
+    "--collection-name",
+    "collectionName",
+    help="Collection name to specify alias.",
+    type=str,
+)
+@click.option(
+    "-p", "--schema-primary-field", "primaryField", help="Primary field name."
+)
+@click.option(
+    "-A",
+    "--schema-auto-id",
+    "autoId",
+    help="[Optional, Flag] - Enable auto id.",
+    default=False,
+    is_flag=True,
+)
+@click.option(
+    "-a",
+    "--alias",
+    "alias",
+    help="The connection alias name.",
+    type=str,
+)
+@click.option(
+    "-desc",
+    "--schema-description",
+    "description",
+    help="[Optional] - Description details.",
+    default="",
+)
+@click.option(
+    "-d",
+    "--is-dynamic",
+    "isDynamic",
+    help="[Optional] - Collection schema supports dynamic fields or not.",
+    default=None,
+)
+@click.option(
+    "-level",
+    "--consistency-level",
+    "consistencyLevel",
+    help="[Optional] - Consistency level: Bounded,Session,Strong, Eventual .",
+    default="Bounded",
+)
+@click.option(
+    "-f",
+    "--schema-field",
+    "fields",
+    help='[Multiple] - FieldSchema. Usage is "<Name>:<DataType>:<Dim(if vector) or Description>"',
+    default=None,
+    multiple=True,
+)
+@click.option(
+    "-s",
+    "--shards-num",
+    "shardsNum",
+    help="[Optional] - Shards number",
+    default=1,
+    type=int,
+)
+@click.pass_obj
+def create_collection(
+    obj,
+    collectionName,
+    primaryField,
+    autoId,
+    description,
+    fields,
+    isDynamic,
+    consistencyLevel,
+    shardsNum,
+    alias,
+):
+    """
+    Create collection.
+
+    Example:
+
+      create collection -c car -f id:INT64:primary_field -f vector:FLOAT_VECTOR:128 -f color:INT64:color -f brand:INT64:brand -p id -A -d 'car_collection'
+    """
+    try:
+        validateCollectionParameter(
+            collectionName,
+            primaryField,
+            fields,
+        )
+    except ParameterException as pe:
+        click.echo("Error!\n{}".format(str(pe)))
+    else:
+        click.echo(
+            obj.collection.create_collection(
+                collectionName,
+                primaryField,
+                fields,
+                autoId,
+                description,
+                isDynamic,
+                consistencyLevel,
+                shardsNum,
+                alias,
+            )
+        )
+        click.echo("Create collection successfully!")
+
+
+@getList.command("collections")
+@click.option(
+    "-a",
+    "--alias",
+    "alias",
+    help="The connection alias name.",
+    type=str,
+)
+@click.pass_obj
+def collections(obj, alias):
+    """List all collections."""
+    try:
+        click.echo(obj.collection.list_collections(alias))
+    except Exception as e:
+        click.echo(message=e, err=True)
+
+
+@delete.command("collection")
+@click.option(
+    "-a",
+    "--alias",
+    "alias",
+    help="[Optional]: The connection alias name.",
+    type=str,
+)
+@click.option(
+    "-c",
+    "--collection-name",
+    "collectionName",
+    help="The name of collection to be deleted.",
+    required=True,
+)
+@click.pass_obj
+def delete_collection(obj, collectionName, alias):
+    """
+    Drops the collection together with its index files.
+
+    Example:
+
+        milvus_cli > delete collection -c car
+    """
+    click.echo(
+        "Warning!\nYou are trying to delete the collection with data. This action cannot be undone!\n"
+    )
+    if not click.confirm("Do you want to continue?"):
+        return
+    try:
+        obj.collection.has_collection(alias, collectionName)
+    except Exception as e:
+        click.echo(f"Error occurred when get collection by name!\n{str(e)}")
+    else:
+        result = obj.collection.drop_collection(alias, collectionName)
+        click.echo(result)
+
+
+@rename.command("collection")
+@click.option(
+    "-a",
+    "--alias",
+    "alias",
+    help="The connection alias name.",
+    type=str,
+)
+@click.option(
+    "-old",
+    "--old-collection-name",
+    "collectionName",
+    help="The old collection name",
+    type=str,
+    required=True,
+)
+@click.option(
+    "-new",
+    "--new-collection-name",
+    "newName",
+    help="The new collection name",
+    type=str,
+    required=True,
+)
+@click.pass_obj
+def rename_collection(obj, alias, collectionName, newName):
+    """
+    Rename collection.
+
+    Example:
+
+        milvus_cli > rename collection -old car -new new_car
+    """
+    try:
+        obj.collection.has_collection(alias, collectionName)
+    except Exception as e:
+        click.echo(f"Error occurred when get collection by name!\n{str(e)}")
+    else:
+        result = obj.collection.rename_collection(alias, collectionName, newName)
+        click.echo(result)
+
+
+@show.command("collection")
+@click.option("-c", "--collection-name", "collection", help="The name of collection.")
+@click.pass_obj
+def describeCollection(obj, collection):
+    """
+    Describe collection.
+
+    Example:
+
+        milvus_cli > describe collection -c test_collection_insert
+    """
+    try:
+        click.echo(obj.collection.get_collection_details(collection))
+    except Exception as e:
+        click.echo(message=e, err=True)

milvus_cli/scripts/helper_cli.py

@@ -37,7 +37,7 @@ def clear():
 @cli.group(no_args_is_help=False)
 @click.pass_obj
 def show(obj):
-    """Show connection, database, loading_progress and index_progress."""
+    """Show connection, database,collection, loading_progress and index_progress."""
     pass
 
 
@@ -48,6 +48,13 @@ def getList(obj):
     pass
 
 
+@cli.group("rename", no_args_is_help=False)
+@click.pass_obj
+def rename(obj):
+    """Rename collection"""
+    pass
+
+
 @cli.group("create", no_args_is_help=False)
 @click.pass_obj
 def create(obj):

milvus_cli/scripts/milvus_cli.py

@@ -2,6 +2,7 @@
 from connection_cli import connect, connection
 from helper_cli import help, version, clear, show, runCliPrompt
 from database_cli import create_database, list_databases, drop_database, use_database
+from collection_cli import create_collection
 
 cli.add_command(help)
 cli.add_command(version)
@@ -16,6 +17,8 @@
 cli.add_command(drop_database)
 cli.add_command(use_database)
 
+cli.add_command(create_collection)
+
 
 if __name__ == "__main__":
     runCliPrompt()