From 533a7d8a45ea88ee47c5e0fa614059444de2d5e5 Mon Sep 17 00:00:00 2001 From: Joan Fontanals Date: Tue, 10 Sep 2024 10:02:04 +0200 Subject: [PATCH] refactor: handle async context manager in clientlet (#6194) --- jina/clients/base/helper.py | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) diff --git a/jina/clients/base/helper.py b/jina/clients/base/helper.py index 0a2d7481164a8..1e869c2599c7c 100644 --- a/jina/clients/base/helper.py +++ b/jina/clients/base/helper.py @@ -173,14 +173,15 @@ async def send_message(self, request: 'Request'): from docarray.base_doc.io.json import orjson_dumps request_kwargs['data'] = JinaJsonPayload(value=req_dict) - response = await self.session.post(**request_kwargs).__aenter__() - try: - r_str = await response.json() - except aiohttp.ContentTypeError: - r_str = await response.text() - handle_response_status(response.status, r_str, self.url) - return response - except (ValueError, ConnectionError, BadClient, aiohttp.ClientError) as err: + async with self.session.post(**request_kwargs) as response: + try: + r_str = await response.json() + except aiohttp.ContentTypeError: + r_str = await response.text() + handle_response_status(response.status, r_str, self.url) + return response + except (ValueError, ConnectionError, BadClient, aiohttp.ClientError, aiohttp.ClientConnectionError) as err: + self.logger.debug(f'Got an error: {err} sending POST to {self.url} in attempt {attempt}/{self.max_attempts}') await retry.wait_or_raise_err( attempt=attempt, err=err, @@ -189,6 +190,10 @@ async def send_message(self, request: 'Request'): initial_backoff=self.initial_backoff, max_backoff=self.max_backoff, ) + except Exception as exc: + self.logger.debug( + f'Got a non-retried error: {exc} sending POST to {self.url}') + raise exc async def send_streaming_message(self, doc: 'Document', on: str): """Sends a GET SSE request to the server