async_client.py

import asyncio
import json
import random
from typing import Any, Dict

import httpx


class AsyncClient:
    """
    Asynchronous client for making authenticated requests to Openapi endpoints.
    Suitable for use with FastAPI, aiohttp, etc.
    """

    def __init__(
        self,
        token: str,
        client: Any = None,
        timeout: float = 30.0,
        max_retries: int = 0,
        backoff_factor: float = 1.0,
        retry_on_status: list[int] = None,
    ):
        self.client = client if client is not None else httpx.AsyncClient(timeout=timeout)
        self.max_retries = max_retries
        self.backoff_factor = backoff_factor
        self.retry_on_status = (
            retry_on_status if retry_on_status is not None else [429, 502, 503, 504]
        )
        self.auth_header: str = f"Bearer {token}"
        self.headers: Dict[str, str] = {
            "Authorization": self.auth_header,
            "Content-Type": "application/json",
        }

    async def __aenter__(self):
        """Enable use as an asynchronous context manager."""
        return self

    async def __aexit__(self, exc_type, exc_val, exc_tb):
        """Ensure the underlying HTTP client is closed on exit (async)."""
        await self.client.aclose()

    async def aclose(self):
        """Manually close the underlying HTTP client (async)."""
        await self.client.aclose()

    async def _request_with_retry(self, request_fn, *args, **kwargs) -> httpx.Response:
        attempts = 0
        while True:
            try:
                resp = await request_fn(*args, **kwargs)
                if resp.status_code in self.retry_on_status and attempts < self.max_retries:
                    attempts += 1
                    sleep_time = self.backoff_factor * (2 ** attempts) + random.uniform(0, 0.5)
                    if resp.status_code == 429:
                        retry_after = resp.headers.get("Retry-After")
                        if retry_after:
                            try:
                                sleep_time = float(retry_after)
                            except ValueError:
                                pass
                    await asyncio.sleep(sleep_time)
                    continue
                return resp
            except httpx.RequestError as exc:
                if attempts < self.max_retries:
                    attempts += 1
                    sleep_time = self.backoff_factor * (2 ** attempts) + random.uniform(0, 0.5)
                    await asyncio.sleep(sleep_time)
                    continue
                raise exc

    async def request(
        self,
        method: str = "GET",
        url: str = None,
        payload: Dict[str, Any] = None,
        params: Dict[str, Any] = None,
    ) -> Dict[str, Any]:
        """
        Make an asynchronous HTTP request to the specified Openapi endpoint.
        """
        payload = payload or {}
        params = params or {}
        url = url or ""

        if params:
            import urllib.parse
            query_string = urllib.parse.urlencode(params, doseq=True)
            url = f"{url}&{query_string}" if "?" in url else f"{url}?{query_string}"
            params = None

        resp = await self._request_with_retry(
            self.client.request,
            method=method,
            url=url,
            headers=self.headers,
            json=payload,
            params=params,
        )
        data = resp.json()

        # Handle cases where the API might return a JSON-encoded string instead of an object
        if isinstance(data, str):
            try:
                data = json.loads(data)
            except json.JSONDecodeError:
                pass

        return data