Repository: jina-ai/jina
File: jina/clients/base/helper.py (Python)

Code before:
        :param kwargs: keyword args"""
        await self.session.__aexit__(*args, **kwargs)


class HTTPClientlet(AioHttpClientlet):
    """HTTP Client to be used with the streamer"""

    async def send_message(self, request: 'Request'):
        """Sends a POST request to the server

        :param request: request as dict
        :return: send post message
        """
        req_dict = request.to_dict()
        req_dict['exec_endpoint'] = req_dict['header']['exec_endpoint']
        if 'target_executor' in req_dict['header']:
            req_dict['target_executor'] = req_dict['header']['target_executor']
        for attempt in range(1, self.max_attempts + 1):
            try:
                response = await self.session.post(
                    url=self.url, json=req_dict
                ).__aenter__()
                r_str = await response.json()
                handle_response_status(response.status, r_str, self.url)
                return response
            except (ValueError, ConnectionError, BadClient, aiohttp.ClientError, aiohttp.ContentTypeError) as err:
                await retry.wait_or_raise_err(
                    attempt=attempt,
                    err=err,
                    max_attempts=self.max_attempts,
                    backoff_multiplier=self.backoff_multiplier,
                    initial_backoff=self.initial_backoff,
                    max_backoff=self.max_backoff,
                )

    async def send_dry_run(self, **kwargs):
        """Query the dry_run endpoint from Gateway
        :param kwargs: keyword arguments to make sure compatible API with other clients
        :return: send get message
        """
        return await self.session.get(
            url=self.url, timeout=kwargs.get('timeout', None)
        ).__aenter__()

    async def recv_message(self):
        """Receive message for HTTP (sleep)

        :return: await sleep
        """
        return await asyncio.sleep(1e10)


Review comment (refactor):
Why don't we check the content-type to extract the `r_str`. Here the status code is more important or should be checked first before extracting the payload. The underlying issue is a `504` status code with some default response from the Ingress component which is used raise a new error. You can check the `handle_response_status` method also which already handles the `504` http error code.



Code after:
    async def start(self):
        """Create ClientSession and enter context

        :return: self
        """
        with ImportExtensions(required=True):
            import aiohttp

        self.session = aiohttp.ClientSession(
            **self._session_kwargs, trace_configs=self._trace_config
        )
        await self.session.__aenter__()
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        await self.close(exc_type, exc_val, exc_tb)

    async def close(self, *args, **kwargs):
        """Close ClientSession

        :param args: positional args
        :param kwargs: keyword args"""
        await self.session.__aexit__(*args, **kwargs)


class HTTPClientlet(AioHttpClientlet):
    """HTTP Client to be used with the streamer"""

    async def send_message(self, request: 'Request'):
        """Sends a POST request to the server

        :param request: request as dict
        :return: send post message
        """
        req_dict = request.to_dict()
        req_dict['exec_endpoint'] = req_dict['header']['exec_endpoint']
        if 'target_executor' in req_dict['header']:
            req_dict['target_executor'] = req_dict['header']['target_executor']
        for attempt in range(1, self.max_attempts + 1):
            try:
                request_kwargs = {'url': self.url}
                if not docarray_v2:
                    request_kwargs['json'] = req_dict
                else:
                    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()