Merge pull request #33 from zilliztech/connection

init connection,database api and test

.gitignore

@@ -3,6 +3,7 @@ __pycache__/
 *.py[cod]
 *$py.class
 .vscode
+.DS_Store
 
 # C extensions
 *.so

milvus_cli/Connection.py

@@ -0,0 +1,63 @@
+from pymilvus import __version__,connections,list_collections
+from Types import ConnectException
+from tabulate import tabulate
+
+class MilvusConnection(object):
+  uri = "127.0.0.1:19530"
+  alias = "default"
+
+  def connect(self,
+          alias='',
+          uri=None,
+          username=None,
+          password=None):
+    self.alias = alias
+    self.uri = uri
+    trimUsername = None if username is None else username.strip()
+    trimPwd = None if password is None else password.strip()
+
+    try:
+      res = connections.connect(alias= self.alias,
+                            uri=self.uri,
+                            user=trimUsername,
+                            password=trimPwd
+                            )  
+      return res
+    except Exception as e:
+        raise ConnectException(f"Connect to Milvus error!{str(e)}")
+
+  def checkConnection(self,alias=None):
+    try:
+      tempAlias = alias if alias else self.alias  
+      collections = list_collections(timeout=10.0, using=tempAlias)
+      return collections
+    except Exception as e:
+      raise ConnectException(f"Connect to Milvus error!{str(e)}")
+
+  def showConnection(self, alias=None, showAll=False):
+    tempAlias = alias if alias else self.alias
+    allConnections = connections.list_connections()
+
+    if showAll:
+        return tabulate(allConnections,
+                        headers=["Alias"],
+                        tablefmt="pretty")
+    aliasList = map(lambda x: x[0], allConnections)
+
+    if tempAlias in aliasList:
+        address, user = connections.get_connection_addr(tempAlias).values()
+        return tabulate(
+            [["Address", address], ["User", user], ["Alias", tempAlias]],
+            tablefmt="pretty",
+        )
+    else:
+        return "Connection not found!"
+
+  def disconnect(self,alias=None):
+    tempAlias = alias if alias else self.alias
+    try:
+      connections.disconnect(alias=tempAlias)
+      return f"Disconnect from {tempAlias} successfully!"
+    except Exception as e:
+      raise f"Disconnect from {tempAlias} error!{str(e)}"
+

milvus_cli/Database.py

@@ -0,0 +1,37 @@
+from pymilvus import db
+from Types import ConnectException
+from tabulate import tabulate
+
+class Database():
+  alias = "default"
+  def create_database(self,dbName=None,alias=None):
+    tempAlias = alias if alias else self.alias
+    try:
+      db.create_database(db_name=dbName,using=tempAlias)
+      return f"Create database {dbName} successfully!"
+    except Exception as e:
+      raise f"Create database error!{str(e)}"
+
+  def list_databases(self,alias=None):
+    tempAlias = alias if alias else self.alias
+    try:
+      res = db.list_database(using=tempAlias)
+      return res
+    except Exception as e:
+      raise f"List database error!{str(e)}"
+
+  def drop_database(self,dbName=None,alias=None):
+    tempAlias = alias if alias else self.alias
+    try:
+      db.drop_database(db_name=dbName,using=tempAlias)
+      return f"Drop database {dbName} successfully!"
+    except Exception as e:
+      raise f"Drop database error!{str(e)}"
+
+  def using_database(self,dbName=None,alias=None):
+    tempAlias = alias if alias else self.alias
+    try:
+      db.using_database(db_name=dbName,using=tempAlias)
+      return f"Using database {dbName} successfully!"
+    except Exception as e:
+      raise f"Using database error!{str(e)}"

milvus_cli/test/test_connection.py

@@ -0,0 +1,42 @@
+import unittest;
+import sys
+import os
+from tabulate import tabulate
+
+currentdir = os.path.dirname(os.path.realpath(__file__))
+parentdir = os.path.dirname(currentdir)
+sys.path.append(parentdir)
+from Connection import MilvusConnection
+
+uri = "http://localhost:19530"
+tempAlias = "zilliz"
+milvusConnection = MilvusConnection()
+
+class TestConnection(unittest.TestCase):
+
+  def setUp(self):
+    milvusConnection.connect(uri=uri ,alias=tempAlias)
+
+  def tearDown(self):
+    milvusConnection.disconnect(alias=tempAlias)
+
+  def test_show_connection(self):
+    res = milvusConnection.showConnection(alias=tempAlias)
+    expectRes = tabulate(
+            [["Address", f'localhost:19530'], ["User", ''], ["Alias", tempAlias]],
+            tablefmt="pretty",
+        )
+    self.assertEqual(res, expectRes)
+
+  def test_disconnect(self):
+    res = milvusConnection.disconnect(alias=tempAlias)
+    expectRes = f"Disconnect from {tempAlias} successfully!"
+    self.assertEqual(res, expectRes)
+
+  def test_check_connection(self):
+    res = milvusConnection.checkConnection(alias=tempAlias)
+    self.assertEqual(len(res), 0)
+
+
+if __name__ == '__main__':
+  unittest.main()

milvus_cli/test/test_database.py

@@ -0,0 +1,34 @@
+import unittest;
+import sys
+import os
+from tabulate import tabulate
+
+currentdir = os.path.dirname(os.path.realpath(__file__))
+parentdir = os.path.dirname(currentdir)
+sys.path.append(parentdir)
+from Connection import MilvusConnection
+from Database import Database
+
+uri = "http://localhost:19530"
+tempAlias = "zilliz"
+dbName="test_db"
+milvusConnection = MilvusConnection()
+database = Database()
+
+class TestDatabase(unittest.TestCase):
+
+  def setUp(self):
+    milvusConnection.connect(uri=uri,alias=tempAlias)
+
+  def tearDown(self):
+    database.drop_database(dbName=dbName,alias=tempAlias)
+    milvusConnection.disconnect(alias=tempAlias)
+
+  def test_create_database(self):
+    database.create_database(dbName=dbName,alias=tempAlias)
+    databaseList = database.list_databases(alias=tempAlias)
+    self.assertIn(dbName,databaseList)
+
+
+if __name__ == '__main__':
+  unittest.main()

milvus_cli/utils.py

@@ -25,12 +25,6 @@ def checkEmpty(x):
     return not not x
 
 
-def getMilvusTimestamp(isSimilar=True):
-    ts = time()
-    if isSimilar:
-        return int(ts) << 18
-
-
 class PyOrm(object):
     host = "127.0.0.1"
     port = 19530

setup.py

@@ -17,7 +17,7 @@
     include_package_data=True,
     install_requires=[
         "Click==8.0.1",
-        "pymilvus==2.2.1",
+        "pymilvus==2.2.13",
         "tabulate==0.8.9",
         "requests==2.26.0",
     ],