feat: enhance upstream fetch with retry logic for rate limits and non-retryable errors

Author: Bornunique911

application/cmd/cre_main.py

@@ -69,13 +69,22 @@ def fetch_upstream_json(
             last_error = exc
 
         if attempt < max_attempts:
+            retry_after_header = None
+            sleep_seconds = backoff_seconds * attempt
+            if "response" in locals() and response.status_code == 429:
+                retry_after_header = response.headers.get("Retry-After")
+                try:
+                    if retry_after_header is not None:
+                        sleep_seconds = float(retry_after_header)
+                except (TypeError, ValueError):
+                    sleep_seconds = backoff_seconds * attempt
             logger.warning(
                 "upstream fetch failed for %s on attempt %s/%s, retrying",
                 url,
                 attempt,
                 max_attempts,
             )
-            time.sleep(backoff_seconds * attempt)
+            time.sleep(sleep_seconds)
 
     if last_error:
         raise RuntimeError(f"upstream fetch failed for {url}") from last_error
@@ -635,11 +644,11 @@ def download_graph_from_upstream(cache: str) -> None:
     collection = db_connect(path=cache).with_graph()
 
     def download_cre_from_upstream(creid: str):
+        if creid in imported_cres:
+            return
         data = fetch_upstream_json(f"/id/{creid}")
         credict = data["data"]
         cre = defs.Document.from_dict(credict)
-        if cre.id in imported_cres:
-            return
 
         register_cre(cre, collection)
         imported_cres[cre.id] = ""

application/tests/cre_main_test.py

@@ -488,6 +488,56 @@ def test_fetch_upstream_json_retries_transient_failures(
         self.assertEqual(mock_get.call_count, 2)
         mock_sleep.assert_called_once()
 
+    @patch("application.cmd.cre_main.time.sleep")
+    @patch("application.cmd.cre_main.requests.get")
+    def test_fetch_upstream_json_fails_fast_on_non_retryable_status(
+        self, mock_get, mock_sleep
+    ) -> None:
+        response = Mock()
+        response.status_code = 404
+        response.headers = {}
+        mock_get.return_value = response
+
+        with self.assertRaises(RuntimeError):
+            main.fetch_upstream_json("/root_cres")
+
+        self.assertEqual(mock_get.call_count, 1)
+        mock_sleep.assert_not_called()
+
+    @patch("application.cmd.cre_main.time.sleep")
+    @patch("application.cmd.cre_main.requests.get")
+    def test_fetch_upstream_json_retries_retryable_status_until_exhausted(
+        self, mock_get, mock_sleep
+    ) -> None:
+        response = Mock()
+        response.status_code = 503
+        response.headers = {}
+        mock_get.return_value = response
+
+        with self.assertRaises(RuntimeError):
+            main.fetch_upstream_json("/root_cres", max_attempts=3, backoff_seconds=2)
+
+        self.assertEqual(mock_get.call_count, 3)
+        self.assertEqual(mock_sleep.call_count, 2)
+        mock_sleep.assert_any_call(2)
+        mock_sleep.assert_any_call(4)
+
+    @patch("application.cmd.cre_main.time.sleep")
+    @patch("application.cmd.cre_main.requests.get")
+    def test_fetch_upstream_json_honors_retry_after_for_rate_limit(
+        self, mock_get, mock_sleep
+    ) -> None:
+        response = Mock()
+        response.status_code = 429
+        response.headers = {"Retry-After": "7"}
+        mock_get.side_effect = [response, response, RuntimeError("should not reach")]
+
+        with self.assertRaises(RuntimeError):
+            main.fetch_upstream_json("/root_cres", max_attempts=2, backoff_seconds=2)
+
+        self.assertEqual(mock_get.call_count, 2)
+        mock_sleep.assert_called_once_with(7.0)
+
     @patch.object(main, "db_connect")
     @patch.object(Queue, "enqueue_call")
     @patch.object(redis, "wait_for_jobs")