From 72e8b5af69037c86c7c84388f8bc40d7047bf860 Mon Sep 17 00:00:00 2001 From: Aryan Pardeshi Date: Tue, 11 Aug 2026 01:13:09 +0530 Subject: [PATCH] feat(client): support with/async with on QdrantClient and AsyncQdrantClient QdrantClient and AsyncQdrantClient own gRPC channels, an httpx client and, in local mode, a SQLite database plus a portalocker lockfile, but neither implemented the context manager protocol, so callers had to remember an explicit close() in a finally block. Adds __enter__/__exit__ to the sync client and __aenter__/__aexit__ to the async one. Both delegate to the existing idempotent close(). The async client is generated from the sync one, so the generator now takes a rename_methods map and rewrites __enter__/__exit__ into their async counterparts, forcing them async (they have no counterpart on AsyncQdrantBase) and remapping the string return annotation, which the name-based transformers do not touch. Closes #1285 --- qdrant_client/async_qdrant_client.py | 12 ++++ qdrant_client/qdrant_client.py | 12 ++++ tests/test_context_manager.py | 67 +++++++++++++++++++ .../client_generator.py | 6 ++ .../client/function_def_transformer.py | 17 +++++ 5 files changed, 114 insertions(+) create mode 100644 tests/test_context_manager.py diff --git a/qdrant_client/async_qdrant_client.py b/qdrant_client/async_qdrant_client.py index 4e58bd04d..bec6377a5 100644 --- a/qdrant_client/async_qdrant_client.py +++ b/qdrant_client/async_qdrant_client.py @@ -10,6 +10,7 @@ # ****** WARNING: THIS FILE IS AUTOGENERATED ****** import warnings +from types import TracebackType from typing import Any, Awaitable, Callable, Iterable, Mapping, Sequence import numpy as np from qdrant_client import grpc as grpc @@ -155,6 +156,17 @@ def __init__( is_local_mode=isinstance(self._client, AsyncQdrantLocal), ) + async def __aenter__(self) -> "AsyncQdrantClient": + return self + + async def __aexit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + await self.close() + async def close(self, grpc_grace: float | None = None, **kwargs: Any) -> None: """Closes the connection to Qdrant diff --git a/qdrant_client/qdrant_client.py b/qdrant_client/qdrant_client.py index 5ee0f5c24..c05b4b8d5 100644 --- a/qdrant_client/qdrant_client.py +++ b/qdrant_client/qdrant_client.py @@ -1,4 +1,5 @@ import warnings +from types import TracebackType from typing import ( Any, Awaitable, @@ -168,6 +169,17 @@ def __init__( def __del__(self) -> None: self.close() + def __enter__(self) -> "QdrantClient": + return self + + def __exit__( + self, + exc_type: type[BaseException] | None, + exc_value: BaseException | None, + traceback: TracebackType | None, + ) -> None: + self.close() + def close(self, grpc_grace: float | None = None, **kwargs: Any) -> None: """Closes the connection to Qdrant diff --git a/tests/test_context_manager.py b/tests/test_context_manager.py new file mode 100644 index 000000000..c3da430a3 --- /dev/null +++ b/tests/test_context_manager.py @@ -0,0 +1,67 @@ +import pytest + +from qdrant_client import AsyncQdrantClient, QdrantClient, models + + +def test_sync_context_manager_closes_on_exit(): + with QdrantClient(":memory:") as client: + assert isinstance(client, QdrantClient) + client.create_collection( + collection_name="ctx", + vectors_config=models.VectorParams(size=2, distance=models.Distance.DOT), + ) + assert client._client.closed is False + + assert client._client.closed is True + + +def test_sync_context_manager_closes_when_body_raises(): + client = QdrantClient(":memory:") + + with pytest.raises(RuntimeError, match="boom"): + with client: + raise RuntimeError("boom") + + assert client._client.closed is True + + +def test_sync_context_manager_exit_is_idempotent(): + with QdrantClient(":memory:") as client: + pass + + # A second close() must not raise: QdrantLocal/QdrantRemote.close() are idempotent. + client.close() + assert client._client.closed is True + + +@pytest.mark.asyncio +async def test_async_context_manager_closes_on_exit(): + async with AsyncQdrantClient(":memory:") as client: + assert isinstance(client, AsyncQdrantClient) + await client.create_collection( + collection_name="ctx", + vectors_config=models.VectorParams(size=2, distance=models.Distance.DOT), + ) + assert client._client.closed is False + + assert client._client.closed is True + + +@pytest.mark.asyncio +async def test_async_context_manager_closes_when_body_raises(): + client = AsyncQdrantClient(":memory:") + + with pytest.raises(RuntimeError, match="boom"): + async with client: + raise RuntimeError("boom") + + assert client._client.closed is True + + +@pytest.mark.asyncio +async def test_async_context_manager_exit_is_idempotent(): + async with AsyncQdrantClient(":memory:") as client: + pass + + await client.close() + assert client._client.closed is True diff --git a/tools/async_client_generator/client_generator.py b/tools/async_client_generator/client_generator.py index 3386d1a60..8c5b861db 100644 --- a/tools/async_client_generator/client_generator.py +++ b/tools/async_client_generator/client_generator.py @@ -21,6 +21,7 @@ def __init__( class_replace_map: dict[str, str] | None = None, import_replace_map: dict[str, str] | None = None, exclude_methods: list[str] | None = None, + rename_methods: dict[str, str] | None = None, ): super().__init__() self._async_methods: list[str] | None = None @@ -33,6 +34,7 @@ def __init__( class_replace_map=class_replace_map, exclude_methods=exclude_methods, async_methods=self.async_methods, + rename_methods=rename_methods, ) ) self.transformers.append(ClassDefTransformer(class_replace_map=class_replace_map)) @@ -94,6 +96,10 @@ def get_async_methods(class_obj: type) -> list[str]: "__del__", "migrate", ], + rename_methods={ + "__enter__": "__aenter__", + "__exit__": "__aexit__", + }, ) modified_code = generator.generate(code) diff --git a/tools/async_client_generator/transformers/client/function_def_transformer.py b/tools/async_client_generator/transformers/client/function_def_transformer.py index 1fe2740cd..58ed06f32 100644 --- a/tools/async_client_generator/transformers/client/function_def_transformer.py +++ b/tools/async_client_generator/transformers/client/function_def_transformer.py @@ -10,17 +10,34 @@ def __init__( class_replace_map: dict[str, str] | None = None, exclude_methods: list[str] | None = None, async_methods: list[str] | None = None, + rename_methods: dict[str, str] | None = None, ): super().__init__(keep_sync) self.class_replace_map = class_replace_map if class_replace_map is not None else {} self.exclude_methods = exclude_methods if exclude_methods is not None else [] self.async_methods = async_methods if async_methods is not None else [] + self.rename_methods = rename_methods if rename_methods is not None else {} def _keep_sync(self, name: str) -> bool: + # Renamed methods (e.g. __enter__ -> __aenter__) have no counterpart on + # AsyncQdrantBase, so the async_methods lookup would keep them sync. + if name in self.rename_methods.values(): + return False return name in self.keep_sync or name not in self.async_methods def visit_FunctionDef(self, sync_node: ast.FunctionDef) -> ast.AST | None: if sync_node.name in self.exclude_methods: return None + if sync_node.name in self.rename_methods: + sync_node.name = self.rename_methods[sync_node.name] + # A string return annotation ("QdrantClient") is a Constant, so the + # name-based transformers leave it alone. + if isinstance(sync_node.returns, ast.Constant) and isinstance( + sync_node.returns.value, str + ): + sync_node.returns.value = self.class_replace_map.get( + sync_node.returns.value, sync_node.returns.value + ) + return super().visit_FunctionDef(sync_node)