diff --git a/.DS_Store b/.DS_Store new file mode 100644 index 0000000..3303398 Binary files /dev/null and b/.DS_Store differ diff --git a/.gitignore b/.gitignore index 0af027c..1a7bbfd 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,7 @@ __pycache__/ *.py[cod] *$py.class .vscode +.DS_Store # C extensions *.so diff --git a/milvus_cli/Connection.py b/milvus_cli/Connection.py new file mode 100644 index 0000000..06dd34f --- /dev/null +++ b/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)}" + diff --git a/milvus_cli/Database.py b/milvus_cli/Database.py new file mode 100644 index 0000000..b995165 --- /dev/null +++ b/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)}" \ No newline at end of file diff --git a/milvus_cli/test/test_connection.py b/milvus_cli/test/test_connection.py new file mode 100644 index 0000000..9498e9d --- /dev/null +++ b/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() \ No newline at end of file diff --git a/milvus_cli/test/test_database.py b/milvus_cli/test/test_database.py new file mode 100644 index 0000000..b6946ff --- /dev/null +++ b/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() \ No newline at end of file diff --git a/milvus_cli/utils.py b/milvus_cli/utils.py index 0542981..21d6b8b 100644 --- a/milvus_cli/utils.py +++ b/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 diff --git a/setup.py b/setup.py index e74c73c..e50e3e1 100644 --- a/setup.py +++ b/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", ],