diff --git a/pymilvus/client/async_grpc_handler.py b/pymilvus/client/async_grpc_handler.py index 11cf4b686..8a3301cbd 100644 --- a/pymilvus/client/async_grpc_handler.py +++ b/pymilvus/client/async_grpc_handler.py @@ -108,9 +108,9 @@ def __enter__(self): def __exit__(self: object, exc_type: object, exc_val: object, exc_tb: object): pass - def close(self): + async def close(self): self.deregister_state_change_callbacks() - self._async_channel.close() + await self._async_channel.close() def _setup_authorization_interceptor(self, user: str, password: str, token: str): keys = [] diff --git a/pymilvus/milvus_client/async_milvus_client.py b/pymilvus/milvus_client/async_milvus_client.py index 61d5f71db..efb507158 100644 --- a/pymilvus/milvus_client/async_milvus_client.py +++ b/pymilvus/milvus_client/async_milvus_client.py @@ -516,8 +516,8 @@ def create_schema(cls, **kwargs): kwargs["check_fields"] = False # do not check fields for now return CollectionSchema([], **kwargs) - def close(self): - connections.disconnect(self._using) + async def close(self): + await connections.async_disconnect(self._using) def _get_connection(self): return connections._fetch_handler(self._using) diff --git a/pymilvus/orm/connections.py b/pymilvus/orm/connections.py index 151b31e7d..5add5cb76 100644 --- a/pymilvus/orm/connections.py +++ b/pymilvus/orm/connections.py @@ -285,6 +285,13 @@ def disconnect(self, alias: str): if alias in self._connected_alias: self._connected_alias.pop(alias).close() + async def async_disconnect(self, alias: str): + if not isinstance(alias, str): + raise ConnectionConfigException(message=ExceptionsMessage.AliasType % type(alias)) + + if alias in self._connected_alias: + await self._connected_alias.pop(alias).close() + def remove_connection(self, alias: str): """Removes connection from the registry.