fix: add grpc_grace param to close in QdrantClient

qdrant · Feb 3, 2024 · 2e61144 · 2e61144
1 parent 5e564e6
commit 2e61144
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 8 deletions.
diff --git a/qdrant_client/async_qdrant_client.py b/qdrant_client/async_qdrant_client.py
@@ -128,12 +128,14 @@ def __init__(
             except ImportError:
                 self._is_fastembed_installed = False
 
-    async def close(self, **kwargs: Any) -> None:
-        """
-        Closes the connection to Qdrant
+    async def close(self, grpc_grace: Optional[float] = None, **kwargs: Any) -> None:
+        """Closes the connection to Qdrant
+
+        Args:
+            grpc_grace: Grace period for gRPC connection close. Default: None
         """
         if hasattr(self, "_client"):
-            await self._client.close(**kwargs)
+            await self._client.close(grpc_grace=grpc_grace, **kwargs)
 
     @property
     def grpc_collections(self) -> grpc.CollectionsStub:

diff --git a/qdrant_client/qdrant_client.py b/qdrant_client/qdrant_client.py
@@ -131,12 +131,14 @@ def __init__(
     def __del__(self) -> None:
         self.close()
 
-    def close(self, **kwargs: Any) -> None:
-        """
-        Closes the connection to Qdrant
+    def close(self, grpc_grace: Optional[float] = None, **kwargs: Any) -> None:
+        """Closes the connection to Qdrant
+
+        Args:
+            grpc_grace: Grace period for gRPC connection close. Default: None
         """
         if hasattr(self, "_client"):
-            self._client.close(**kwargs)
+            self._client.close(grpc_grace=grpc_grace, **kwargs)
 
     @property
     def grpc_collections(self) -> grpc.CollectionsStub: