Extend retry_api_call with stop_on instead of a private variant

retry_api_call and async_retry_api_call take a keyword-only stop_on
predicate; the private _retry_api_call and _async_retry_api_call go.
Existing callers that pass **request built from a single-key literal now
annotate it as dict[str, Any], since mypy would otherwise match its
values against stop_on; the values are unchanged.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Author: laughingman7743

pyathena/aio/common.py

@@ -9,7 +9,7 @@
 
 from botocore.exceptions import BotoCoreError, ClientError
 
-from pyathena.aio.util import _async_retry_api_call, async_retry_api_call
+from pyathena.aio.util import async_retry_api_call
 from pyathena.common import BaseCursor, CursorIterator
 from pyathena.error import DatabaseError, OperationalError, ProgrammingError
 from pyathena.glue import GlueMetadataClient
@@ -88,7 +88,7 @@ async def _execute(  # type: ignore[override]
         return query_id
 
     async def _get_query_execution(self, query_id: str) -> AthenaQueryExecution:  # type: ignore[override]
-        request = {"QueryExecutionId": query_id}
+        request: dict[str, Any] = {"QueryExecutionId": query_id}
         try:
             response = await async_retry_api_call(
                 self._connection.client.get_query_execution,
@@ -128,7 +128,7 @@ async def _poll(self, query_id: str) -> AthenaQueryExecution:  # type: ignore[ov
         return query_execution
 
     async def _cancel(self, query_id: str) -> None:  # type: ignore[override]
-        request = {"QueryExecutionId": query_id}
+        request: dict[str, Any] = {"QueryExecutionId": query_id}
         try:
             await async_retry_api_call(
                 self._connection.client.stop_query_execution,
@@ -283,11 +283,11 @@ async def _list_databases(  # type: ignore[override]
             max_results=max_results,
         )
         try:
-            response = await _async_retry_api_call(
+            response = await async_retry_api_call(
                 self.connection._client.list_databases,
-                self._retry_config,
-                _logger,
-                stop_on,
+                config=self._retry_config,
+                logger=_logger,
+                stop_on=stop_on,
                 **request,
             )
         except Exception as e:
@@ -342,11 +342,11 @@ async def _get_table_metadata(  # type: ignore[override]
             schema_name=schema_name,
         )
         try:
-            response = await _async_retry_api_call(
+            response = await async_retry_api_call(
                 self._connection.client.get_table_metadata,
-                self._retry_config,
-                _logger,
-                stop_on,
+                config=self._retry_config,
+                logger=_logger,
+                stop_on=stop_on,
                 **request,
             )
         except Exception as e:
@@ -397,11 +397,11 @@ async def _list_table_metadata(  # type: ignore[override]
             max_results=max_results,
         )
         try:
-            response = await _async_retry_api_call(
+            response = await async_retry_api_call(
                 self.connection._client.list_table_metadata,
-                self._retry_config,
-                _logger,
-                stop_on,
+                config=self._retry_config,
+                logger=_logger,
+                stop_on=stop_on,
                 **request,
             )
         except Exception as e:

pyathena/aio/util.py

@@ -12,14 +12,15 @@
 from collections.abc import Callable
 from typing import Any
 
-from pyathena.util import RetryConfig, _retry_api_call, retry_api_call
+from pyathena.util import RetryConfig, retry_api_call
 
 
 async def async_retry_api_call(
     func: Callable[..., Any],
     config: RetryConfig,
     logger: logging.Logger | None = None,
     *args: Any,
+    stop_on: Callable[[BaseException], bool] | None = None,
     **kwargs: Any,
 ) -> Any:
     """Execute a function with retry logic in a thread to avoid blocking the event loop.
@@ -32,21 +33,12 @@ async def async_retry_api_call(
         config: RetryConfig instance specifying retry behavior.
         logger: Optional logger for retry attempt logging.
         *args: Positional arguments to pass to ``retry_api_call``.
+        stop_on: Passed to ``retry_api_call``.
         **kwargs: Keyword arguments to pass to the function.
 
     Returns:
         The result of the successful function call.
     """
-    return await asyncio.to_thread(retry_api_call, func, config, logger, *args, **kwargs)
-
-
-async def _async_retry_api_call(
-    func: Callable[..., Any],
-    config: RetryConfig,
-    logger: logging.Logger | None,
-    stop_on: Callable[[BaseException], bool] | None,
-    *args: Any,
-    **kwargs: Any,
-) -> Any:
-    """``async_retry_api_call`` that raises an exception ``stop_on`` accepts at once."""
-    return await asyncio.to_thread(_retry_api_call, func, config, logger, stop_on, *args, **kwargs)
+    return await asyncio.to_thread(
+        retry_api_call, func, config, logger, *args, stop_on=stop_on, **kwargs
+    )

pyathena/common.py

@@ -29,7 +29,6 @@
     RetryConfig,
     _get_error_code,
     _is_throttling_error,
-    _retry_api_call,
     retry_api_call,
 )
 
@@ -424,11 +423,11 @@ def _list_databases(
             max_results=max_results,
         )
         try:
-            response = _retry_api_call(
+            response = retry_api_call(
                 self.connection._client.list_databases,
-                self._retry_config,
-                _logger,
-                stop_on,
+                config=self._retry_config,
+                logger=_logger,
+                stop_on=stop_on,
                 **request,
             )
         except Exception as e:
@@ -498,11 +497,11 @@ def _get_table_metadata(
             schema_name=schema_name,
         )
         try:
-            response = _retry_api_call(
+            response = retry_api_call(
                 self._connection.client.get_table_metadata,
-                self._retry_config,
-                _logger,
-                stop_on,
+                config=self._retry_config,
+                logger=_logger,
+                stop_on=stop_on,
                 **request,
             )
         except Exception as e:
@@ -553,11 +552,11 @@ def _list_table_metadata(
             max_results=max_results,
         )
         try:
-            response = _retry_api_call(
+            response = retry_api_call(
                 self.connection._client.list_table_metadata,
-                self._retry_config,
-                _logger,
-                stop_on,
+                config=self._retry_config,
+                logger=_logger,
+                stop_on=stop_on,
                 **request,
             )
         except Exception as e:
@@ -610,7 +609,7 @@ def athena_request(
         )
 
     def _get_query_execution(self, query_id: str) -> AthenaQueryExecution:
-        request = {"QueryExecutionId": query_id}
+        request: dict[str, Any] = {"QueryExecutionId": query_id}
         try:
             response = retry_api_call(
                 self._connection.client.get_query_execution,
@@ -625,7 +624,7 @@ def _get_query_execution(self, query_id: str) -> AthenaQueryExecution:
             return AthenaQueryExecution(response)
 
     def _get_calculation_execution_status(self, query_id: str) -> AthenaCalculationExecutionStatus:
-        request = {"CalculationExecutionId": query_id}
+        request: dict[str, Any] = {"CalculationExecutionId": query_id}
         try:
             response = retry_api_call(
                 self._connection.client.get_calculation_execution_status,
@@ -640,7 +639,7 @@ def _get_calculation_execution_status(self, query_id: str) -> AthenaCalculationE
             return AthenaCalculationExecutionStatus(response)
 
     def _get_calculation_execution(self, query_id: str) -> AthenaCalculationExecution:
-        request = {"CalculationExecutionId": query_id}
+        request: dict[str, Any] = {"CalculationExecutionId": query_id}
         try:
             response = retry_api_call(
                 self._connection.client.get_calculation_execution,
@@ -942,7 +941,7 @@ def close(self) -> None:
         raise NotImplementedError  # pragma: no cover
 
     def _cancel(self, query_id: str) -> None:
-        request = {"QueryExecutionId": query_id}
+        request: dict[str, Any] = {"QueryExecutionId": query_id}
         try:
             retry_api_call(
                 self._connection.client.stop_query_execution,

pyathena/filesystem/s3.py

@@ -168,7 +168,9 @@ def __init__(
         self.version_aware = version_aware
 
         requester_pays = kwargs.pop("requester_pays", False)
-        self.request_kwargs = {"RequestPayer": "requester"} if requester_pays else {}
+        self.request_kwargs: dict[str, Any] = (
+            {"RequestPayer": "requester"} if requester_pays else {}
+        )
 
     def _get_client_compatible_with_s3fs(self, **kwargs) -> BaseClient:
         """Build a boto3 S3 client from s3fs-compatible constructor arguments.

pyathena/spark/common.py

@@ -147,7 +147,7 @@ def _wait_for_idle_session(self, session_id: str):
             time.sleep(self._poll_interval)
 
     def _exists_session(self, session_id: str) -> bool:
-        request = {"SessionId": session_id}
+        request: dict[str, Any] = {"SessionId": session_id}
         try:
             retry_api_call(
                 self._connection.client.get_session,
@@ -193,7 +193,7 @@ def _start_session(self) -> str:
             return session_id
 
     def _terminate_session(self) -> None:
-        request = {"SessionId": self._session_id}
+        request: dict[str, Any] = {"SessionId": self._session_id}
         try:
             retry_api_call(
                 self._connection.client.terminate_session,
@@ -231,7 +231,7 @@ def _poll(self, query_id: str) -> AthenaQueryExecution | AthenaCalculationExecut
         return query_execution
 
     def _cancel(self, query_id: str) -> None:
-        request = {"CalculationExecutionId": query_id}
+        request: dict[str, Any] = {"CalculationExecutionId": query_id}
         try:
             retry_api_call(
                 self._connection.client.stop_calculation_execution,

pyathena/util.py

@@ -215,6 +215,7 @@ def retry_api_call(
     config: RetryConfig,
     logger: logging.Logger | None = None,
     *args,
+    stop_on: Callable[[BaseException], bool] | None = None,
     **kwargs,
 ) -> Any:
     """Execute a function with automatic retry logic for AWS API calls.
@@ -228,6 +229,8 @@ def retry_api_call(
         config: RetryConfig instance specifying retry behavior.
         logger: Optional logger for retry attempt logging.
         *args: Positional arguments to pass to the function.
+        stop_on: Optional predicate; an exception it accepts is raised at once
+            instead of being retried.
         **kwargs: Keyword arguments to pass to the function.
 
     Returns:
@@ -251,18 +254,6 @@ def retry_api_call(
         This includes recognized Glue error codes wrapped in MetadataException.
         Other errors are propagated without retrying.
     """
-    return _retry_api_call(func, config, logger, None, *args, **kwargs)
-
-
-def _retry_api_call(
-    func: Callable[..., Any],
-    config: RetryConfig,
-    logger: logging.Logger | None,
-    stop_on: Callable[[BaseException], bool] | None,
-    *args,
-    **kwargs,
-) -> Any:
-    """``retry_api_call`` that raises an exception ``stop_on`` accepts at once."""
     retry = tenacity.Retrying(
         retry=retry_if_exception(
             lambda ex: is_retryable_error(ex, config) and not (stop_on and stop_on(ex))

tests/pyathena/test_util.py

@@ -8,7 +8,6 @@
 from pyathena.util import (
     RetryConfig,
     _is_throttling_error,
-    _retry_api_call,
     _without_retries,
     is_retryable_error,
     parse_output_location,
@@ -338,7 +337,7 @@ def call():
     )
 
     with pytest.raises(ClientError) as caught:
-        _retry_api_call(call, config, None, _is_throttling_error)
+        retry_api_call(call, config, stop_on=_is_throttling_error)
 
     assert caught.value is error
     assert calls == expected_calls