Retry API calls on transient connection errors

n2y/utils.py

@@ -7,6 +7,7 @@
 from time import sleep
 
 import pandoc
+import requests
 import yaml
 from pandoc.types import Meta, MetaBool, MetaList, MetaMap, MetaString, Space, Str
 from plumbum import ProcessExecutionError
@@ -325,6 +326,22 @@ def wrapper(*args, retry_count=0, **kwargs):
         assert "retry_count" not in kwargs, "retry_count is a reserved keyword"
         try:
             return api_call(*args, **kwargs)
+        except requests.exceptions.ConnectionError as err:
+            # Transient network failures (e.g. connection reset by peer during
+            # long pulls with many file downloads) are worth retrying too.
+            if retry_count >= max_api_retries:
+                raise err
+            retry_count += 1
+            retry_after = 2 * retry_count
+            client.logger.info(
+                "This API call failed with a connection error and "
+                "will be retried in %f seconds. Attempt %d of %d.",
+                retry_after,
+                retry_count,
+                max_api_retries,
+            )
+            sleep(retry_after)
+            return wrapper(*args, retry_count=retry_count, **kwargs)
         except APIResponseError as err:
             if err.code not in APIErrorCode.RetryableCodes:
                 raise err

tests/test_utils.py

@@ -1,7 +1,9 @@
 from datetime import datetime, timedelta, timezone
 from math import isclose
+from unittest.mock import patch
 
 import pytest
+import requests
 from pandoc.types import MetaBool, MetaList, MetaMap, MetaString
 from pytest import raises
 
@@ -124,6 +126,34 @@ def tester(_):
     assert call_count == 2
 
 
+def test_retry_api_call_connection_error():
+    client = Client(foo_token)
+    call_count = 0
+
+    @retry_api_call
+    def tester(_):
+        nonlocal call_count
+        call_count += 1
+        if call_count == 1:
+            raise requests.exceptions.ConnectionError("Connection reset by peer")
+        return True
+
+    with patch("n2y.utils.sleep"):
+        assert tester(client)
+    assert call_count == 2
+
+
+def test_retry_api_call_connection_error_max_retries():
+    client = Client(foo_token)
+
+    @retry_api_call
+    def tester(_):
+        raise requests.exceptions.ConnectionError("Connection reset by peer")
+
+    with patch("n2y.utils.sleep"), raises(requests.exceptions.ConnectionError):
+        tester(client)
+
+
 def test_retry_api_call_max_errors():
     client = Client(foo_token)