minor: making retry_if_exception more defensive

Author: chollinger93

pyathena/util.py

@@ -3,7 +3,7 @@
 
 import logging
 import re
-from typing import Any, Callable, Iterable, Optional, Pattern, Tuple
+from typing import Any, Callable, Iterable, Optional, Pattern, Tuple, cast
 
 import tenacity
 from tenacity import after_log, retry_if_exception, stop_after_attempt, wait_exponential
@@ -172,12 +172,18 @@ def retry_api_call(
         Only retries on AWS exceptions listed in the RetryConfig.exceptions.
         Does not retry on client errors or non-AWS exceptions.
     """
+
+    def _extract_code(ex: BaseException) -> Optional[str]:
+        resp = cast(Optional[dict[str, Any]], getattr(ex, "response", None))
+        err = cast(Optional[dict[str, Any]], (resp or {}).get("Error"))
+        return cast(Optional[str], (err or {}).get("Code"))
+
+    def _is_retryable(ex: BaseException) -> bool:
+        code = _extract_code(ex)
+        return code is not None and code in config.exceptions
+
     retry = tenacity.Retrying(
-        retry=retry_if_exception(
-            lambda e: getattr(e, "response", {}).get("Error", {}).get("Code") in config.exceptions
-            if e
-            else False
-        ),
+        retry=retry_if_exception(_is_retryable),
         stop=stop_after_attempt(config.attempt),
         wait=wait_exponential(
             multiplier=config.multiplier,

tests/pyathena/test_util.py

@@ -1,8 +1,10 @@
 # -*- coding: utf-8 -*-
+from typing import Any
+
 import pytest
 
 from pyathena import DataError
-from pyathena.util import parse_output_location, strtobool
+from pyathena.util import RetryConfig, parse_output_location, retry_api_call, strtobool
 
 
 def test_parse_output_location():
@@ -25,3 +27,38 @@ def test_strtobool():
 
     for n in no:
         assert not strtobool(n)
+
+
+class _WithCodeError(Exception):
+    def __init__(self, code: int) -> None:
+        super().__init__(f"error:{code}")
+        self.response = {"Error": {"Code": code}}
+
+
+class _NoResponseError(Exception):
+    def __init__(self) -> None:
+        super().__init__("error")
+        self.response = None
+
+
+def _test_retry(ex: Exception) -> None:
+    calls = {"n": 0}
+
+    def fn() -> Any:
+        calls["n"] += 1
+        raise ex
+
+    cfg = RetryConfig(attempt=1, max_delay=1)
+
+    with pytest.raises(type(ex)):
+        retry_api_call(fn, config=cfg)
+
+    assert calls["n"] == 1
+
+
+def test_retry_api_call():
+    _test_retry(_WithCodeError(500))
+
+
+def test_retry_api_call_with_none_error():
+    _test_retry(_NoResponseError())