parent class of aiohttp.ClientSession to add retries

streamrip/client/client.py

@@ -8,6 +8,7 @@ import aiohttp
 import aiolimiter
 
 from .downloadable import Downloadable
+from .ripClient import StreamRipClient
 
 logger = logging.getLogger("streamrip")
 
@@ -52,7 +53,12 @@ class Client(ABC):
     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,
         )

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")