Rate limit retry: initial logic

.gitignore

@@ -5,6 +5,7 @@ site/
 .pdm-python
 .token
 .cred
+.env
 .cache_ggshield
 
 # common name for local Sony Ci config file 

pyproject.toml

@@ -7,7 +7,7 @@ dependencies = [
     "requests-oauth2client~=1.6",
     "loguru~=0.7",
 ]
-requires-python = '>=3.8'
+requires-python = '>=3.9'
 readme = 'README.md'
 license = { text = 'MIT' }
 dynamic = ['version']

sonyci/_version.py

@@ -1 +1 @@
-__version__ = '0.3.0'
+__version__ = '0.4.0'

sonyci/cli.py

@@ -1,12 +1,11 @@
 from json import dumps, loads
 from pathlib import Path
-from typing import Optional
+from typing import Annotated, Optional
 from urllib.request import urlretrieve
 
 from requests_oauth2client.tokens import BearerToken, BearerTokenSerializer
 from typer import Argument, Context, Exit, Option, Typer
 from typer.main import get_group
-from typing_extensions import Annotated
 
 from sonyci import SonyCi
 from sonyci.log import log

sonyci/exceptions.py

@@ -0,0 +1,2 @@
+class RetryError(Exception):
+    """Raised when a function fails after retrying a set number of times."""

sonyci/sonyci.py

@@ -6,7 +6,7 @@
 
 from sonyci.config import Config
 from sonyci.log import log
-from sonyci.utils import get_token, json
+from sonyci.utils import get_token, json, retry
 
 
 class SonyCi(Config):
@@ -95,12 +95,15 @@ def asset_download(self, asset_id: str, **kwargs) -> dict:
         return self.get(f'/assets/{asset_id}/download', params=kwargs)
 
     @json
+    @retry
     def get(self, *args, **kwargs):
         log.debug(f'GET {args} {kwargs}')
         return self.client.get(*args, **kwargs)
 
     @json
+    @retry
     def post(self, *args, **kwargs):
+        log.debug(f'POST {args} {kwargs}')
         return self.client.post(*args, **kwargs)
 
     def __call__(self, path: str, **kwds: Any) -> Any:

sonyci/utils.py

@@ -2,6 +2,7 @@
 from requests_oauth2client.tokens import BearerToken, BearerTokenSerializer
 
 from sonyci.config import TOKEN_URL
+from sonyci.exceptions import RetryError
 from sonyci.log import log
 
 
@@ -48,3 +49,33 @@ def inner(*args, **kwargs):
         return func(*args, **kwargs).json()
 
     return inner
+
+
+def retry(func):
+    """Decorator for retrying a function call after a rate limit error."""
+    from requests import HTTPError
+
+    def inner(*args, **kwargs):
+        tries: int = 0
+        max_tries: int = 5
+        while tries < max_tries:
+            try:
+                return func(*args, **kwargs)
+            except HTTPError as e:
+                if e.response.status_code != 429:
+                    log.error(f'HTTPError {e.response.status_code}: {e}')
+                    raise e
+                from time import sleep
+
+                # Get the retry-after header, if it exists
+                retry_after = e.response.headers.get('Retry-After')
+                if not retry_after:
+                    log.error('No Retry-After header found')
+                    raise e
+                log.warning(f'Rate limited. Retrying after {retry_after} seconds...')
+                sleep(int(retry_after + 1))
+                tries += 1
+        log.error(f'Failed after {max_tries} tries')
+        raise RetryError(f'Failed after {max_tries} tries')
+
+    return inner

tests/cli/test_cli_login.py

@@ -45,8 +45,8 @@ def test_bad_login(error_runner):
     assert 'invalid_client' in str(result.exception)
 
 
-def test_missing_username(runner):
-    result = runner.invoke(
+def test_missing_username(error_runner):
+    result = error_runner.invoke(
         app,
         [
             '--client-id',
@@ -59,4 +59,4 @@ def test_missing_username(runner):
         ],
     )
     assert result.exit_code == 2
-    assert 'username' in result.stdout
+    assert '--username' in result.stderr