feat(utils): retry transient upstream errors in rest_query/post_query

gget/utils.py

@@ -742,6 +742,44 @@ def wrap_cols_func(df: pd.DataFrame, cols: list[str]) -> Any:
     return display(HTML(df.to_html().replace("\\n", "<br>")))
 
 
+# Retry the shared REST query helpers on transient upstream errors. Ensembl's REST API (which
+# many gget modules depend on) returns 429 (rate limit, with a Retry-After header) or 5xx under
+# load; a short backoff keeps a transient hiccup from failing the whole call.
+_QUERY_RETRIES = 4
+_QUERY_BACKOFF_SEC = 2
+_RETRYABLE_STATUS = frozenset({429, 500, 502, 503, 504})
+
+
+def _query_with_retry(method: str, url: str, **kwargs: Any) -> requests.Response:
+    """GET/POST with retry on transient errors (429/5xx and connection/timeout errors).
+
+    Honors a Retry-After header when present. Non-retryable client errors (e.g. 404) and the
+    final response after exhausting retries are returned to the caller to handle as before.
+    """
+    kwargs.setdefault("timeout", DEFAULT_REQUESTS_TIMEOUT)
+    last_exc: Exception | None = None
+    response: requests.Response | None = None
+    for attempt in range(1, _QUERY_RETRIES + 1):
+        try:
+            response = requests.request(method, url, **kwargs)
+        except requests.RequestException as e:
+            last_exc = e
+            response = None
+        else:
+            if response.ok or response.status_code not in _RETRYABLE_STATUS:
+                return response
+        if attempt < _QUERY_RETRIES:
+            wait = _QUERY_BACKOFF_SEC * attempt
+            if response is not None:
+                retry_after = response.headers.get("Retry-After", "")
+                if retry_after.isdigit():
+                    wait = min(int(retry_after), 30)
+            time.sleep(wait)
+    if response is not None:
+        return response
+    raise requests.RequestException(f"Request to {url} failed after {_QUERY_RETRIES} attempts: {last_exc}")
+
+
 def rest_query(server: str, query: str, content_type: str) -> Any:
     """Function to perform a REST API query.
 
@@ -752,7 +790,7 @@ def rest_query(server: str, query: str, content_type: str) -> Any:
 
     Returns server output.
     """
-    r = requests.get(server + query, headers={"Content-Type": content_type})
+    r = _query_with_retry("GET", server + query, headers={"Content-Type": content_type})
 
     if not r.ok:
         raise RuntimeError(
@@ -774,7 +812,7 @@ def post_query(server: str, endpoint: str, query: Any) -> Any:
 
     :return: server output
     """
-    r = requests.post(server + endpoint, json=query, headers={"Content-Type": "application/json"})
+    r = _query_with_retry("POST", server + endpoint, json=query, headers={"Content-Type": "application/json"})
 
     if not r.ok:
         raise RuntimeError(

pyproject.toml

@@ -145,7 +145,7 @@ show_error_codes = true
 
 [tool.pytest]
 ini_options.testpaths = [ "tests" ]
-ini_options.addopts = [ "-ra" ]
+ini_options.addopts = [ "-ra", "--durations=25" ]
 
 [tool.coverage]
 run.omit = [

tests/test_utils.py

@@ -1,5 +1,8 @@
 import unittest
+from unittest.mock import Mock, patch
 
+import gget.utils as gget_utils
+import requests
 from gget.constants import ENSEMBL_FTP_URL_NV, ENSEMBL_REST_API, UNIPROT_REST_API
 from gget.utils import (
     aa_colors,
@@ -182,3 +185,64 @@ def test_ref_iv_species_options(self):
     def test_ref_species_options_bad_type(self):
         with self.assertRaises(RuntimeError):
             ref_species_options("gtf", release=2000)
+
+
+def _resp(status_code, headers=None):
+    """Minimal mock requests.Response with .status_code/.ok/.headers."""
+    m = Mock()
+    m.status_code = status_code
+    m.ok = 200 <= status_code < 300
+    m.headers = headers or {}
+    return m
+
+
+class TestQueryRetry(unittest.TestCase):
+    """Network-free tests for the transient-error retry in the REST query helpers."""
+
+    @patch("gget.utils.time.sleep")
+    @patch("gget.utils.requests.request")
+    def test_retries_on_503_then_succeeds(self, mock_request, _sleep):
+        mock_request.side_effect = [_resp(503), _resp(503), _resp(200)]
+        r = gget_utils._query_with_retry("GET", "http://x")
+        self.assertEqual(r.status_code, 200)
+        self.assertEqual(mock_request.call_count, 3)
+
+    @patch("gget.utils.time.sleep")
+    @patch("gget.utils.requests.request")
+    def test_retries_on_connection_error_then_succeeds(self, mock_request, _sleep):
+        mock_request.side_effect = [requests.ConnectionError("boom"), _resp(200)]
+        r = gget_utils._query_with_retry("GET", "http://x")
+        self.assertEqual(r.status_code, 200)
+        self.assertEqual(mock_request.call_count, 2)
+
+    @patch("gget.utils.time.sleep")
+    @patch("gget.utils.requests.request")
+    def test_no_retry_on_client_error(self, mock_request, _sleep):
+        # A 404 is a real client error, not transient -> returned immediately (caller handles it).
+        mock_request.return_value = _resp(404)
+        r = gget_utils._query_with_retry("GET", "http://x")
+        self.assertEqual(r.status_code, 404)
+        self.assertEqual(mock_request.call_count, 1)
+
+    @patch("gget.utils.time.sleep")
+    @patch("gget.utils.requests.request")
+    def test_exhausts_retries_on_persistent_503(self, mock_request, _sleep):
+        mock_request.return_value = _resp(503)
+        r = gget_utils._query_with_retry("GET", "http://x")
+        # After exhausting retries the last (503) response is returned so the caller raises as before.
+        self.assertEqual(r.status_code, 503)
+        self.assertEqual(mock_request.call_count, gget_utils._QUERY_RETRIES)
+
+    @patch("gget.utils.time.sleep")
+    @patch("gget.utils.requests.request")
+    def test_persistent_connection_error_raises(self, mock_request, _sleep):
+        mock_request.side_effect = requests.ConnectionError("down")
+        with self.assertRaises(requests.RequestException):
+            gget_utils._query_with_retry("GET", "http://x")
+
+    @patch("gget.utils.time.sleep")
+    @patch("gget.utils.requests.request")
+    def test_retry_after_header_used(self, mock_request, mock_sleep):
+        mock_request.side_effect = [_resp(429, headers={"Retry-After": "7"}), _resp(200)]
+        gget_utils._query_with_retry("GET", "http://x")
+        mock_sleep.assert_called_once_with(7)