support cloud

milvus_client.py

@@ -25,7 +25,8 @@ class MilvusClient(PyOperator):
                 ANNOY: {"params": {"search_k": 10}}.
     """
 
-    def __init__(self, host: str = 'localhost', port: int = 19530, collection_name: str = None, **kwargs):
+    def __init__(self, host: str = 'localhost', port: int = 19530, collection_name: str = None,
+                 user: str = None, password: str = None, **kwargs):
         """
         Get an existing collection.
         """
@@ -33,8 +34,12 @@ class MilvusClient(PyOperator):
         self._port = port
         self._collection_name = collection_name
         self._connect_name = uuid.uuid4().hex
-        connections.connect(alias=self._connect_name, host=self._host, port=self._port)
-        self._collection = Collection(self._collection_name, using=self._connect_name)
+        if None in [user, password]:
+            connections.connect(alias=self._connect_name, host=self._host, port=self._port)
+        else:
+            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)        
 
         self.kwargs = kwargs
         if 'anns_field' not in self.kwargs:
@@ -55,7 +60,8 @@ class MilvusClient(PyOperator):
             'RHNSW_SQ': {'params': {'ef': 10}},
             'RHNSW_PQ': {'params': {'ef': 10}},
             'IVF_HNSW': {'params': {'nprobe': 10, 'ef': 10}},
-            'ANNOY': {'params': {'search_k': 10}}
+            'ANNOY': {'params': {'search_k': 10}},
+            'AUTOINDEX': {}
         }
 
         if 'param' not in self.kwargs: