|
1 | 1 | import asyncio |
| 2 | +import io |
2 | 3 | import json |
3 | 4 | import logging |
4 | 5 | import platform |
|
13 | 14 |
|
14 | 15 | import aiohttp |
15 | 16 | import requests |
16 | | -from aiohttp.client_exceptions import ClientError |
17 | 17 | from lib.core.exceptions import AppException |
18 | 18 | from lib.core.service_types import ServiceResponse |
19 | 19 | from lib.core.serviceproviders import BaseClient |
@@ -229,19 +229,39 @@ class AIOHttpSession(aiohttp.ClientSession): |
229 | 229 | RETRY_LIMIT = 3 |
230 | 230 | BACKOFF_FACTOR = 0.3 |
231 | 231 |
|
232 | | - async def request(self, *args, **kwargs) -> aiohttp.ClientResponse: |
233 | | - attempts = self.RETRY_LIMIT |
234 | | - delay = 0 |
235 | | - for _ in range(attempts): |
236 | | - delay += self.BACKOFF_FACTOR |
237 | | - attempts -= 1 |
238 | | - try: |
239 | | - response = await super()._request(*args, **kwargs) |
240 | | - except ClientError: |
241 | | - if not attempts: |
242 | | - raise |
| 232 | + class AIOHttpSession(aiohttp.ClientSession): |
| 233 | + RETRY_STATUS_CODES = [401, 403, 502, 503, 504] |
| 234 | + RETRY_LIMIT = 3 |
| 235 | + BACKOFF_FACTOR = 0.3 |
| 236 | + |
| 237 | + @staticmethod |
| 238 | + def _copy_form_data(data: aiohttp.FormData) -> aiohttp.FormData: |
| 239 | + form_data = aiohttp.FormData(quote_fields=False) |
| 240 | + for field in data._fields: # noqa |
| 241 | + if isinstance(field[2], io.IOBase): |
| 242 | + field[2].seek(0) |
| 243 | + form_data.add_field( |
| 244 | + value=field[2], |
| 245 | + content_type=field[1].get("Content-Type", ""), |
| 246 | + **field[0], |
| 247 | + ) |
| 248 | + return form_data |
| 249 | + |
| 250 | + async def request(self, *args, **kwargs) -> aiohttp.ClientResponse: |
| 251 | + attempts = self.RETRY_LIMIT |
| 252 | + delay = 0 |
| 253 | + for _ in range(attempts): |
| 254 | + delay += self.BACKOFF_FACTOR |
| 255 | + try: |
| 256 | + response = await super()._request(*args, **kwargs) |
| 257 | + if attempts <= 1 or response.status not in self.RETRY_STATUS_CODES: |
| 258 | + return response |
| 259 | + except (aiohttp.ClientError, RuntimeError) as e: |
| 260 | + if attempts <= 1: |
| 261 | + raise |
| 262 | + if isinstance(e, RuntimeError): |
| 263 | + data = kwargs["data"] |
| 264 | + if isinstance(data, aiohttp.FormData): |
| 265 | + kwargs["data"] = self._copy_form_data(data) |
| 266 | + attempts -= 1 |
243 | 267 | await asyncio.sleep(delay) |
244 | | - continue |
245 | | - if response.status not in self.RETRY_STATUS_CODES or not attempts: |
246 | | - return response |
247 | | - await asyncio.sleep(delay) |
|
0 commit comments