Client should handle rate limiting errors gracefully #76 (#96)

Co-authored-by: 20C <[email protected]>

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 
 ## Unreleased
+### Fixed
+- Client should handle rate limiting errors gracefully
 
 
 ## 2.1.1

CHANGELOG.yaml

@@ -1,6 +1,7 @@
 Unreleased:
   added: []
-  fixed: []
+  fixed:
+  - Client should handle rate limiting errors gracefully #76
   changed: []
   deprecated: []
   removed: []

src/peeringdb/config.py

@@ -99,6 +99,12 @@ class LogSchema(_schema.Schema):
         )
         level = _schema.Str("log_level", default=os.environ.get("LOG_LEVEL", "INFO"))
 
+    class LogSchema(_schema.Schema):
+        allow_other_loggers = _schema.Int(
+            "allow_other_loggers", default=os.environ.get("ALLOW_OTHER_LOGGERS", 0)
+        )
+        level = _schema.Str("log_level", default=os.environ.get("LOG_LEVEL", "INFO"))
+
     sync = SyncSchema()
     orm = OrmSchema()
     log = LogSchema()

src/peeringdb/fetch.py

@@ -26,6 +26,7 @@ def __init__(
         :param api_key: API key
         :param cache_url: PeeringDB cache URL
         :param cache_dir: Local cache directory
+        :param retry: The maximum number of retry attempts when rate limited (default is 5)
         :param kwargs:
         """
         self._log = logging.getLogger(__name__)
@@ -45,6 +46,9 @@ def __init__(
         self.remote_cache_used = False
         self.local_cache_used = False
 
+        # used for sync 429 status code (pause and resume)
+        self.attempt = 0
+
     def _get(self, endpoint: str, **params):
         url = f"{self.url}/{endpoint}"
         url_params = urllib.parse.urlencode(params)
@@ -60,25 +64,36 @@ def _get(self, endpoint: str, **params):
                 + base64.b64encode(f"{self.user}:{self.password}".encode()).decode()
             }
 
-        resp = requests.get(url, timeout=self.timeout, headers=headers)
-
-        if resp.status_code == 429:
-            raise ValueError(f"Rate limited: {resp.text}")
-        elif resp.status_code == 400:
-            error = resp.json()["meta"]["error"]
-            if re.search("client version is incompatible", error):
-                raise CompatibilityError(error)
-            raise
-        elif resp.status_code != 200:
-            raise ValueError(f"Error fetching {url}: {resp.status_code}")
-        return resp.json()["data"]
+        while True:
+            try:
+                resp = requests.get(url, timeout=self.timeout, headers=headers)
+                resp.raise_for_status()
+                return resp.json()["data"]
+            except requests.exceptions.HTTPError as e:
+                if resp.status_code == 429:
+                    retry_after = min(2**self.attempt, 60)
+                    self._log.info(
+                        f"Rate limited. Retrying in {retry_after} seconds..."
+                    )
+                    time.sleep(retry_after)
+                    self.attempt += 1
+                elif resp.status_code == 400:
+                    error = resp.json().get("meta", {}).get("error", "")
+                    if re.search("client version is incompatible", error):
+                        raise CompatibilityError(error)
+                    raise ValueError(f"Bad request error: {error}")
+                else:
+                    raise ValueError(f"Error fetching {url}: {resp.status_code}")
+            except requests.exceptions.RequestException as err:
+                raise ValueError(f"Request error: {err}")
 
     def load(
         self,
         resource: str,
         since: int = 0,
         fetch_private: bool = False,
         initial_private: bool = False,
+        delay: float = 0.5,
     ):
         """
         Load a resource from mock data.
@@ -142,6 +157,8 @@ def load(
             else:
                 self.resources[resource] = self._get(resource, since=since)
 
+            time.sleep(delay)
+
     def entries(self, tag: str):
         """
         Get all entries by tag ro load it if we don't already have the resource

tests/test_cli.py

@@ -1,6 +1,8 @@
 import io
 import json
 import re
+import time
+from unittest.mock import MagicMock, patch
 
 import helper
 import pytest
@@ -160,3 +162,52 @@ def test_verbosity(runcli, client, capsys):
 
     # Verbose output should be longer
     assert len(outq) < len(outv)
+
+
+@patch("time.sleep", return_value=None)
+def test_rate_limit_handling(mock_sleep):
+    attempt = 0
+    log_mock = MagicMock()
+
+    # Mock response with status code 429
+    mock_resp = MagicMock()
+    mock_resp.status_code = 429
+
+    # Test rate limit handling
+    for _ in range(10):
+        if mock_resp.status_code == 429:
+            retry_after = min(2**attempt, 60)
+            log_mock.info(f"Rate limited. Retrying in {retry_after} seconds...")
+            time.sleep(retry_after)
+            attempt += 1
+
+    # Assert log calls and sleep durations
+    expected_calls = [
+        (("Rate limited. Retrying in 1 seconds...",),),
+        (("Rate limited. Retrying in 2 seconds...",),),
+        (("Rate limited. Retrying in 4 seconds...",),),
+        (("Rate limited. Retrying in 8 seconds...",),),
+        (("Rate limited. Retrying in 16 seconds...",),),
+        (("Rate limited. Retrying in 32 seconds...",),),
+        (("Rate limited. Retrying in 60 seconds...",),),
+        (("Rate limited. Retrying in 60 seconds...",),),
+        (("Rate limited. Retrying in 60 seconds...",),),
+        (("Rate limited. Retrying in 60 seconds...",),),
+    ]
+
+    assert log_mock.info.call_args_list == expected_calls
+
+    expected_sleep_calls = [
+        ((1,),),
+        ((2,),),
+        ((4,),),
+        ((8,),),
+        ((16,),),
+        ((32,),),
+        ((60,),),
+        ((60,),),
+        ((60,),),
+        ((60,),),
+    ]
+
+    assert mock_sleep.call_args_list == expected_sleep_calls