|
| 1 | +# Copyright (c) 2023 EPAM Systems |
| 2 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 3 | +# you may not use this file except in compliance with the License. |
| 4 | +# You may obtain a copy of the License at |
| 5 | +# |
| 6 | +# https://www.apache.org/licenses/LICENSE-2.0 |
| 7 | +# |
| 8 | +# Unless required by applicable law or agreed to in writing, software |
| 9 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 10 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 11 | +# See the License for the specific language governing permissions and |
| 12 | +# limitations under the License |
| 13 | +# |
| 14 | +# https://www.apache.org/licenses/LICENSE-2.0 |
| 15 | +# |
| 16 | +# Unless required by applicable law or agreed to in writing, software |
| 17 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 18 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 19 | +# See the License for the specific language governing permissions and |
| 20 | +# limitations under the License |
| 21 | + |
| 22 | +"""This module designed to help with asynchronous HTTP request/response handling.""" |
| 23 | + |
| 24 | +import asyncio |
| 25 | +import sys |
| 26 | +from types import TracebackType |
| 27 | +from typing import Coroutine, Any, Optional, Type, Callable |
| 28 | + |
| 29 | +from aenum import Enum |
| 30 | +from aiohttp import ClientSession, ClientResponse, ServerConnectionError, \ |
| 31 | + ClientResponseError |
| 32 | + |
| 33 | +DEFAULT_RETRY_NUMBER: int = 5 |
| 34 | +DEFAULT_RETRY_DELAY: float = 0.005 |
| 35 | +THROTTLING_STATUSES: set = {425, 429} |
| 36 | +RETRY_STATUSES: set = {408, 500, 502, 503, 504, 507}.union(THROTTLING_STATUSES) |
| 37 | + |
| 38 | + |
| 39 | +class RetryClass(int, Enum): |
| 40 | + """Enum contains error types and their retry delay multiply factor as values.""" |
| 41 | + |
| 42 | + SERVER_ERROR = 1 |
| 43 | + CONNECTION_ERROR = 2 |
| 44 | + THROTTLING = 3 |
| 45 | + |
| 46 | + |
| 47 | +class RetryingClientSession: |
| 48 | + """Class uses aiohttp.ClientSession.request method and adds request retry logic.""" |
| 49 | + |
| 50 | + _client: ClientSession |
| 51 | + __retry_number: int |
| 52 | + __retry_delay: float |
| 53 | + |
| 54 | + def __init__( |
| 55 | + self, |
| 56 | + *args, |
| 57 | + max_retry_number: int = DEFAULT_RETRY_NUMBER, |
| 58 | + base_retry_delay: float = DEFAULT_RETRY_DELAY, |
| 59 | + **kwargs |
| 60 | + ): |
| 61 | + """Initialize an instance of the session with arguments. |
| 62 | +
|
| 63 | + To obtain the full list of arguments please see aiohttp.ClientSession.__init__() method. This class |
| 64 | + just bypass everything to the base method, except two local arguments 'max_retry_number' and |
| 65 | + 'base_retry_delay'. |
| 66 | +
|
| 67 | + :param max_retry_number: the maximum number of the request retries if it was unsuccessful |
| 68 | + :param base_retry_delay: base value for retry delay, determine how much time the class will wait after |
| 69 | + an error. Real value highly depends on Retry Class and Retry attempt number, |
| 70 | + since retries are performed in exponential delay manner |
| 71 | + """ |
| 72 | + self._client = ClientSession(*args, **kwargs) |
| 73 | + self.__retry_number = max_retry_number |
| 74 | + self.__retry_delay = base_retry_delay |
| 75 | + |
| 76 | + async def __nothing(self): |
| 77 | + pass |
| 78 | + |
| 79 | + def __sleep(self, retry_num: int, retry_factor: int) -> Coroutine: |
| 80 | + if retry_num > 0: # don't wait at the first retry attempt |
| 81 | + delay = (((retry_factor * self.__retry_delay) * 1000) ** retry_num) / 1000 |
| 82 | + return asyncio.sleep(delay) |
| 83 | + else: |
| 84 | + return self.__nothing() |
| 85 | + |
| 86 | + async def __request( |
| 87 | + self, method: Callable, url, **kwargs: Any |
| 88 | + ) -> ClientResponse: |
| 89 | + """Make a request and retry if necessary. |
| 90 | +
|
| 91 | + The method retries requests depending on error class and retry number. For no-retry errors, such as |
| 92 | + 400 Bad Request it just returns result, for cases where it's reasonable to retry it does it in |
| 93 | + exponential manner. |
| 94 | + """ |
| 95 | + result = None |
| 96 | + exceptions = [] |
| 97 | + for i in range(self.__retry_number + 1): # add one for the first attempt, which is not a retry |
| 98 | + retry_factor = None |
| 99 | + try: |
| 100 | + result = await method(url, **kwargs) |
| 101 | + except Exception as exc: |
| 102 | + exceptions.append(exc) |
| 103 | + if isinstance(exc, ServerConnectionError) or isinstance(exc, ClientResponseError): |
| 104 | + retry_factor = RetryClass.CONNECTION_ERROR |
| 105 | + |
| 106 | + if not retry_factor: |
| 107 | + raise exc |
| 108 | + |
| 109 | + if result: |
| 110 | + if result.ok or result.status not in RETRY_STATUSES: |
| 111 | + return result |
| 112 | + |
| 113 | + if result.status in THROTTLING_STATUSES: |
| 114 | + retry_factor = RetryClass.THROTTLING |
| 115 | + else: |
| 116 | + retry_factor = RetryClass.SERVER_ERROR |
| 117 | + |
| 118 | + if i + 1 < self.__retry_number: |
| 119 | + # don't wait at the last attempt |
| 120 | + await self.__sleep(i, retry_factor) |
| 121 | + |
| 122 | + if exceptions: |
| 123 | + if len(exceptions) > 1: |
| 124 | + if sys.version_info > (3, 10): |
| 125 | + # noinspection PyCompatibility |
| 126 | + raise ExceptionGroup( # noqa: F821 |
| 127 | + 'During retry attempts the following exceptions happened', |
| 128 | + exceptions |
| 129 | + ) |
| 130 | + else: |
| 131 | + raise exceptions[-1] |
| 132 | + else: |
| 133 | + raise exceptions[0] |
| 134 | + return result |
| 135 | + |
| 136 | + def get(self, url: str, *, allow_redirects: bool = True, |
| 137 | + **kwargs: Any) -> Coroutine[Any, Any, ClientResponse]: |
| 138 | + """Perform HTTP GET request.""" |
| 139 | + return self.__request(self._client.get, url, allow_redirects=allow_redirects, **kwargs) |
| 140 | + |
| 141 | + def post(self, url: str, *, data: Any = None, **kwargs: Any) -> Coroutine[Any, Any, ClientResponse]: |
| 142 | + """Perform HTTP POST request.""" |
| 143 | + return self.__request(self._client.post, url, data=data, **kwargs) |
| 144 | + |
| 145 | + def put(self, url: str, *, data: Any = None, **kwargs: Any) -> Coroutine[Any, Any, ClientResponse]: |
| 146 | + """Perform HTTP PUT request.""" |
| 147 | + return self.__request(self._client.put, url, data=data, **kwargs) |
| 148 | + |
| 149 | + def close(self) -> Coroutine: |
| 150 | + """Gracefully close internal aiohttp.ClientSession class instance.""" |
| 151 | + return self._client.close() |
| 152 | + |
| 153 | + async def __aenter__(self) -> "RetryingClientSession": |
| 154 | + """Auxiliary method which controls what `async with` construction does on block enter.""" |
| 155 | + return self |
| 156 | + |
| 157 | + async def __aexit__( |
| 158 | + self, |
| 159 | + exc_type: Optional[Type[BaseException]], |
| 160 | + exc_val: Optional[BaseException], |
| 161 | + exc_tb: Optional[TracebackType], |
| 162 | + ) -> None: |
| 163 | + """Auxiliary method which controls what `async with` construction does on block exit.""" |
| 164 | + await self.close() |
0 commit comments