enhance: Store alias before wait for ready (#2894)

Signed-off-by: yangxuan <[email protected]>

Author: XuanYang-cn

pymilvus/milvus_client/_utils.py

@@ -0,0 +1,52 @@
+import hashlib
+import logging
+
+from pymilvus.orm.connections import connections
+
+logger = logging.getLogger(__name__)
+
+
+def create_connection(
+    uri: str,
+    token: str = "",
+    db_name: str = "",
+    use_async: bool = False,
+    *,
+    user: str = "",
+    password: str = "",
+    **kwargs,
+) -> str:
+    """Get or create the connection to the Milvus server.
+
+    Returns:
+        str: The alias of the connection
+    """
+    using = kwargs.pop("alias", None)
+    if not using:
+        use_async_fmt = "async" if use_async else ""
+
+        auth_fmt = ""
+        if user:
+            auth_fmt = f"{user}"
+        elif token:
+            md5 = hashlib.new("md5", usedforsecurity=False)
+            md5.update(token.encode())
+            auth_fmt = f"{md5.hexdigest()}"
+
+        # different uri, auth, db_name cannot share the same connection
+        not_empty = [v for v in [use_async_fmt, uri, db_name, auth_fmt] if v]
+        using = "-".join(not_empty)
+
+    if connections.has_connection(using):
+        return using
+
+    try:
+        connections.connect(
+            using, user, password, db_name, token, uri=uri, _async=use_async, **kwargs
+        )
+    except Exception as ex:
+        logger.error("Failed to create new connection using: %s", using)
+        raise ex from ex
+    else:
+        logger.debug("Created new connection using: %s", using)
+        return using

pymilvus/milvus_client/async_milvus_client.py

@@ -22,6 +22,7 @@
 from pymilvus.orm.schema import FieldSchema
 from pymilvus.orm.types import DataType
 
+from ._utils import create_connection
 from .check import validate_param
 from .index import IndexParam, IndexParams
 
@@ -43,8 +44,15 @@ def __init__(
         timeout: Optional[float] = None,
         **kwargs,
     ) -> None:
-        self._using = self._create_connection(
-            uri, user, password, db_name, token, timeout=timeout, **kwargs
+        self._using = create_connection(
+            uri,
+            token,
+            db_name,
+            use_async=True,
+            user=user,
+            password=password,
+            timeout=timeout,
+            **kwargs,
         )
         self.is_self_hosted = bool(utility.get_server_type(using=self._using) == "milvus")
 
@@ -815,30 +823,6 @@ async def close(self):
     def _get_connection(self):
         return connections._fetch_handler(self._using)
 
-    def _create_connection(
-        self,
-        uri: str,
-        user: str = "",
-        password: str = "",
-        db_name: str = "",
-        token: str = "",
-        **kwargs,
-    ) -> str:
-        """Create the connection to the Milvus server."""
-        using = kwargs.pop("alias", None)
-        if not using or using == "":
-            using = f"async-{uri}{user}"
-        try:
-            connections.connect(
-                using, user, password, db_name, token, uri=uri, _async=True, **kwargs
-            )
-        except Exception as ex:
-            logger.error("Failed to create new connection using: %s", using)
-            raise ex from ex
-        else:
-            logger.debug("Created new connection using: %s", using)
-            return using
-
     def _extract_primary_field(self, schema_dict: Dict) -> dict:
         fields = schema_dict.get("fields", [])
         if not fields:

pymilvus/milvus_client/milvus_client.py

@@ -1,4 +1,3 @@
-import hashlib
 import logging
 from typing import Dict, List, Optional, Union
 
@@ -28,6 +27,7 @@
 from pymilvus.orm.iterator import QueryIterator, SearchIterator
 from pymilvus.orm.types import DataType
 
+from ._utils import create_connection
 from .check import validate_param
 from .index import IndexParam, IndexParams
 
@@ -62,8 +62,8 @@ def __init__(
                 to None.
                 Unit: second
         """
-        self._using = self._create_connection(
-            uri, user, password, db_name, token, timeout=timeout, **kwargs
+        self._using = create_connection(
+            uri, token, db_name, user=user, password=password, timeout=timeout, **kwargs
         )
         self.is_self_hosted = bool(utility.get_server_type(using=self._using) == "milvus")
 
@@ -919,41 +919,6 @@ def close(self):
     def _get_connection(self):
         return connections._fetch_handler(self._using)
 
-    def _create_connection(
-        self,
-        uri: str,
-        user: str = "",
-        password: str = "",
-        db_name: str = "",
-        token: str = "",
-        **kwargs,
-    ) -> str:
-        """Create the connection to the Milvus server."""
-        using = kwargs.pop("alias", None)
-        if not using or using == "":
-            base_using = f"{uri}-{db_name}"
-            # different user cannot share the same connection
-            if user:
-                using = f"{base_using}-{user}"
-            elif token:
-                # make md5 of token
-                md5 = hashlib.new("md5", usedforsecurity=False)
-                md5.update(token.encode())
-                using = f"{base_using}-{md5.hexdigest()}"
-            else:
-                using = f"{base_using}"
-
-        if connections.has_connection(using):
-            return using
-        try:
-            connections.connect(using, user, password, db_name, token, uri=uri, **kwargs)
-        except Exception as ex:
-            logger.error("Failed to create new connection using: %s", using)
-            raise ex from ex
-        else:
-            logger.debug("Created new connection using: %s", using)
-            return using
-
     def _extract_primary_field(self, schema_dict: Dict) -> dict:
         fields = schema_dict.get("fields", [])
         if not fields: