From 201a4b781fbcb39a81d421e4419e0cceb3c98075 Mon Sep 17 00:00:00 2001 From: "Kaiyuan.Hu" Date: Sat, 22 Jul 2023 18:57:58 +0800 Subject: [PATCH] Support uri and token Signed-off-by: Kaiyuan.Hu --- milvus_client.py | 15 ++++++++++----- 1 file changed, 10 insertions(+), 5 deletions(-) diff --git a/milvus_client.py b/milvus_client.py index 1fb09e8..93ad29e 100644 --- a/milvus_client.py +++ b/milvus_client.py @@ -26,20 +26,25 @@ class MilvusClient(PyOperator): """ def __init__(self, host: str = 'localhost', port: int = 19530, collection_name: str = None, - user: str = None, password: str = None, **kwargs): + uri: str = None, user: str = None, password: str = None, token: str = None, **kwargs): """ Get an existing collection. """ self._host = host self._port = port + self._uri = uri self._collection_name = collection_name self._connect_name = uuid.uuid4().hex - if None in [user, password]: - connections.connect(alias=self._connect_name, host=self._host, port=self._port) - else: + if uri and token: + connections.connect(alias=self._connect_name, uri=self._uri, token=token, secure=True) + elif user and password: connections.connect(alias=self._connect_name, host=self._host, port=self._port, user=user, password=password, secure=True) - self._collection = Collection(self._collection_name, using=self._connect_name) + else: + connections.connect(alias=self._connect_name, host=self._host, port=self._port) + + + self._collection = Collection(self._collection_name, using=self._connect_name) self.kwargs = kwargs if 'anns_field' not in self.kwargs: