feat: retry transient gateway 5xx via urllib3 Retry (#139)

datamaxi/api.py

@@ -3,6 +3,8 @@
 from json import JSONDecodeError
 import logging
 import requests
+from requests.adapters import HTTPAdapter
+from urllib3.util.retry import Retry
 from .__version__ import __version__
 from datamaxi.error import ClientError, ServerError
 from datamaxi.lib.utils import cleanNoneValue
@@ -24,6 +26,9 @@ def __init__(
         proxies=None,
         show_limit_usage=False,
         show_header=False,
+        max_retries=3,
+        retry_backoff=0.5,
+        retry_statuses=(502, 503, 504),
     ):
         """Client API constructor. `api_key` can be set
         as an environment variable `DATAMAXI_API_KEY`.
@@ -35,6 +40,12 @@ def __init__(
             proxies (dict): The proxies for the requests.
             show_limit_usage (bool): Show the limit usage.
             show_header (bool): Show the header.
+            max_retries (int): Retry attempts for transient gateway 5xx
+                and connection/read errors. Set to 0 to disable.
+            retry_backoff (float): Backoff factor between retries (seconds);
+                see urllib3 ``Retry(backoff_factor=...)``.
+            retry_statuses (tuple): HTTP status codes treated as transient
+                and retried (GET only).
         """
         self.api_key = api_key or os.environ.get("DATAMAXI_API_KEY")
         self.base_url = base_url
@@ -51,10 +62,37 @@ def __init__(
                 "X-DTMX-APIKEY": str(self.api_key),
             }
         )
+        self._mount_retries(max_retries, retry_backoff, retry_statuses)
 
         self._logger = logging.getLogger(__name__)
         return
 
+    def _mount_retries(self, max_retries, retry_backoff, retry_statuses):
+        """Mount a urllib3 retry policy on the session's HTTP adapters.
+
+        Retries transient gateway 5xx (``retry_statuses``) and
+        connection/read failures on idempotent GETs, honoring any
+        ``Retry-After`` header. ``raise_on_status`` is False so an
+        exhausted retry returns the final response and the SDK's own
+        ``_handle_exception`` still raises ``ServerError`` — preserving
+        the existing error contract instead of leaking urllib3's
+        ``MaxRetryError``.
+        """
+        retry = Retry(
+            total=max_retries,
+            connect=max_retries,
+            read=max_retries,
+            status=max_retries,
+            status_forcelist=tuple(retry_statuses),
+            allowed_methods=frozenset(["GET"]),
+            backoff_factor=retry_backoff,
+            respect_retry_after_header=True,
+            raise_on_status=False,
+        )
+        adapter = HTTPAdapter(max_retries=retry)
+        self.session.mount("https://", adapter)
+        self.session.mount("http://", adapter)
+
     def query(self, url_path, payload=None):
         return self.send_request("GET", url_path, payload=payload)
 

tests/test_retry.py

@@ -0,0 +1,60 @@
+"""Local tests for the transient-5xx retry policy mounted on the session.
+
+`responses` intercepts at the adapter-send level, which sits *above* urllib3's
+Retry, so retries can't be exercised through it. These tests assert the retry
+policy is wired onto the session adapters with the expected configuration.
+"""
+
+from datamaxi.api import API
+from datamaxi import Datamaxi
+
+BASE_URL = "https://api.datamaxiplus.com"
+
+
+def _retry(client):
+    return client.session.get_adapter(BASE_URL).max_retries
+
+
+def test_default_retry_policy_mounted():
+    r = _retry(API(api_key="k", base_url=BASE_URL))
+    assert r.total == 3
+    assert tuple(r.status_forcelist) == (502, 503, 504)
+    assert r.backoff_factor == 0.5
+    assert r.respect_retry_after_header is True
+    assert r.raise_on_status is False
+    assert "GET" in r.allowed_methods
+    assert "POST" not in r.allowed_methods
+
+
+def test_retries_disabled_when_zero():
+    r = _retry(API(api_key="k", base_url=BASE_URL, max_retries=0))
+    assert r.total == 0
+
+
+def test_retry_params_are_tunable():
+    r = _retry(
+        API(
+            api_key="k",
+            base_url=BASE_URL,
+            max_retries=5,
+            retry_backoff=1.5,
+            retry_statuses=(500, 502),
+        )
+    )
+    assert r.total == 5
+    assert r.backoff_factor == 1.5
+    assert tuple(r.status_forcelist) == (500, 502)
+
+
+def test_same_adapter_mounted_for_http_and_https():
+    client = API(api_key="k", base_url=BASE_URL)
+    assert client.session.get_adapter("https://x") is client.session.get_adapter(
+        "http://x"
+    )
+
+
+def test_retry_config_propagates_through_datamaxi():
+    c = Datamaxi(api_key="k", max_retries=7)
+    # The whole tree shares one API/session (see #137), so the policy is shared.
+    assert _retry(c.cex._api).total == 7
+    assert _retry(c.premium._api).total == 7