Skip to content

Commit 44b08a1

Browse files
committed
feat(async-client): add retry mechanism to AsyncClient
1 parent 11618d4 commit 44b08a1

1 file changed

Lines changed: 44 additions & 2 deletions

File tree

openapi_python_sdk/async_client.py

Lines changed: 44 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
1+
import asyncio
12
import json
3+
import random
24
from typing import Any, Dict
35

46
import httpx
@@ -10,8 +12,21 @@ class AsyncClient:
1012
Suitable for use with FastAPI, aiohttp, etc.
1113
"""
1214

13-
def __init__(self, token: str, client: Any = None, timeout: float = 30.0):
15+
def __init__(
16+
self,
17+
token: str,
18+
client: Any = None,
19+
timeout: float = 30.0,
20+
max_retries: int = 0,
21+
backoff_factor: float = 1.0,
22+
retry_on_status: list[int] = None,
23+
):
1424
self.client = client if client is not None else httpx.AsyncClient(timeout=timeout)
25+
self.max_retries = max_retries
26+
self.backoff_factor = backoff_factor
27+
self.retry_on_status = (
28+
retry_on_status if retry_on_status is not None else [429, 502, 503, 504]
29+
)
1530
self.auth_header: str = f"Bearer {token}"
1631
self.headers: Dict[str, str] = {
1732
"Authorization": self.auth_header,
@@ -30,6 +45,32 @@ async def aclose(self):
3045
"""Manually close the underlying HTTP client (async)."""
3146
await self.client.aclose()
3247

48+
async def _request_with_retry(self, request_fn, *args, **kwargs) -> httpx.Response:
49+
attempts = 0
50+
while True:
51+
try:
52+
resp = await request_fn(*args, **kwargs)
53+
if resp.status_code in self.retry_on_status and attempts < self.max_retries:
54+
attempts += 1
55+
sleep_time = self.backoff_factor * (2 ** attempts) + random.uniform(0, 0.5)
56+
if resp.status_code == 429:
57+
retry_after = resp.headers.get("Retry-After")
58+
if retry_after:
59+
try:
60+
sleep_time = float(retry_after)
61+
except ValueError:
62+
pass
63+
await asyncio.sleep(sleep_time)
64+
continue
65+
return resp
66+
except httpx.RequestError as exc:
67+
if attempts < self.max_retries:
68+
attempts += 1
69+
sleep_time = self.backoff_factor * (2 ** attempts) + random.uniform(0, 0.5)
70+
await asyncio.sleep(sleep_time)
71+
continue
72+
raise exc
73+
3374
async def request(
3475
self,
3576
method: str = "GET",
@@ -50,7 +91,8 @@ async def request(
5091
url = f"{url}&{query_string}" if "?" in url else f"{url}?{query_string}"
5192
params = None
5293

53-
resp = await self.client.request(
94+
resp = await self._request_with_retry(
95+
self.client.request,
5496
method=method,
5597
url=url,
5698
headers=self.headers,

0 commit comments

Comments
 (0)