diff --git a/streamrip/client/client.py b/streamrip/client/client.py index f56523b8..dd983c02 100644 --- a/streamrip/client/client.py +++ b/streamrip/client/client.py @@ -8,6 +8,7 @@ import aiolimiter from .downloadable import Downloadable +from .ripClient import StreamRipClient logger = logging.getLogger("streamrip") @@ -52,7 +53,12 @@ def get_rate_limiter( async def get_session(headers: dict | None = None) -> aiohttp.ClientSession: if headers is None: headers = {} - return aiohttp.ClientSession( + # return aiohttp.ClientSession( + # headers={"User-Agent": DEFAULT_USER_AGENT}, + # **headers, + # ) + + return StreamRipClient( headers={"User-Agent": DEFAULT_USER_AGENT}, **headers, ) diff --git a/streamrip/client/ripClient.py b/streamrip/client/ripClient.py new file mode 100644 index 00000000..a1201349 --- /dev/null +++ b/streamrip/client/ripClient.py @@ -0,0 +1,33 @@ +import asyncio +from typing import Any + +import aiohttp +import aiohttp.log + + +class StreamRipClient(aiohttp.ClientSession): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.max_retries = 5 + self.retry_delay = 5 + + async def _request(self, method: str, url: str, **kwargs) -> aiohttp.ClientResponse: + for attempt in range(self.max_retries): + try: + response = await super()._request(method, url, **kwargs) + if response.status != 429: + return response + else: + aiohttp.log.client_logger.warning( + f"Rate limited. Retrying in {self.retry_delay} seconds..." + ) + await asyncio.sleep(self.retry_delay) + except aiohttp.ClientError: + if attempt == self.max_retries - 1: + raise + aiohttp.log.client_logger.warning( + f"Request failed. Retrying in {self.retry_delay} seconds..." + ) + await asyncio.sleep(self.retry_delay) + + raise aiohttp.ClientError(f"Max retries ({self.max_retries}) exceeded") \ No newline at end of file