Stop the retries of closed result sets and cancelled asyncio calls

docs/aio.md

@@ -166,6 +166,13 @@ async with await aio_connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
             print(f"Query timed out: {cursor.query_id}")
 ```
 
+Cancelling a task also stops the retries of the Athena API requests that the asyncio cursors make in a thread for it, such as the status polling of `execute()` and the page requests of `AioCursor`.
+No further attempt is made, and the wait before the next attempt ends.
+A request that has already been sent still runs to its end in its thread.
+The requests that stop a query or a calculation and the request that terminates a Spark session keep their retries.
+The start request of `execute()` with `kill_on_interrupt` enabled keeps its retries after a cancellation, and stops them when the task is cancelled again.
+The fetch methods of `AioPandasCursor`, `AioArrowCursor`, `AioPolarsCursor`, and `AioS3FSCursor` run to their end in their thread.
+
 (aio-dict-cursor)=
 
 ## AioDictCursor

docs/api/utilities.rst

@@ -24,6 +24,10 @@ Utility Functions
 
 .. autofunction:: pyathena.util.retry_api_call
 
+.. autofunction:: pyathena.util.stop_retries_on
+
+.. autofunction:: pyathena.util.unstoppable_retries
+
 .. autofunction:: pyathena.util.is_retryable_error
 
 .. autofunction:: pyathena.util.parse_output_location

pyathena/aio/common.py

@@ -267,6 +267,7 @@ async def _cancel(self, query_id: str) -> None:  # type: ignore[override]
                 self._connection.client.stop_query_execution,
                 config=self._retry_config,
                 logger=_logger,
+                stoppable=False,
                 **request,
             )
         except Exception as e:

pyathena/aio/result_set.py

@@ -123,7 +123,8 @@ async def _async_get_query_results(
 
         Raises:
             ProgrammingError: If the query ID is missing, the query has not
-                succeeded, or the result set is closed.
+                succeeded, or the result set is closed, also during the
+                request, which then makes no further attempt.
             OperationalError: If the request fails.
         """
         if self.is_closed:
@@ -134,9 +135,13 @@ async def _async_get_query_results(
                 self.connection.client.get_query_results,
                 config=self._retry_config,
                 logger=_logger,
+                stop_when=lambda: self.is_closed,
                 **request,
             )
         except Exception as e:
+            if self.is_closed:
+                # close() during the request stopped its retries.
+                raise ProgrammingError("AthenaAioResultSet is closed.") from e
             _logger.exception("Failed to fetch result set.")
             raise OperationalError(*e.args) from e
         else:
@@ -145,19 +150,26 @@ async def _async_get_query_results(
     async def _async_fetch(self) -> None:
         """Fetch the next page of rows, if any, into the buffer.
 
-        Nothing is fetched when no page remains, as after ``close()``, and the page
-        is discarded when ``close()`` runs during its request.
+        Nothing is fetched when no page remains or the result set is closed. A
+        ``close()`` during the request stops its retries, and the page is
+        discarded.
 
         Raises:
             ProgrammingError: If the next page is requested and the query ID is
-                missing, the query has not succeeded, or the result set is closed.
+                missing or the query has not succeeded.
             OperationalError: If the request fails.
         """
         if not self._next_token:
             return
-        response = await self._async_get_query_results(self._arraysize, self._next_token)
-        rows, self._next_token = self._parse_result_rows(response)
-        self._process_rows(rows)
+        try:
+            response = await self._async_get_query_results(self._arraysize, self._next_token)
+        except ProgrammingError:
+            # Raised for a close() before the request or during its retries.
+            if not self.is_closed:
+                raise
+        else:
+            rows, self._next_token = self._parse_result_rows(response)
+            self._process_rows(rows)
         # close() clears the connection before the rest, so checking after the
         # assignments catches a close() during the request or this method.
         if self.is_closed:

pyathena/aio/spark/cursor.py

@@ -278,6 +278,7 @@ async def _cancel(self, query_id: str) -> None:  # type: ignore[override]
                 self._connection.client.stop_calculation_execution,
                 config=self._retry_config,
                 logger=_logger,
+                stoppable=False,
                 **request,
             )
         except Exception as e:
@@ -292,6 +293,7 @@ async def _terminate_session(self) -> None:  # type: ignore[override]
                 self._connection.client.terminate_session,
                 config=self._retry_config,
                 logger=_logger,
+                stoppable=False,
                 **request,
             )
         except Exception as e:

pyathena/aio/util.py

@@ -10,13 +10,16 @@
 from __future__ import annotations
 
 import asyncio
+import contextlib
 import contextvars
 import logging
 import threading
 from collections.abc import Callable
 from typing import Any, Protocol, TypeVar, cast
 
-from pyathena.util import RetryConfig, retry_api_call
+from pyathena.util import RetryConfig, retry_api_call, stop_retries_on, unstoppable_retries
+
+T = TypeVar("T")
 
 
 class _Closable(Protocol):
@@ -26,33 +29,82 @@ def close(self) -> Any: ...
 _ClosableT = TypeVar("_ClosableT", bound=_Closable)
 
 
+async def run_in_thread(func: Callable[..., T], /, *args: Any, **kwargs: Any) -> T:
+    """Run a blocking function in the default executor, stopping its retries on cancellation.
+
+    Like ``asyncio.to_thread()``, the function runs in the event loop's default
+    executor with a copy of the caller's context, and a cancellation of the
+    caller skips a function that has not started. A cancellation also stops
+    the retries of the AWS API calls that the running function makes (see
+    :func:`~pyathena.util.stop_retries_on`). A request that has been sent
+    still runs to its end, and the function returns or raises in its
+    thread, where nobody awaits it.
+
+    Args:
+        func: The blocking function.
+        *args: Positional arguments for ``func``.
+        **kwargs: Keyword arguments for ``func``.
+
+    Returns:
+        The return value of ``func``.
+    """
+    stop = threading.Event()
+
+    def call() -> T:
+        """Run the function in a scope that the cancellation stops."""
+        with stop_retries_on(stop):
+            return func(*args, **kwargs)
+
+    future = asyncio.get_running_loop().run_in_executor(None, contextvars.copy_context().run, call)
+    try:
+        return await future
+    except asyncio.CancelledError:
+        stop.set()
+        raise
+
+
 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,
+    stop_when: Callable[[], bool] | None = None,
+    stoppable: bool = True,
     **kwargs: Any,
 ) -> Any:
     """Execute a function with retry logic in a thread to avoid blocking the event loop.
 
-    Wraps ``retry_api_call`` with ``asyncio.to_thread()`` so that blocking
-    boto3 calls do not block the asyncio event loop.
+    Runs ``retry_api_call`` with :func:`run_in_thread` so that blocking boto3
+    calls do not block the asyncio event loop. Cancelling the caller makes no
+    further attempt and ends the backoff wait, unless ``stoppable`` is false.
 
     Args:
         func: The AWS API function to 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``.
+        stop_when: Passed to ``retry_api_call``.
+        stoppable: Whether a cancellation of the caller stops the retries. A
+            request that cleans up, such as one that stops a query or
+            terminates a session, passes False: its retries then run to their
+            end in the thread (see :func:`~pyathena.util.unstoppable_retries`),
+            as with ``asyncio.to_thread()``.
         **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, stop_on=stop_on, **kwargs
-    )
+
+    def call() -> Any:
+        """Make the call, ignoring the stop scopes unless it is stoppable."""
+        with contextlib.nullcontext() if stoppable else unstoppable_retries():
+            return retry_api_call(
+                func, config, logger, *args, stop_on=stop_on, stop_when=stop_when, **kwargs
+            )
+
+    return await (run_in_thread(call) if stoppable else asyncio.to_thread(call))
 
 
 async def close_in_thread(close: Callable[[], Any]) -> None:
@@ -76,7 +128,8 @@ async def create_in_thread(
 
     Like ``asyncio.to_thread()``, the constructor runs in the event loop's
     default executor with the caller's context. If the caller is cancelled, a
-    construction that has not started is skipped, an object constructed after
+    construction that has not started is skipped, the retries of a running
+    one are stopped as by :func:`run_in_thread`, an object constructed after
     the cancellation is closed by the constructing thread, and an object
     constructed just before it is closed in a new thread: nothing else can
     close it, and a close by the garbage collector could block the event
@@ -96,14 +149,22 @@ async def create_in_thread(
     constructed: list[_ClosableT] = []
 
     def construct() -> _ClosableT | None:
+        """Construct the object in a scope that the cancellation stops.
+
+        Returns:
+            The object, or None if the caller was cancelled, after which an
+            object constructed here is closed here.
+        """
         if cancelled.is_set():
             return None
-        obj = factory(*args, **kwargs)
+        with stop_retries_on(cancelled):
+            obj = factory(*args, **kwargs)
         with lock:
             if not cancelled.is_set():
                 constructed.append(obj)
                 return obj
-        obj.close()
+        with unstoppable_retries():
+            obj.close()
         return None
 
     future = asyncio.get_running_loop().run_in_executor(

pyathena/common.py

@@ -35,6 +35,7 @@
     _get_error_code,
     _is_throttling_error,
     retry_api_call,
+    unstoppable_retries,
 )
 
 if TYPE_CHECKING:
@@ -1407,6 +1408,9 @@ def close(self) -> None:
     def _cancel(self, query_id: str) -> None:
         """Stop a query execution with ``StopQueryExecution``.
 
+        The request cleans up, so a stop scope does not stop its retries (see
+        :func:`~pyathena.util.unstoppable_retries`).
+
         Args:
             query_id: The query execution ID.
 
@@ -1415,12 +1419,13 @@ def _cancel(self, query_id: str) -> None:
         """
         request: dict[str, Any] = {"QueryExecutionId": query_id}
         try:
-            retry_api_call(
-                self._connection.client.stop_query_execution,
-                config=self._retry_config,
-                logger=_logger,
-                **request,
-            )
+            with unstoppable_retries():
+                retry_api_call(
+                    self._connection.client.stop_query_execution,
+                    config=self._retry_config,
+                    logger=_logger,
+                    **request,
+                )
         except Exception as e:
             _logger.exception("Failed to cancel query.")
             raise OperationalError(*e.args) from e

pyathena/pandas/result_set.py

@@ -10,6 +10,7 @@
 from collections.abc import Callable, Iterable, Iterator
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import ExitStack, contextmanager
+from contextvars import copy_context
 from functools import partial
 from io import BufferedReader, BytesIO, IOBase, StringIO, TextIOWrapper
 from itertools import chain, islice
@@ -1625,11 +1626,20 @@ def _read_parquet_footers(self) -> list[FileMetaData]:
             self.connection.s3_client.meta.config.max_pool_connections,
             len(self._data_manifest),
         )
+        # Each read runs with a copy of this context, so that a stop of the retries in
+        # this thread (see stop_retries_on()) stops those of the reads too.
+        contexts = [copy_context() for _ in self._data_manifest]
         try:
             # When the failure of a read reaches map()'s iterator, it cancels the reads
             # that have not started.
             with ThreadPoolExecutor(max_workers=workers) as executor:
-                footers = list(executor.map(self._read_parquet_footer, self._data_manifest))
+                footers = list(
+                    executor.map(
+                        lambda context, path: context.run(self._read_parquet_footer, path),
+                        contexts,
+                        self._data_manifest,
+                    )
+                )
         except Exception as e:
             _logger.exception(f"Failed to read {self._unload_location}.")
             raise OperationalError(*e.args) from e