From e5e741631b736c611c1d0d295cdd8447a3f1bd85 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 01:18:04 +0900 Subject: [PATCH 1/9] Let retry_api_call() stop its retries on a stop scope or condition A stop scope (stop_retries_on()) carries a threading.Event in a context variable, which threads that run a copy of the context inherit, and stop_when is an explicit condition. Once either holds, retry_api_call() makes no further attempt, ends its backoff wait, and raises the error of the last attempt. shield_retries() keeps the scopes around a cleanup from stopping it. Without a scope or condition, tenacity's own sleep is kept. Refs #1202 Co-Authored-By: Claude Opus 5.5 --- pyathena/util.py | 124 ++++++++++++++++++++++++++++++- tests/pyathena/test_util.py | 143 ++++++++++++++++++++++++++++++++++++ 2 files changed, 263 insertions(+), 4 deletions(-) diff --git a/pyathena/util.py b/pyathena/util.py index a88dc6bf5..6a2c50161 100644 --- a/pyathena/util.py +++ b/pyathena/util.py @@ -4,12 +4,17 @@ import logging import re -from collections.abc import Callable, Iterable +import threading +import time +from collections.abc import Callable, Iterable, Iterator +from contextlib import contextmanager +from contextvars import ContextVar from re import Pattern from typing import TYPE_CHECKING, Any import tenacity from tenacity import ( + RetryCallState, after_log, retry_if_exception, stop_after_attempt, @@ -265,12 +270,65 @@ def is_retryable_error(ex: BaseException, config: RetryConfig) -> bool: ) +# The events of the enclosing stop scopes of retry_api_call(); see stop_retries_on(). +_RETRY_STOP_EVENTS: ContextVar[tuple[threading.Event, ...]] = ContextVar( + "pyathena_retry_stop_events", default=() +) +# The longest that a backoff wait sleeps before it checks again whether the +# retries are stopped, unless it can wait for the single event of a scope. +_STOP_POLL_INTERVAL = 0.1 + + +@contextmanager +def stop_retries_on(event: threading.Event) -> Iterator[None]: + """Stop the retries of AWS API calls in this context once an event is set. + + A call of :func:`retry_api_call` in the block makes no further attempt + once the event is set, ends its backoff wait, and raises the error of its + last attempt. That includes the calls in threads that run a copy of this + context, as ``asyncio.to_thread()`` runs its function. The scopes nest: a + call stops when the event of any enclosing scope is set. The stop is + cooperative: a request that has been sent, and the retries of botocore + itself, still run to their end. + + Args: + event: The event that stops the retries. + + Yields: + None. + """ + token = _RETRY_STOP_EVENTS.set((*_RETRY_STOP_EVENTS.get(), event)) + try: + yield + finally: + _RETRY_STOP_EVENTS.reset(token) + + +@contextmanager +def shield_retries() -> Iterator[None]: + """Keep the events of the enclosing stop scopes from stopping retries in this block. + + For cleanup that has to run to its end even when the work that it cleans + up was stopped, such as the abort of a multipart upload. A scope that + :func:`stop_retries_on` opens in the block stops the retries again. + + Yields: + None. + """ + token = _RETRY_STOP_EVENTS.set(()) + try: + yield + finally: + _RETRY_STOP_EVENTS.reset(token) + + def retry_api_call( func: Callable[..., Any], config: RetryConfig, logger: logging.Logger | None = None, *args, stop_on: Callable[[BaseException], bool] | None = None, + stop_when: Callable[[], bool] | None = None, **kwargs, ) -> Any: """Execute a function with automatic retry logic for AWS API calls. @@ -286,13 +344,18 @@ def retry_api_call( *args: Positional arguments to pass to the function. stop_on: Optional predicate; an exception it accepts is raised at once instead of being retried. + stop_when: Optional condition, checked from the thread of the call, + that stops the retries once it holds, as the event of + :func:`stop_retries_on` does. The backoff wait checks it at least + every 0.1 seconds. **kwargs: Keyword arguments to pass to the function. Returns: The result of the successful function call. Raises: - The original exception if all retry attempts are exhausted. + The original exception if all retry attempts are exhausted, or the + error of the last attempt if the retries are stopped. Example: >>> from pyathena.util import RetryConfig, retry_api_call @@ -309,9 +372,56 @@ def retry_api_call( This includes recognized Glue error codes wrapped in MetadataException. Other errors are propagated without retrying. """ + events = _RETRY_STOP_EVENTS.get() + + def stopped() -> bool: + """Whether an event of the enclosing scopes is set or ``stop_when`` holds.""" + return any(event.is_set() for event in events) or (stop_when is not None and stop_when()) + + # Without a stop, tenacity's own sleep is kept. + stop_kwargs: dict[str, Any] = {} + # The error of the attempt before a backoff wait, raised if the wait is stopped. + failures: list[BaseException] = [] + if events or stop_when is not None: + + def record_failure(retry_state: RetryCallState) -> None: + """Keep the error of the failed attempt, which a stopped wait raises. + + Args: + retry_state: The state of the call before the backoff wait. + """ + outcome = retry_state.outcome + if outcome is not None and (error := outcome.exception()) is not None: + failures[:] = [error] + + def sleep(seconds: float) -> None: + """Wait before the next attempt, ending the wait once the retries are stopped. + + Args: + seconds: The backoff time. + + Raises: + BaseException: The error of the last attempt, instead of + making the next one, if the retries are stopped. + """ + deadline = time.monotonic() + seconds + while not stopped(): + remaining = deadline - time.monotonic() + if remaining <= 0: + return + if len(events) == 1 and stop_when is None: + events[0].wait(remaining) + else: + time.sleep(min(remaining, _STOP_POLL_INTERVAL)) + raise failures[0] + + stop_kwargs = {"before_sleep": record_failure, "sleep": sleep} + retry = tenacity.Retrying( retry=retry_if_exception( - lambda ex: is_retryable_error(ex, config) and not (stop_on and stop_on(ex)) + lambda ex: ( + is_retryable_error(ex, config) and not (stop_on and stop_on(ex)) and not stopped() + ) ), stop=stop_after_attempt(config.attempt), # Uniform jitter of up to one multiplier keeps concurrent clients from @@ -324,5 +434,11 @@ def retry_api_call( + wait_random(0, config.multiplier), after=after_log(logger, logger.getEffectiveLevel()) if logger else None, # type: ignore[arg-type] reraise=True, + **stop_kwargs, ) - return retry(func, *args, **kwargs) + try: + return retry(func, *args, **kwargs) + finally: + # The traceback of a raised error keeps this frame, which would keep + # the error in a reference cycle through the list. + failures.clear() diff --git a/tests/pyathena/test_util.py b/tests/pyathena/test_util.py index 1c5eed6fa..83dc2ef36 100644 --- a/tests/pyathena/test_util.py +++ b/tests/pyathena/test_util.py @@ -1,3 +1,5 @@ +import threading +import time from typing import Any import pytest @@ -12,8 +14,11 @@ is_retryable_error, parse_output_location, retry_api_call, + shield_retries, + stop_retries_on, strtobool, ) +from tests.pyathena.util import EVENT_TIMEOUT def test_parse_output_location(): @@ -372,3 +377,141 @@ def call(): assert caught.value is error assert calls == expected_calls + + +class _ThrottledCall: + """A call that always throttles and stops its retries at a given point.""" + + def __init__(self, stop, stop_in): + """Initialize the call. + + Args: + stop: Stops the retries; called once. + stop_in: "request" to stop during the first attempt, "backoff" to + stop from another thread during the backoff wait after it. + """ + self.stop = stop + self.stop_in = stop_in + self.error = _throttling_error() + self.calls = 0 + self.timer = None + + def __call__(self): + if self.calls == 0: + if self.stop_in == "request": + self.stop() + else: + self.timer = threading.Timer(0.1, self.stop) + self.timer.start() + self.calls += 1 + raise self.error + + def join(self): + """Wait for the thread that stops the backoff wait, if any.""" + if self.timer is not None: + self.timer.join() + + +@pytest.mark.parametrize("stop_in", ["request", "backoff"]) +def test_retry_api_call_stopped_by_event(stop_in): + """GH-1202: the event of a stop scope ends the retries and their backoff wait. + + Without the stop, the backoff wait would take at least EVENT_TIMEOUT seconds. + """ + event = threading.Event() + call = _ThrottledCall(event.set, stop_in) + started = time.monotonic() + with stop_retries_on(event), pytest.raises(ClientError) as caught: + retry_api_call(call, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) + call.join() + assert time.monotonic() - started < EVENT_TIMEOUT + assert caught.value is call.error + assert call.calls == 1 + + +@pytest.mark.parametrize("stop_in", ["request", "backoff"]) +def test_retry_api_call_stopped_by_stop_when(stop_in): + """GH-1202: ``stop_when`` ends the retries and their backoff wait.""" + stopped = threading.Event() + call = _ThrottledCall(stopped.set, stop_in) + started = time.monotonic() + with pytest.raises(ClientError) as caught: + retry_api_call( + call, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT), stop_when=stopped.is_set + ) + call.join() + assert time.monotonic() - started < EVENT_TIMEOUT + assert caught.value is call.error + assert call.calls == 1 + + +def test_retry_api_call_stopped_before_call(): + """A call in a stopped scope makes its first attempt, but no retry.""" + event = threading.Event() + event.set() + call = _ThrottledCall(event.set, "request") + with stop_retries_on(event), pytest.raises(ClientError): + retry_api_call(call, RetryConfig(multiplier=0)) + assert call.calls == 1 + + +def test_retry_api_call_stopped_by_enclosing_scope(): + """The event of any enclosing scope stops the retries.""" + outer = threading.Event() + call = _ThrottledCall(outer.set, "backoff") + with ( + stop_retries_on(outer), + stop_retries_on(threading.Event()), + pytest.raises(ClientError), + ): + retry_api_call(call, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) + call.join() + assert call.calls == 1 + + +def test_retry_api_call_shielded(): + """Retries in a shielded block ignore the events of the enclosing scopes.""" + event = threading.Event() + event.set() + error = _throttling_error() + calls = 0 + + def call(): + nonlocal calls + calls += 1 + if calls < 3: + raise error + return "success" + + with stop_retries_on(event), shield_retries(): + assert retry_api_call(call, RetryConfig(multiplier=0)) == "success" + assert calls == 3 + # A scope opened in the shielded block stops the retries again. + calls = 0 + with ( + stop_retries_on(event), + shield_retries(), + stop_retries_on(event), + pytest.raises(ClientError), + ): + retry_api_call(call, RetryConfig(multiplier=0)) + assert calls == 1 + + +def test_retry_api_call_scope_ends_with_block(): + """The event of a scope does not stop the retries after its block.""" + event = threading.Event() + with stop_retries_on(event): + event.set() + error = _throttling_error() + calls = 0 + + def call(): + nonlocal calls + calls += 1 + if calls < 3: + raise error + return "success" + + assert retry_api_call(call, RetryConfig(multiplier=0)) == "success" + assert calls == 3 From 1c5debe104e4cc8d0f7245378d3cf6ebd160ae80 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 01:18:05 +0900 Subject: [PATCH 2/9] Stop the page requests of a result set when it is closed The result sets pass stop_when=is_closed to their GetQueryResults requests. A close() during the retries of a page fetch now ends them, and the fetch returns nothing, as after close(); a whole read through the GetQueryResults fallback raises ProgrammingError for the closed result set instead of OperationalError after the remaining attempts. Refs #1202 Co-Authored-By: Claude Opus 5.5 --- pyathena/aio/result_set.py | 26 ++++++--- pyathena/result_set.py | 38 ++++++++++--- tests/pyathena/aio/test_result_set.py | 25 ++++++++- tests/pyathena/test_result_set.py | 45 ++++++++++++++- tests/pyathena/util.py | 80 +++++++++++++++++++++++++++ 5 files changed, 196 insertions(+), 18 deletions(-) diff --git a/pyathena/aio/result_set.py b/pyathena/aio/result_set.py index 40bb2ad3c..f07e87420 100644 --- a/pyathena/aio/result_set.py +++ b/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: diff --git a/pyathena/result_set.py b/pyathena/result_set.py index c0a7cf9e3..817088629 100644 --- a/pyathena/result_set.py +++ b/pyathena/result_set.py @@ -439,7 +439,8 @@ def _get_query_results(self, max_results: int, next_token: str | None = None) -> 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: @@ -450,9 +451,13 @@ def _get_query_results(self, max_results: int, next_token: str | None = None) -> 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("AthenaResultSet is closed.") from e _logger.exception("Failed to fetch result set.") raise OperationalError(*e.args) from e else: @@ -462,19 +467,25 @@ def _fetch(self) -> None: """Fetch the next rows, if any, into the buffer. The default result set fetches the next page; it fetches nothing when no - page remains, as after ``close()``, and discards the page when ``close()`` - runs during its request. + 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 = self._get_query_results(self._arraysize, self._next_token) - rows, self._next_token = self._parse_result_rows(response) - self._process_rows(rows) + try: + response = self._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: @@ -752,6 +763,11 @@ def _iter_all_row_pages(self) -> Iterator[tuple[list[dict[str, Any]], int]]: Yields: The rows of each page and the offset of their first data row, which is 1 when the first page starts with the column labels. + + Raises: + ProgrammingError: If the result set is closed, also during a request, + which then makes no further attempt. + OperationalError: If a request fails. """ _logger.warning( "output_location is not available (e.g. managed query result storage). " @@ -785,7 +801,9 @@ def _fetch_all_rows( List of converted row tuples. Raises: - ProgrammingError: If the metadata is not available. + ProgrammingError: If the metadata is not available or the result set + is closed. + OperationalError: If a request fails. """ metadata = self._metadata if metadata is None: @@ -814,7 +832,9 @@ def _fetch_all_rows_as_csv(self) -> bytes: The CSV text encoded in UTF-8, or empty bytes when the result has no columns. Raises: - ProgrammingError: If the metadata is not available. + ProgrammingError: If the metadata is not available or the result set + is closed. + OperationalError: If a request fails. """ if self._metadata is None: raise ProgrammingError("Metadata is not available.") diff --git a/tests/pyathena/aio/test_result_set.py b/tests/pyathena/aio/test_result_set.py index dc0292ba8..205be0b38 100644 --- a/tests/pyathena/aio/test_result_set.py +++ b/tests/pyathena/aio/test_result_set.py @@ -4,6 +4,7 @@ # See LICENSE or https://opensource.org/licenses/MIT. # # SPDX-License-Identifier: MIT +import time from unittest.mock import MagicMock import pytest @@ -13,7 +14,12 @@ from pyathena.error import OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.util import RetryConfig -from tests.pyathena.util import close_during_next_request, get_query_results_responses +from tests.pyathena.util import ( + EVENT_TIMEOUT, + close_during_next_request, + close_during_next_retries, + get_query_results_responses, +) async def _create_result_set(result_set_class, pages=(("1", "2"),)): @@ -117,6 +123,23 @@ async def test_fetch_after_close(self): assert await result_set.fetchall() == [] assert result_set.rownumber is None + @pytest.mark.parametrize("close_in", ["request", "backoff"]) + @pytest.mark.parametrize( + ("method", "expected"), [("fetchone", None), ("fetchmany", []), ("fetchall", [])] + ) + async def test_fetch_with_close_during_retries(self, close_in, method, expected): + """GH-1202: a close() during the retries of a page request stops them.""" + result_set = await _create_result_set(AthenaAioResultSet, [("1", "2"), ("3",)]) + assert await result_set.fetchmany(2) == [(1,), (2,)] + get_query_results, closer = close_during_next_retries(result_set, close_in) + started = time.monotonic() + assert await getattr(result_set, method)() == expected + assert time.monotonic() - started < EVENT_TIMEOUT + if close_in == "backoff": + closer.join() + assert await result_set.fetchone() is None + assert get_query_results.call_count == 2 + async def test_get_query_results_after_close(self): result_set = await _create_result_set(AthenaAioResultSet, [("1", "2"), ("3",)]) result_set.close() diff --git a/tests/pyathena/test_result_set.py b/tests/pyathena/test_result_set.py index 83df71187..89fdda0f8 100644 --- a/tests/pyathena/test_result_set.py +++ b/tests/pyathena/test_result_set.py @@ -6,6 +6,7 @@ # SPDX-License-Identifier: MIT import collections import sys +import time from unittest.mock import MagicMock, patch import pytest @@ -17,7 +18,12 @@ from pyathena.model import AthenaQueryExecution from pyathena.result_set import AthenaResultSet, WithFetch, WithResultSet from pyathena.util import RetryConfig, override -from tests.pyathena.util import close_during_next_request, get_query_results_responses +from tests.pyathena.util import ( + EVENT_TIMEOUT, + close_during_next_request, + close_during_next_retries, + get_query_results_responses, +) def _page(values, next_token=None): @@ -152,6 +158,43 @@ def test_fetch_with_close_during_request(self, method, expected): assert result_set.rownumber is None assert get_query_results.call_count == 2 + @pytest.mark.parametrize("close_in", ["request", "backoff"]) + @pytest.mark.parametrize( + ("method", "expected"), [("fetchone", None), ("fetchmany", []), ("fetchall", [])] + ) + def test_fetch_with_close_during_retries(self, close_in, method, expected): + """GH-1202: a close() during the retries of a page request stops them. + + The fetch returns nothing, as after a close(), instead of retrying the + throttled request with the policy's attempts and backoff. + """ + result_set = _fetching_result_set(("1", "2"), ("3",)) + assert result_set.fetchmany(2) == [(1,), (2,)] + get_query_results, closer = close_during_next_retries(result_set, close_in) + started = time.monotonic() + assert getattr(result_set, method)() == expected + assert time.monotonic() - started < EVENT_TIMEOUT + if close_in == "backoff": + closer.join() + assert result_set.fetchone() is None + assert get_query_results.call_count == 2 + + @pytest.mark.parametrize("close_in", ["request", "backoff"]) + def test_fetch_all_rows_with_close_during_retries(self, close_in): + """GH-1202: a close() during the retries of a whole read stops them. + + ``_fetch_all_rows_as_csv()`` requests the pages in the same way. + """ + result_set = _fetching_result_set(("1", "2"), ("3",)) + get_query_results, closer = close_during_next_retries(result_set, close_in) + started = time.monotonic() + with pytest.raises(ProgrammingError, match=r"AthenaResultSet is closed\."): + result_set._fetch_all_rows() + assert time.monotonic() - started < EVENT_TIMEOUT + if close_in == "backoff": + closer.join() + assert get_query_results.call_count == 2 + def test_get_query_results_after_close(self): result_set = _fetching_result_set(("1", "2"), ("3",)) result_set.close() diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index cf227edb4..ea038e8af 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -25,6 +25,7 @@ from pyathena.filesystem.s3_core import S3Core from pyathena.glue import GlueMetadataClient from pyathena.model import AthenaCalculationExecutionStatus, AthenaQueryExecution +from pyathena.util import RetryConfig _queries = Environment( loader=FileSystemLoader(Path(__file__).parents[1].resolve() / "resources" / "queries") @@ -79,6 +80,85 @@ def close_and_respond(**kwargs): return get_query_results +def throttling_error(operation_name="GetQueryResults"): + """Build the ``ThrottlingException`` error of an AWS API operation. + + Args: + operation_name: The name of the operation. + + Returns: + The error. + """ + return ClientError( + {"Error": {"Code": "ThrottlingException", "Message": "Rate exceeded"}}, operation_name + ) + + +class ThrottledOperation: + """A stubbed AWS API operation that throttles, signalling its first attempt. + + The first attempt waits for ``release`` if ``block`` is set, so that a test + can stop the retries during the request. + """ + + def __init__(self, block=False, failures=None, operation_name="GetQueryResults"): + """Initialize the operation. + + Args: + block: Whether the first attempt waits for ``release``. + failures: The number of attempts that throttle before one succeeds; + all of them if None. + operation_name: The name of the operation in its error. + """ + self.block = block + self.failures = failures + self.operation_name = operation_name + self.started = threading.Event() + self.release = threading.Event() + self.calls = 0 + + def __call__(self, *args, **kwargs): + self.calls += 1 + if self.calls == 1: + self.started.set() + if self.block: + assert self.release.wait(EVENT_TIMEOUT) + if self.failures is None or self.calls <= self.failures: + raise throttling_error(self.operation_name) + return "success" + + +def close_during_next_retries(result_set, close_in): + """Throttle the next ``GetQueryResults`` request and close the result set during its retries. + + The retry policy of the result set is replaced with one whose first backoff + wait takes at least ``EVENT_TIMEOUT`` seconds, so that a test can tell + whether the close ended it. + + Args: + result_set: A result set whose ``GetQueryResults`` responses are stubbed. + close_in: "request" to close during the request, "backoff" to close from + another thread during the backoff wait after it. + + Returns: + The stubbed ``GetQueryResults`` mock, and the thread that closes the + result set during the backoff wait, if any, which the caller joins. + """ + result_set._retry_config = RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT) + get_query_results = result_set.connection.client.get_query_results + closer = threading.Timer(0.1, result_set.close) + + def throttle(**kwargs): + if close_in == "request": + result_set.close() + else: + closer.start() + raise throttling_error() + + get_query_results.side_effect = throttle + return get_query_results, closer + + def overlapping_calls(calls, timeout=5): """Return a callable that records the peak number of its overlapping calls. From 4b584efed99e35377f8d9edd7b8d632c4aa48611 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 01:18:12 +0900 Subject: [PATCH 3/9] Stop the retries of cancelled asyncio calls in their threads run_in_thread() runs a blocking function as asyncio.to_thread() does, in a stop scope whose event a cancellation of the caller sets. async_retry_api_call() and the to_thread() fetches of the asyncio cursors use it, and create_in_thread() stops a running construction with its existing cancellation event, while the close of an object constructed after the cancellation is shielded. The shielded start requests of kill_on_interrupt keep their retries, since the event exists only in the context copies of the threads. The pandas footer reads run in a copy of the context of their caller, so that a cancelled construction stops them too. Closes #1202 Co-Authored-By: Claude Opus 5.5 --- docs/aio.md | 5 + docs/api/utilities.rst | 4 + pyathena/aio/common.py | 25 +++-- pyathena/aio/s3fs/cursor.py | 2 +- pyathena/aio/util.py | 65 +++++++++-- pyathena/pandas/result_set.py | 15 ++- tests/pyathena/aio/test_cursor.py | 32 ++++++ tests/pyathena/aio/test_util.py | 134 ++++++++++++++++++++++- tests/pyathena/pandas/test_result_set.py | 32 +++++- 9 files changed, 289 insertions(+), 25 deletions(-) diff --git a/docs/aio.md b/docs/aio.md index 9841119b7..ce59924cb 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -166,6 +166,11 @@ 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 PyAthena runs in a thread for it, such as the requests of a fetch or of the status polling. +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 start request of `execute()` with `kill_on_interrupt` enabled keeps its retries, so that the query that it starts can be cancelled as described above. + (aio-dict-cursor)= ## AioDictCursor diff --git a/docs/api/utilities.rst b/docs/api/utilities.rst index bb90d570a..93fff21f2 100644 --- a/docs/api/utilities.rst +++ b/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.shield_retries + .. autofunction:: pyathena.util.is_retryable_error .. autofunction:: pyathena.util.parse_output_location diff --git a/pyathena/aio/common.py b/pyathena/aio/common.py index 217405e8f..885151742 100644 --- a/pyathena/aio/common.py +++ b/pyathena/aio/common.py @@ -9,7 +9,7 @@ from botocore.exceptions import BotoCoreError, ClientError -from pyathena.aio.util import async_retry_api_call, close_in_thread +from pyathena.aio.util import async_retry_api_call, close_in_thread, run_in_thread from pyathena.common import BaseCursor, CursorIterator from pyathena.error import DatabaseError, OperationalError, ProgrammingError from pyathena.glue import GlueMetadataClient @@ -708,7 +708,7 @@ class WithAsyncFetch(WithResultSet, AioBaseCursor, CursorIterator): Combines ``WithResultSet`` with ``AioBaseCursor`` and ``CursorIterator``, and provides async fetch, ``executemany``, and ``cancel``, async iteration, and the async context manager protocol. The fetch methods run - the result set's synchronous fetch with ``asyncio.to_thread``; a subclass + the result set's synchronous fetch with ``run_in_thread()``; a subclass whose result set fetches asynchronously overrides them. Closing the previous result set when a query is executed, and closing the cursor at the end of an ``async with`` block, also run in a thread, because closing @@ -792,8 +792,9 @@ async def fetchone( ) -> tuple[Any | None, ...] | dict[Any, Any | None] | None: """Fetch the next row of the result set. - Wraps the synchronous fetch in ``asyncio.to_thread`` to avoid - blocking the event loop. + Runs the synchronous fetch with + :func:`~pyathena.aio.util.run_in_thread` to avoid blocking the event + loop; cancelling the caller stops the retries of its requests. Returns: A tuple representing the next row, or None if no more rows. @@ -804,7 +805,7 @@ async def fetchone( if not self.has_result_set: raise ProgrammingError("No result set.") result_set = cast(AthenaResultSet, self.result_set) - return await asyncio.to_thread(result_set.fetchone) + return await run_in_thread(result_set.fetchone) @override async def fetchmany( @@ -812,8 +813,9 @@ async def fetchmany( ) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]: """Fetch multiple rows from the result set. - Wraps the synchronous fetch in ``asyncio.to_thread`` to avoid - blocking the event loop. + Runs the synchronous fetch with + :func:`~pyathena.aio.util.run_in_thread` to avoid blocking the event + loop; cancelling the caller stops the retries of its requests. Args: size: Maximum number of rows to fetch. If None or not positive, @@ -828,7 +830,7 @@ async def fetchmany( if not self.has_result_set: raise ProgrammingError("No result set.") result_set = cast(AthenaResultSet, self.result_set) - return await asyncio.to_thread(result_set.fetchmany, size) + return await run_in_thread(result_set.fetchmany, size) @override async def fetchall( @@ -836,8 +838,9 @@ async def fetchall( ) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]: """Fetch all remaining rows from the result set. - Wraps the synchronous fetch in ``asyncio.to_thread`` to avoid - blocking the event loop. + Runs the synchronous fetch with + :func:`~pyathena.aio.util.run_in_thread` to avoid blocking the event + loop; cancelling the caller stops the retries of its requests. Returns: List of tuples representing all remaining rows. @@ -848,7 +851,7 @@ async def fetchall( if not self.has_result_set: raise ProgrammingError("No result set.") result_set = cast(AthenaResultSet, self.result_set) - return await asyncio.to_thread(result_set.fetchall) + return await run_in_thread(result_set.fetchall) @override def __iter__(self) -> NoReturn: diff --git a/pyathena/aio/s3fs/cursor.py b/pyathena/aio/s3fs/cursor.py index d1653658b..e6fd89d79 100644 --- a/pyathena/aio/s3fs/cursor.py +++ b/pyathena/aio/s3fs/cursor.py @@ -33,7 +33,7 @@ class AioS3FSCursor(WithAsyncFetch): Uses ``AioS3FileSystem`` for S3 operations, which replaces ``ThreadPoolExecutor`` parallelism with ``asyncio.gather`` + ``asyncio.to_thread``. Fetch operations are wrapped in - ``asyncio.to_thread()`` because CSV reading is blocking I/O. + ``run_in_thread()`` because CSV reading is blocking I/O. Example: >>> async with await pyathena.aio_connect(...) as conn: diff --git a/pyathena/aio/util.py b/pyathena/aio/util.py index b52c6d741..84b54c110 100644 --- a/pyathena/aio/util.py +++ b/pyathena/aio/util.py @@ -16,7 +16,9 @@ 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, shield_retries, stop_retries_on + +T = TypeVar("T") class _Closable(Protocol): @@ -26,18 +28,54 @@ 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, **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. Args: func: The AWS API function to call. @@ -45,13 +83,21 @@ async def async_retry_api_call( 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``. **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 + return await run_in_thread( + retry_api_call, + func, + config, + logger, + *args, + stop_on=stop_on, + stop_when=stop_when, + **kwargs, ) @@ -76,7 +122,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 @@ -98,12 +145,14 @@ async def create_in_thread( def construct() -> _ClosableT | None: 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 shield_retries(): + obj.close() return None future = asyncio.get_running_loop().run_in_executor( diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 2762b9b86..17949c1a5 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -10,9 +10,10 @@ from collections.abc import Callable, Iterable, Iterator from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack, contextmanager +from contextvars import Context, copy_context from functools import partial from io import BufferedReader, BytesIO, IOBase, StringIO, TextIOWrapper -from itertools import chain, islice +from itertools import chain, islice, repeat from multiprocessing import cpu_count from types import GeneratorType from typing import ( @@ -1625,11 +1626,21 @@ 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( + Context.run, + contexts, + repeat(self._read_parquet_footer), + self._data_manifest, + ) + ) except Exception as e: _logger.exception(f"Failed to read {self._unload_location}.") raise OperationalError(*e.args) from e diff --git a/tests/pyathena/aio/test_cursor.py b/tests/pyathena/aio/test_cursor.py index c20e81747..101988f7a 100644 --- a/tests/pyathena/aio/test_cursor.py +++ b/tests/pyathena/aio/test_cursor.py @@ -20,6 +20,7 @@ from tests.pyathena.aio.conftest import _aio_connect from tests.pyathena.util import ( EVENT_TIMEOUT, + ThrottledOperation, succeeded_query_execution, throttle_metadata_api, ) @@ -369,6 +370,37 @@ async def test_execute_cancelled_while_starting(self): assert cursor.query_id == "query_id" assert cursor.result_set is None + async def test_execute_cancelled_while_start_is_retried(self): + """GH-1202: a cancellation does not stop the retries of the shielded start (no AWS). + + The query that the retried request starts is then cancelled. + """ + cursor, cancel, _, _ = _starting_cursor() + operation = ThrottledOperation(block=True, failures=1, operation_name="StartQueryExecution") + + def start_query_execution(**kwargs): + operation(**kwargs) + return {"QueryExecutionId": "query_id"} + + cursor._connection.client.start_query_execution.side_effect = start_query_execution + task = asyncio.create_task(cursor.execute("SELECT 1")) + try: + assert await asyncio.to_thread(operation.started.wait, EVENT_TIMEOUT) + task.cancel() + # Let the task handle the cancellation before the first attempt fails. + for _ in range(5): + await asyncio.sleep(0) + assert not task.done() + finally: + operation.release.set() + with pytest.raises(asyncio.CancelledError): + await task + + assert task.cancelled() + assert operation.calls == 2 + cancel.assert_awaited_once_with("query_id") + assert cursor.query_id == "query_id" + async def test_execute_cancelled_before_request_is_sent(self): """Cancellation before the start task begins sends no request (no AWS).""" cursor, cancel, started, release = _starting_cursor() diff --git a/tests/pyathena/aio/test_util.py b/tests/pyathena/aio/test_util.py index 666c5a8b4..aeb233fca 100644 --- a/tests/pyathena/aio/test_util.py +++ b/tests/pyathena/aio/test_util.py @@ -8,16 +8,42 @@ import contextvars import gc import threading +import time from concurrent.futures import ThreadPoolExecutor import pytest +from botocore.exceptions import ClientError -from pyathena.aio.util import create_in_thread -from tests.pyathena.util import EVENT_TIMEOUT +from pyathena.aio.util import async_retry_api_call, create_in_thread, run_in_thread +from pyathena.util import RetryConfig, retry_api_call +from tests.pyathena.util import EVENT_TIMEOUT, ThrottledOperation _VALUE: contextvars.ContextVar[str] = contextvars.ContextVar("_VALUE", default="unset") +async def _cancel_once_started(task, operation, stop_in): + """Cancel a task once its operation has made its first attempt. + + Args: + task: The task that runs the operation in a thread. + operation: The blocking ``ThrottledOperation`` of the task. + stop_in: "request" to cancel during the first attempt, "backoff" to + cancel during the backoff wait after it. + """ + # Polled, since the default executor may have no free thread. + deadline = time.monotonic() + EVENT_TIMEOUT + while not operation.started.is_set(): + assert time.monotonic() < deadline + await asyncio.sleep(0.01) + if stop_in == "backoff": + # Lets the first attempt fail and the backoff wait begin. + await asyncio.sleep(0.1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + operation.release.set() + + class _Resource: """A closable object that records its constructor and close threads.""" @@ -180,3 +206,107 @@ def factory(): finally: release.set() executor.shutdown(wait=True) + + +async def test_run_in_thread(): + """The function runs in another thread with the caller's context.""" + _VALUE.set("caller") + thread, value = await run_in_thread(lambda: (threading.get_ident(), _VALUE.get())) + assert thread != threading.get_ident() + assert value == "caller" + + +@pytest.mark.parametrize("stop_in", ["request", "backoff"]) +async def test_async_retry_api_call_cancelled(stop_in): + """GH-1202: cancelling the caller makes no further attempt and ends the backoff wait. + + Without the stop, the thread would wait at least EVENT_TIMEOUT seconds before + the next attempt, holding the only thread of the default executor. + """ + loop = asyncio.get_running_loop() + executor = ThreadPoolExecutor(max_workers=1) + loop.set_default_executor(executor) + operation = ThrottledOperation(block=stop_in == "request") + try: + task = asyncio.create_task( + async_retry_api_call(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) + ) + await _cancel_once_started(task, operation, stop_in) + started = time.monotonic() + # Runs once the thread of the cancelled call is free. + await loop.run_in_executor(None, lambda: None) + assert time.monotonic() - started < EVENT_TIMEOUT + assert operation.calls == 1 + finally: + operation.release.set() + executor.shutdown(wait=True) + + +async def test_async_retry_api_call_stop_when(): + """``stop_when`` is passed to ``retry_api_call()``.""" + operation = ThrottledOperation() + with pytest.raises(ClientError): + await async_retry_api_call( + operation, RetryConfig(attempt=3, multiplier=0), stop_when=lambda: True + ) + assert operation.calls == 1 + + +async def test_async_retry_api_call_retries(): + """Without a cancellation, the retries continue.""" + operation = ThrottledOperation(failures=2) + assert await async_retry_api_call(operation, RetryConfig(attempt=3, multiplier=0)) == "success" + assert operation.calls == 3 + + +@pytest.mark.parametrize("stop_in", ["request", "backoff"]) +async def test_create_in_thread_cancelled_stops_retries(stop_in): + """GH-1202: cancelling the caller stops the retries of a running construction.""" + operation = ThrottledOperation(block=stop_in == "request") + finished = threading.Event() + + def factory(): + try: + retry_api_call(operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) + finally: + finished.set() + + task = asyncio.create_task(create_in_thread(factory)) + try: + await _cancel_once_started(task, operation, stop_in) + started = time.monotonic() + assert await asyncio.to_thread(finished.wait, EVENT_TIMEOUT) + assert time.monotonic() - started < EVENT_TIMEOUT + assert operation.calls == 1 + finally: + operation.release.set() + + +async def test_create_in_thread_cleanup_not_stopped(): + """The close of an object constructed after a cancellation retries as usual.""" + started = threading.Event() + release = threading.Event() + close_operation = ThrottledOperation(failures=2) + closed = threading.Event() + + class _ClosedWithRequest(_Resource): + def close(self): + retry_api_call(close_operation, RetryConfig(attempt=3, multiplier=0)) + closed.set() + + def factory(): + started.set() + assert release.wait(EVENT_TIMEOUT) + return _ClosedWithRequest() + + task = asyncio.create_task(create_in_thread(factory)) + try: + assert await asyncio.to_thread(started.wait, EVENT_TIMEOUT) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + release.set() + assert await asyncio.to_thread(closed.wait, EVENT_TIMEOUT) + assert close_operation.calls == 3 + finally: + release.set() diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index 6dee08c4b..186bb8842 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -33,7 +33,8 @@ _no_trunc_date, _read_csv_with_pyarrow, ) -from pyathena.util import RetryConfig +from pyathena.util import RetryConfig, retry_api_call, stop_retries_on +from tests.pyathena.util import ThrottledOperation def _typed_values(rows): @@ -617,6 +618,35 @@ def test_iter_parquet_chunks_after_close(self, tmp_path): result_set.close() assert [chunk["i"].tolist() for chunk in reader] == [[1], [2]] + def test_read_parquet_footers_stopped(self, tmp_path): + """GH-1202: the footer reads run in the stop scope of their caller. + + A cancelled ``create_in_thread()`` of AioPandasCursor stops their retries too. + """ + result_set, paths = _unload_result_set( + tmp_path, [pa.table({"i": [1]}), pa.table({"i": [2]})], chunksize=1 + ) + stop = threading.Event() + stop.set() + operations = [] + + def read_parquet_footer(path): + operation = ThrottledOperation() + operations.append(operation) + retry_api_call(operation, RetryConfig(attempt=3, multiplier=0)) + + with ( + patch.object(AthenaPandasResultSet, "_read_data_manifest", return_value=paths), + patch.object( + AthenaPandasResultSet, "_read_parquet_footer", side_effect=read_parquet_footer + ), + stop_retries_on(stop), + pytest.raises(OperationalError, match="ThrottlingException"), + ): + result_set._read_parquet_footers() + assert operations + assert [operation.calls for operation in operations] == [1] * len(operations) + def test_as_pandas_parquet_chunks_engine(self, tmp_path): """Chunks of an UNLOAD result need the pyarrow engine, as the whole result does.""" result_set, paths = _unload_result_set(tmp_path, [pa.table({"i": [1]})], chunksize=1) From 78f65558c47b0e5850c6dda1dff4fc1a620647e3 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 07:49:27 +0900 Subject: [PATCH 4/9] Keep the retries of cleanup requests and of the cursors' fetches StopQueryExecution, StopCalculationExecution, and TerminateSession clean up, so async_retry_api_call(stoppable=False) keeps their retries after a cancellation, as asyncio.to_thread() did. The fetches of the asyncio pandas, Arrow, Polars, and S3FS cursors go back to asyncio.to_thread(): their retries are mostly S3 reads of a reader that later fetches continue, which a stopped read would leave part-way; #1206 defines a cancelled fetch. Also documents construct() and reuses the shared throttling helper. Refs #1202 Co-Authored-By: Claude Opus 5.5 --- docs/aio.md | 5 +++-- pyathena/aio/common.py | 26 ++++++++++++-------------- pyathena/aio/s3fs/cursor.py | 2 +- pyathena/aio/spark/cursor.py | 2 ++ pyathena/aio/util.py | 16 ++++++++++++++-- tests/pyathena/aio/test_cursor.py | 22 ++++++++++++++++++++++ tests/pyathena/aio/test_util.py | 16 ++++++++++++++++ tests/pyathena/test_util.py | 19 ++++++------------- 8 files changed, 76 insertions(+), 32 deletions(-) diff --git a/docs/aio.md b/docs/aio.md index ce59924cb..acd5b4f15 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -166,10 +166,11 @@ 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 PyAthena runs in a thread for it, such as the requests of a fetch or of the status polling. +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 start request of `execute()` with `kill_on_interrupt` enabled keeps its retries, so that the query that it starts can be cancelled as described above. +The start request of `execute()` with `kill_on_interrupt` enabled, the requests that stop a query or a calculation, and the request that terminates a Spark session keep their retries. +The fetch methods of `AioPandasCursor`, `AioArrowCursor`, `AioPolarsCursor`, and `AioS3FSCursor` run to their end in their thread. (aio-dict-cursor)= diff --git a/pyathena/aio/common.py b/pyathena/aio/common.py index 885151742..219a7598b 100644 --- a/pyathena/aio/common.py +++ b/pyathena/aio/common.py @@ -9,7 +9,7 @@ from botocore.exceptions import BotoCoreError, ClientError -from pyathena.aio.util import async_retry_api_call, close_in_thread, run_in_thread +from pyathena.aio.util import async_retry_api_call, close_in_thread from pyathena.common import BaseCursor, CursorIterator from pyathena.error import DatabaseError, OperationalError, ProgrammingError from pyathena.glue import GlueMetadataClient @@ -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: @@ -708,7 +709,7 @@ class WithAsyncFetch(WithResultSet, AioBaseCursor, CursorIterator): Combines ``WithResultSet`` with ``AioBaseCursor`` and ``CursorIterator``, and provides async fetch, ``executemany``, and ``cancel``, async iteration, and the async context manager protocol. The fetch methods run - the result set's synchronous fetch with ``run_in_thread()``; a subclass + the result set's synchronous fetch with ``asyncio.to_thread``; a subclass whose result set fetches asynchronously overrides them. Closing the previous result set when a query is executed, and closing the cursor at the end of an ``async with`` block, also run in a thread, because closing @@ -792,9 +793,8 @@ async def fetchone( ) -> tuple[Any | None, ...] | dict[Any, Any | None] | None: """Fetch the next row of the result set. - Runs the synchronous fetch with - :func:`~pyathena.aio.util.run_in_thread` to avoid blocking the event - loop; cancelling the caller stops the retries of its requests. + Wraps the synchronous fetch in ``asyncio.to_thread`` to avoid + blocking the event loop. Returns: A tuple representing the next row, or None if no more rows. @@ -805,7 +805,7 @@ async def fetchone( if not self.has_result_set: raise ProgrammingError("No result set.") result_set = cast(AthenaResultSet, self.result_set) - return await run_in_thread(result_set.fetchone) + return await asyncio.to_thread(result_set.fetchone) @override async def fetchmany( @@ -813,9 +813,8 @@ async def fetchmany( ) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]: """Fetch multiple rows from the result set. - Runs the synchronous fetch with - :func:`~pyathena.aio.util.run_in_thread` to avoid blocking the event - loop; cancelling the caller stops the retries of its requests. + Wraps the synchronous fetch in ``asyncio.to_thread`` to avoid + blocking the event loop. Args: size: Maximum number of rows to fetch. If None or not positive, @@ -830,7 +829,7 @@ async def fetchmany( if not self.has_result_set: raise ProgrammingError("No result set.") result_set = cast(AthenaResultSet, self.result_set) - return await run_in_thread(result_set.fetchmany, size) + return await asyncio.to_thread(result_set.fetchmany, size) @override async def fetchall( @@ -838,9 +837,8 @@ async def fetchall( ) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]: """Fetch all remaining rows from the result set. - Runs the synchronous fetch with - :func:`~pyathena.aio.util.run_in_thread` to avoid blocking the event - loop; cancelling the caller stops the retries of its requests. + Wraps the synchronous fetch in ``asyncio.to_thread`` to avoid + blocking the event loop. Returns: List of tuples representing all remaining rows. @@ -851,7 +849,7 @@ async def fetchall( if not self.has_result_set: raise ProgrammingError("No result set.") result_set = cast(AthenaResultSet, self.result_set) - return await run_in_thread(result_set.fetchall) + return await asyncio.to_thread(result_set.fetchall) @override def __iter__(self) -> NoReturn: diff --git a/pyathena/aio/s3fs/cursor.py b/pyathena/aio/s3fs/cursor.py index e6fd89d79..d1653658b 100644 --- a/pyathena/aio/s3fs/cursor.py +++ b/pyathena/aio/s3fs/cursor.py @@ -33,7 +33,7 @@ class AioS3FSCursor(WithAsyncFetch): Uses ``AioS3FileSystem`` for S3 operations, which replaces ``ThreadPoolExecutor`` parallelism with ``asyncio.gather`` + ``asyncio.to_thread``. Fetch operations are wrapped in - ``run_in_thread()`` because CSV reading is blocking I/O. + ``asyncio.to_thread()`` because CSV reading is blocking I/O. Example: >>> async with await pyathena.aio_connect(...) as conn: diff --git a/pyathena/aio/spark/cursor.py b/pyathena/aio/spark/cursor.py index cc9e6502c..e6f0c6668 100644 --- a/pyathena/aio/spark/cursor.py +++ b/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: diff --git a/pyathena/aio/util.py b/pyathena/aio/util.py index 84b54c110..329f7b54d 100644 --- a/pyathena/aio/util.py +++ b/pyathena/aio/util.py @@ -69,13 +69,14 @@ async def async_retry_api_call( *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. 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. + further attempt and ends the backoff wait, unless ``stoppable`` is false. Args: func: The AWS API function to call. @@ -84,12 +85,17 @@ async def async_retry_api_call( *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, as with ``asyncio.to_thread()``. **kwargs: Keyword arguments to pass to the function. Returns: The result of the successful function call. """ - return await run_in_thread( + run = run_in_thread if stoppable else asyncio.to_thread + return await run( retry_api_call, func, config, @@ -143,6 +149,12 @@ 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 with stop_retries_on(cancelled): diff --git a/tests/pyathena/aio/test_cursor.py b/tests/pyathena/aio/test_cursor.py index 101988f7a..2d899b77a 100644 --- a/tests/pyathena/aio/test_cursor.py +++ b/tests/pyathena/aio/test_cursor.py @@ -401,6 +401,28 @@ def start_query_execution(**kwargs): cancel.assert_awaited_once_with("query_id") assert cursor.query_id == "query_id" + async def test_cancel_cancelled_while_retrying(self): + """GH-1202: a cancellation does not stop the retries of StopQueryExecution (no AWS).""" + cursor, _, _ = _offline_cursor( + kill_on_interrupt=True, final_state=AthenaQueryExecution.STATE_CANCELLED + ) + del cursor._cancel # use the real _cancel() + cursor._retry_config = RetryConfig(attempt=3, multiplier=0) + operation = ThrottledOperation(block=True, failures=2, operation_name="StopQueryExecution") + cursor._connection.client.stop_query_execution.side_effect = operation + task = asyncio.create_task(cursor._cancel("query_id")) + try: + assert await asyncio.to_thread(operation.started.wait, EVENT_TIMEOUT) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + finally: + operation.release.set() + deadline = asyncio.get_running_loop().time() + EVENT_TIMEOUT + while operation.calls < 3: + assert asyncio.get_running_loop().time() < deadline + await asyncio.sleep(0.01) + async def test_execute_cancelled_before_request_is_sent(self): """Cancellation before the start task begins sends no request (no AWS).""" cursor, cancel, started, release = _starting_cursor() diff --git a/tests/pyathena/aio/test_util.py b/tests/pyathena/aio/test_util.py index aeb233fca..876bbcc00 100644 --- a/tests/pyathena/aio/test_util.py +++ b/tests/pyathena/aio/test_util.py @@ -252,6 +252,22 @@ async def test_async_retry_api_call_stop_when(): assert operation.calls == 1 +async def test_async_retry_api_call_not_stoppable(): + """GH-1202: a cleanup request that is not stoppable keeps its retries after a cancellation.""" + operation = ThrottledOperation(block=True, failures=2) + task = asyncio.create_task( + async_retry_api_call(operation, RetryConfig(attempt=3, multiplier=0), stoppable=False) + ) + try: + await _cancel_once_started(task, operation, "request") + deadline = time.monotonic() + EVENT_TIMEOUT + while operation.calls < 3: + assert time.monotonic() < deadline + await asyncio.sleep(0.01) + finally: + operation.release.set() + + async def test_async_retry_api_call_retries(): """Without a cancellation, the retries continue.""" operation = ThrottledOperation(failures=2) diff --git a/tests/pyathena/test_util.py b/tests/pyathena/test_util.py index 83dc2ef36..62b7da7d5 100644 --- a/tests/pyathena/test_util.py +++ b/tests/pyathena/test_util.py @@ -18,7 +18,7 @@ stop_retries_on, strtobool, ) -from tests.pyathena.util import EVENT_TIMEOUT +from tests.pyathena.util import EVENT_TIMEOUT, throttling_error def test_parse_output_location(): @@ -207,13 +207,6 @@ def call(): assert calls == 3 -def _throttling_error() -> ClientError: - return ClientError( - {"Error": {"Code": "ThrottlingException", "Message": "Rate exceeded"}}, - "GetTableMetadata", - ) - - class TestRetryConfig: def test_default_attempts(self): assert RetryConfig().attempt == 8 @@ -240,7 +233,7 @@ def test_captures_exception_names(self, policy_type): def test_retry_api_call_waits_with_jitter(monkeypatch, multiplier, max_delay, expected_bases): sleeps = [] monkeypatch.setattr(tenacity.nap.time, "sleep", sleeps.append) - error = _throttling_error() + error = throttling_error("GetTableMetadata") calls = 0 def call(): @@ -260,7 +253,7 @@ def call(): def test_retry_api_call_jitter_varies(monkeypatch): sleeps = [] monkeypatch.setattr(tenacity.nap.time, "sleep", sleeps.append) - error = _throttling_error() + error = throttling_error("GetTableMetadata") def call(): raise error @@ -392,7 +385,7 @@ def __init__(self, stop, stop_in): """ self.stop = stop self.stop_in = stop_in - self.error = _throttling_error() + self.error = throttling_error("GetTableMetadata") self.calls = 0 self.timer = None @@ -473,7 +466,7 @@ def test_retry_api_call_shielded(): """Retries in a shielded block ignore the events of the enclosing scopes.""" event = threading.Event() event.set() - error = _throttling_error() + error = throttling_error("GetTableMetadata") calls = 0 def call(): @@ -503,7 +496,7 @@ def test_retry_api_call_scope_ends_with_block(): event = threading.Event() with stop_retries_on(event): event.set() - error = _throttling_error() + error = throttling_error("GetTableMetadata") calls = 0 def call(): From 27641102be8b19cdff8ac83d8d8efe194ba881a7 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 07:54:51 +0900 Subject: [PATCH 5/9] Shield non-stoppable requests and simplify the stop helpers and tests async_retry_api_call(stoppable=False) now runs its retries in shield_retries(), so a stop scope of the caller does not stop a cleanup request either. retry_api_call() defines its stop hooks at function level, and the pandas footer reads pass each copied context with a lambda. The tests share one throttling stub that stops its retries (StoppingOperation), reuse ThrottledOperation, and wait with wait_until(). Refs #1202 Co-Authored-By: Claude Opus 5.5 --- pyathena/aio/util.py | 26 ++++----- pyathena/pandas/result_set.py | 7 +-- pyathena/util.py | 74 +++++++++++------------ tests/pyathena/aio/test_cursor.py | 6 +- tests/pyathena/aio/test_result_set.py | 7 +-- tests/pyathena/aio/test_util.py | 27 +++++---- tests/pyathena/test_result_set.py | 14 ++--- tests/pyathena/test_util.py | 84 ++++++--------------------- tests/pyathena/util.py | 77 ++++++++++++++++++------ 9 files changed, 160 insertions(+), 162 deletions(-) diff --git a/pyathena/aio/util.py b/pyathena/aio/util.py index 329f7b54d..32740b1d3 100644 --- a/pyathena/aio/util.py +++ b/pyathena/aio/util.py @@ -10,6 +10,7 @@ from __future__ import annotations import asyncio +import contextlib import contextvars import logging import threading @@ -87,24 +88,23 @@ async def async_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, as with ``asyncio.to_thread()``. + terminates a session, passes False: its retries then ignore the + stop scopes (see :func:`~pyathena.util.shield_retries`) and run to + their end in the thread, as with ``asyncio.to_thread()``. **kwargs: Keyword arguments to pass to the function. Returns: The result of the successful function call. """ - run = run_in_thread if stoppable else asyncio.to_thread - return await run( - retry_api_call, - func, - config, - logger, - *args, - stop_on=stop_on, - stop_when=stop_when, - **kwargs, - ) + + def call() -> Any: + """Make the call, ignoring the stop scopes unless it is stoppable.""" + with contextlib.nullcontext() if stoppable else shield_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: diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 17949c1a5..290c8a9a2 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -10,10 +10,10 @@ from collections.abc import Callable, Iterable, Iterator from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack, contextmanager -from contextvars import Context, copy_context +from contextvars import copy_context from functools import partial from io import BufferedReader, BytesIO, IOBase, StringIO, TextIOWrapper -from itertools import chain, islice, repeat +from itertools import chain, islice from multiprocessing import cpu_count from types import GeneratorType from typing import ( @@ -1635,9 +1635,8 @@ def _read_parquet_footers(self) -> list[FileMetaData]: with ThreadPoolExecutor(max_workers=workers) as executor: footers = list( executor.map( - Context.run, + lambda context, path: context.run(self._read_parquet_footer, path), contexts, - repeat(self._read_parquet_footer), self._data_manifest, ) ) diff --git a/pyathena/util.py b/pyathena/util.py index 6a2c50161..1e396e754 100644 --- a/pyathena/util.py +++ b/pyathena/util.py @@ -378,44 +378,46 @@ def stopped() -> bool: """Whether an event of the enclosing scopes is set or ``stop_when`` holds.""" return any(event.is_set() for event in events) or (stop_when is not None and stop_when()) - # Without a stop, tenacity's own sleep is kept. - stop_kwargs: dict[str, Any] = {} # The error of the attempt before a backoff wait, raised if the wait is stopped. failures: list[BaseException] = [] - if events or stop_when is not None: - - def record_failure(retry_state: RetryCallState) -> None: - """Keep the error of the failed attempt, which a stopped wait raises. - - Args: - retry_state: The state of the call before the backoff wait. - """ - outcome = retry_state.outcome - if outcome is not None and (error := outcome.exception()) is not None: - failures[:] = [error] - - def sleep(seconds: float) -> None: - """Wait before the next attempt, ending the wait once the retries are stopped. - - Args: - seconds: The backoff time. - - Raises: - BaseException: The error of the last attempt, instead of - making the next one, if the retries are stopped. - """ - deadline = time.monotonic() + seconds - while not stopped(): - remaining = deadline - time.monotonic() - if remaining <= 0: - return - if len(events) == 1 and stop_when is None: - events[0].wait(remaining) - else: - time.sleep(min(remaining, _STOP_POLL_INTERVAL)) - raise failures[0] - - stop_kwargs = {"before_sleep": record_failure, "sleep": sleep} + # A single event can be waited for; anything else is polled. + wait_for_event = len(events) == 1 and stop_when is None + + def record_failure(retry_state: RetryCallState) -> None: + """Keep the error of the failed attempt, which a stopped wait raises. + + Args: + retry_state: The state of the call before the backoff wait. + """ + outcome = retry_state.outcome + if outcome is not None and (error := outcome.exception()) is not None: + failures[:] = [error] + + def sleep(seconds: float) -> None: + """Wait before the next attempt, ending the wait once the retries are stopped. + + Args: + seconds: The backoff time. + + Raises: + BaseException: The error of the last attempt, instead of + making the next one, if the retries are stopped. + """ + deadline = time.monotonic() + seconds + while not stopped(): + remaining = deadline - time.monotonic() + if remaining <= 0: + return + if wait_for_event: + events[0].wait(remaining) + else: + time.sleep(min(remaining, _STOP_POLL_INTERVAL)) + raise failures[0] + + # Without a stop, tenacity's own sleep is kept. + stop_kwargs: dict[str, Any] = ( + {"before_sleep": record_failure, "sleep": sleep} if events or stop_when is not None else {} + ) retry = tenacity.Retrying( retry=retry_if_exception( diff --git a/tests/pyathena/aio/test_cursor.py b/tests/pyathena/aio/test_cursor.py index 2d899b77a..c522c5b35 100644 --- a/tests/pyathena/aio/test_cursor.py +++ b/tests/pyathena/aio/test_cursor.py @@ -23,6 +23,7 @@ ThrottledOperation, succeeded_query_execution, throttle_metadata_api, + wait_until, ) @@ -418,10 +419,7 @@ async def test_cancel_cancelled_while_retrying(self): await task finally: operation.release.set() - deadline = asyncio.get_running_loop().time() + EVENT_TIMEOUT - while operation.calls < 3: - assert asyncio.get_running_loop().time() < deadline - await asyncio.sleep(0.01) + await wait_until(lambda: operation.calls == 3) async def test_execute_cancelled_before_request_is_sent(self): """Cancellation before the start task begins sends no request (no AWS).""" diff --git a/tests/pyathena/aio/test_result_set.py b/tests/pyathena/aio/test_result_set.py index 205be0b38..6b22afcb9 100644 --- a/tests/pyathena/aio/test_result_set.py +++ b/tests/pyathena/aio/test_result_set.py @@ -131,14 +131,13 @@ async def test_fetch_with_close_during_retries(self, close_in, method, expected) """GH-1202: a close() during the retries of a page request stops them.""" result_set = await _create_result_set(AthenaAioResultSet, [("1", "2"), ("3",)]) assert await result_set.fetchmany(2) == [(1,), (2,)] - get_query_results, closer = close_during_next_retries(result_set, close_in) + operation = close_during_next_retries(result_set, close_in) started = time.monotonic() assert await getattr(result_set, method)() == expected assert time.monotonic() - started < EVENT_TIMEOUT - if close_in == "backoff": - closer.join() + operation.join() assert await result_set.fetchone() is None - assert get_query_results.call_count == 2 + assert operation.calls == 1 async def test_get_query_results_after_close(self): result_set = await _create_result_set(AthenaAioResultSet, [("1", "2"), ("3",)]) diff --git a/tests/pyathena/aio/test_util.py b/tests/pyathena/aio/test_util.py index 876bbcc00..35764cb62 100644 --- a/tests/pyathena/aio/test_util.py +++ b/tests/pyathena/aio/test_util.py @@ -15,8 +15,8 @@ from botocore.exceptions import ClientError from pyathena.aio.util import async_retry_api_call, create_in_thread, run_in_thread -from pyathena.util import RetryConfig, retry_api_call -from tests.pyathena.util import EVENT_TIMEOUT, ThrottledOperation +from pyathena.util import RetryConfig, retry_api_call, stop_retries_on +from tests.pyathena.util import EVENT_TIMEOUT, ThrottledOperation, wait_until _VALUE: contextvars.ContextVar[str] = contextvars.ContextVar("_VALUE", default="unset") @@ -31,10 +31,7 @@ async def _cancel_once_started(task, operation, stop_in): cancel during the backoff wait after it. """ # Polled, since the default executor may have no free thread. - deadline = time.monotonic() + EVENT_TIMEOUT - while not operation.started.is_set(): - assert time.monotonic() < deadline - await asyncio.sleep(0.01) + await wait_until(operation.started.is_set) if stop_in == "backoff": # Lets the first attempt fail and the backoff wait begin. await asyncio.sleep(0.1) @@ -260,14 +257,24 @@ async def test_async_retry_api_call_not_stoppable(): ) try: await _cancel_once_started(task, operation, "request") - deadline = time.monotonic() + EVENT_TIMEOUT - while operation.calls < 3: - assert time.monotonic() < deadline - await asyncio.sleep(0.01) + await wait_until(lambda: operation.calls == 3) finally: operation.release.set() +async def test_async_retry_api_call_not_stoppable_in_stopped_scope(): + """A request that is not stoppable ignores the stop scopes of its caller.""" + stop = threading.Event() + stop.set() + operation = ThrottledOperation(failures=2) + with stop_retries_on(stop): + result = await async_retry_api_call( + operation, RetryConfig(attempt=3, multiplier=0), stoppable=False + ) + assert result == "success" + assert operation.calls == 3 + + async def test_async_retry_api_call_retries(): """Without a cancellation, the retries continue.""" operation = ThrottledOperation(failures=2) diff --git a/tests/pyathena/test_result_set.py b/tests/pyathena/test_result_set.py index 89fdda0f8..6d87f73cb 100644 --- a/tests/pyathena/test_result_set.py +++ b/tests/pyathena/test_result_set.py @@ -170,14 +170,13 @@ def test_fetch_with_close_during_retries(self, close_in, method, expected): """ result_set = _fetching_result_set(("1", "2"), ("3",)) assert result_set.fetchmany(2) == [(1,), (2,)] - get_query_results, closer = close_during_next_retries(result_set, close_in) + operation = close_during_next_retries(result_set, close_in) started = time.monotonic() assert getattr(result_set, method)() == expected assert time.monotonic() - started < EVENT_TIMEOUT - if close_in == "backoff": - closer.join() + operation.join() assert result_set.fetchone() is None - assert get_query_results.call_count == 2 + assert operation.calls == 1 @pytest.mark.parametrize("close_in", ["request", "backoff"]) def test_fetch_all_rows_with_close_during_retries(self, close_in): @@ -186,14 +185,13 @@ def test_fetch_all_rows_with_close_during_retries(self, close_in): ``_fetch_all_rows_as_csv()`` requests the pages in the same way. """ result_set = _fetching_result_set(("1", "2"), ("3",)) - get_query_results, closer = close_during_next_retries(result_set, close_in) + operation = close_during_next_retries(result_set, close_in) started = time.monotonic() with pytest.raises(ProgrammingError, match=r"AthenaResultSet is closed\."): result_set._fetch_all_rows() assert time.monotonic() - started < EVENT_TIMEOUT - if close_in == "backoff": - closer.join() - assert get_query_results.call_count == 2 + operation.join() + assert operation.calls == 1 def test_get_query_results_after_close(self): result_set = _fetching_result_set(("1", "2"), ("3",)) diff --git a/tests/pyathena/test_util.py b/tests/pyathena/test_util.py index 62b7da7d5..307d966ec 100644 --- a/tests/pyathena/test_util.py +++ b/tests/pyathena/test_util.py @@ -18,7 +18,12 @@ stop_retries_on, strtobool, ) -from tests.pyathena.util import EVENT_TIMEOUT, throttling_error +from tests.pyathena.util import ( + EVENT_TIMEOUT, + StoppingOperation, + ThrottledOperation, + throttling_error, +) def test_parse_output_location(): @@ -372,39 +377,6 @@ def call(): assert calls == expected_calls -class _ThrottledCall: - """A call that always throttles and stops its retries at a given point.""" - - def __init__(self, stop, stop_in): - """Initialize the call. - - Args: - stop: Stops the retries; called once. - stop_in: "request" to stop during the first attempt, "backoff" to - stop from another thread during the backoff wait after it. - """ - self.stop = stop - self.stop_in = stop_in - self.error = throttling_error("GetTableMetadata") - self.calls = 0 - self.timer = None - - def __call__(self): - if self.calls == 0: - if self.stop_in == "request": - self.stop() - else: - self.timer = threading.Timer(0.1, self.stop) - self.timer.start() - self.calls += 1 - raise self.error - - def join(self): - """Wait for the thread that stops the backoff wait, if any.""" - if self.timer is not None: - self.timer.join() - - @pytest.mark.parametrize("stop_in", ["request", "backoff"]) def test_retry_api_call_stopped_by_event(stop_in): """GH-1202: the event of a stop scope ends the retries and their backoff wait. @@ -412,7 +384,7 @@ def test_retry_api_call_stopped_by_event(stop_in): Without the stop, the backoff wait would take at least EVENT_TIMEOUT seconds. """ event = threading.Event() - call = _ThrottledCall(event.set, stop_in) + call = StoppingOperation(event.set, stop_in) started = time.monotonic() with stop_retries_on(event), pytest.raises(ClientError) as caught: retry_api_call(call, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT)) @@ -426,7 +398,7 @@ def test_retry_api_call_stopped_by_event(stop_in): def test_retry_api_call_stopped_by_stop_when(stop_in): """GH-1202: ``stop_when`` ends the retries and their backoff wait.""" stopped = threading.Event() - call = _ThrottledCall(stopped.set, stop_in) + call = StoppingOperation(stopped.set, stop_in) started = time.monotonic() with pytest.raises(ClientError) as caught: retry_api_call( @@ -442,7 +414,7 @@ def test_retry_api_call_stopped_before_call(): """A call in a stopped scope makes its first attempt, but no retry.""" event = threading.Event() event.set() - call = _ThrottledCall(event.set, "request") + call = StoppingOperation(event.set, "request") with stop_retries_on(event), pytest.raises(ClientError): retry_api_call(call, RetryConfig(multiplier=0)) assert call.calls == 1 @@ -451,7 +423,7 @@ def test_retry_api_call_stopped_before_call(): def test_retry_api_call_stopped_by_enclosing_scope(): """The event of any enclosing scope stops the retries.""" outer = threading.Event() - call = _ThrottledCall(outer.set, "backoff") + call = StoppingOperation(outer.set, "backoff") with ( stop_retries_on(outer), stop_retries_on(threading.Event()), @@ -466,29 +438,20 @@ def test_retry_api_call_shielded(): """Retries in a shielded block ignore the events of the enclosing scopes.""" event = threading.Event() event.set() - error = throttling_error("GetTableMetadata") - calls = 0 - - def call(): - nonlocal calls - calls += 1 - if calls < 3: - raise error - return "success" - + operation = ThrottledOperation(failures=2) with stop_retries_on(event), shield_retries(): - assert retry_api_call(call, RetryConfig(multiplier=0)) == "success" - assert calls == 3 + assert retry_api_call(operation, RetryConfig(multiplier=0)) == "success" + assert operation.calls == 3 # A scope opened in the shielded block stops the retries again. - calls = 0 + operation = ThrottledOperation(failures=2) with ( stop_retries_on(event), shield_retries(), stop_retries_on(event), pytest.raises(ClientError), ): - retry_api_call(call, RetryConfig(multiplier=0)) - assert calls == 1 + retry_api_call(operation, RetryConfig(multiplier=0)) + assert operation.calls == 1 def test_retry_api_call_scope_ends_with_block(): @@ -496,15 +459,6 @@ def test_retry_api_call_scope_ends_with_block(): event = threading.Event() with stop_retries_on(event): event.set() - error = throttling_error("GetTableMetadata") - calls = 0 - - def call(): - nonlocal calls - calls += 1 - if calls < 3: - raise error - return "success" - - assert retry_api_call(call, RetryConfig(multiplier=0)) == "success" - assert calls == 3 + operation = ThrottledOperation(failures=2) + assert retry_api_call(operation, RetryConfig(multiplier=0)) == "success" + assert operation.calls == 3 diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index ea038e8af..e3fc620f7 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -5,6 +5,7 @@ # # SPDX-License-Identifier: MIT +import asyncio import io import threading import time @@ -32,6 +33,10 @@ ) +# Seconds a test waits for an event from a helper thread before failing. +EVENT_TIMEOUT = 10 + + def get_query_results_responses(*pages): """Build stubbed ``GetQueryResults`` responses of an integer column ``a``. @@ -128,6 +133,43 @@ def __call__(self, *args, **kwargs): return "success" +class StoppingOperation: + """A stubbed AWS API operation that always throttles and stops its retries. + + The retries are stopped during the first request, or from another thread + during the backoff wait after it. + """ + + def __init__(self, stop, stop_in, operation_name="GetQueryResults"): + """Initialize the operation. + + Args: + stop: Stops the retries; called once. + stop_in: "request" to stop during the first request, "backoff" to + stop from another thread 0.1 seconds into the backoff wait. + operation_name: The name of the operation in its error. + """ + self.stop = stop + self.stop_in = stop_in + self.error = throttling_error(operation_name) + self.calls = 0 + self._stopper = threading.Timer(0.1, stop) + + def __call__(self, *args, **kwargs): + if self.calls == 0: + if self.stop_in == "request": + self.stop() + else: + self._stopper.start() + self.calls += 1 + raise self.error + + def join(self): + """Wait for the thread that stops the retries during the backoff wait, if any.""" + if self._stopper.ident is not None: + self._stopper.join() + + def close_during_next_retries(result_set, close_in): """Throttle the next ``GetQueryResults`` request and close the result set during its retries. @@ -137,26 +179,29 @@ def close_during_next_retries(result_set, close_in): Args: result_set: A result set whose ``GetQueryResults`` responses are stubbed. - close_in: "request" to close during the request, "backoff" to close from - another thread during the backoff wait after it. + close_in: "request" or "backoff"; see :class:`StoppingOperation`. Returns: - The stubbed ``GetQueryResults`` mock, and the thread that closes the - result set during the backoff wait, if any, which the caller joins. + The stubbed ``GetQueryResults`` operation, whose ``join()`` the caller + calls before the end of the test. """ result_set._retry_config = RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT) - get_query_results = result_set.connection.client.get_query_results - closer = threading.Timer(0.1, result_set.close) + operation = StoppingOperation(result_set.close, close_in) + result_set.connection.client.get_query_results.side_effect = operation + return operation + - def throttle(**kwargs): - if close_in == "request": - result_set.close() - else: - closer.start() - raise throttling_error() +async def wait_until(predicate, timeout=EVENT_TIMEOUT): + """Wait, without blocking the event loop, until a condition holds. - get_query_results.side_effect = throttle - return get_query_results, closer + Args: + predicate: The condition. + timeout: The longest wait, in seconds. + """ + deadline = time.monotonic() + timeout + while not predicate(): + assert time.monotonic() < deadline + await asyncio.sleep(0.01) def overlapping_calls(calls, timeout=5): @@ -388,10 +433,6 @@ def wait_for_spark_session_state(client, session_id, state, timeout=120): raise AssertionError(f"Session {session_id} did not become {state} in {timeout} seconds.") -# Seconds a test waits for an event from a helper thread before failing. -EVENT_TIMEOUT = 10 - - class ImmediateExecutor: """An executor that runs each call when it is submitted.""" From dfbf5d8e49cd9d5ef09c2180833d62a502ce8cde Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 07:56:18 +0900 Subject: [PATCH 6/9] Name cleanups of this change as the examples of shield_retries() Co-Authored-By: Claude Opus 5.5 --- pyathena/util.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyathena/util.py b/pyathena/util.py index 1e396e754..02c68023d 100644 --- a/pyathena/util.py +++ b/pyathena/util.py @@ -309,7 +309,8 @@ def shield_retries() -> Iterator[None]: """Keep the events of the enclosing stop scopes from stopping retries in this block. For cleanup that has to run to its end even when the work that it cleans - up was stopped, such as the abort of a multipart upload. A scope that + up was stopped, such as the close of an object whose construction was + cancelled, or a request that stops a query. A scope that :func:`stop_retries_on` opens in the block stops the retries again. Yields: From 090655c79ec7798072747337779c1e04e3103e97 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 08:06:19 +0900 Subject: [PATCH 7/9] Shield synchronous cleanup requests and check close after a fallback page The synchronous StopQueryExecution, StopCalculationExecution, and TerminateSession requests run in shield_retries(), as the asyncio ones do, so that a stop scope of their caller, such as a construction in create_in_thread(), does not stop them. A whole read through the GetQueryResults fallback checks is_closed after each page, so that a close() during a request that succeeds raises ProgrammingError instead of returning the page. The docs state that a second cancellation stops the retries of a shielded start request. Refs #1202 Co-Authored-By: Claude Opus 5.5 --- docs/aio.md | 3 +- pyathena/common.py | 17 +++++++---- pyathena/result_set.py | 3 ++ pyathena/spark/common.py | 35 ++++++++++++++--------- tests/pyathena/aio/test_cursor.py | 27 +++++++++++++++++ tests/pyathena/spark/test_spark_cursor.py | 23 +++++++++++++++ tests/pyathena/test_cursor.py | 16 ++++++++++- tests/pyathena/test_result_set.py | 9 ++++++ 8 files changed, 111 insertions(+), 22 deletions(-) diff --git a/docs/aio.md b/docs/aio.md index acd5b4f15..37a43752a 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -169,7 +169,8 @@ async with await aio_connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/", 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 start request of `execute()` with `kill_on_interrupt` enabled, the requests that stop a query or a calculation, and the request that terminates a Spark session keep their retries. +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)= diff --git a/pyathena/common.py b/pyathena/common.py index ae09d313e..a37c4350b 100644 --- a/pyathena/common.py +++ b/pyathena/common.py @@ -35,6 +35,7 @@ _get_error_code, _is_throttling_error, retry_api_call, + shield_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.shield_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 shield_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 diff --git a/pyathena/result_set.py b/pyathena/result_set.py index 817088629..81b9eadd1 100644 --- a/pyathena/result_set.py +++ b/pyathena/result_set.py @@ -779,6 +779,9 @@ def _iter_all_row_pages(self) -> Iterator[tuple[list[dict[str, Any]], int]]: first_page = True while True: response = self._get_query_results(self.DEFAULT_FETCH_SIZE, next_token) + if self.is_closed: + # close() ran during the request; its page is not returned. + raise ProgrammingError("AthenaResultSet is closed.") rows, next_token = self._parse_result_rows(response) # Only the first page starts with the column labels. offset = 1 if first_page and rows and self._is_first_row_column_labels(rows) else 0 diff --git a/pyathena/spark/common.py b/pyathena/spark/common.py index b3341d645..daf1b815a 100644 --- a/pyathena/spark/common.py +++ b/pyathena/spark/common.py @@ -27,7 +27,7 @@ AthenaQueryExecution, AthenaSessionStatus, ) -from pyathena.util import override, parse_output_location, retry_api_call +from pyathena.util import override, parse_output_location, retry_api_call, shield_retries _logger = logging.getLogger(__name__) @@ -287,7 +287,9 @@ def _terminate_session_by_id(self, session_id: str) -> None: Session startup calls this synchronously in every cursor variant, including those that override ``_terminate_session`` with a coroutine, - so subclasses must not override this method with a coroutine. + so subclasses must not override this method with a coroutine. The + request cleans up, so a stop scope does not stop its retries (see + :func:`~pyathena.util.shield_retries`). Args: session_id: The session ID. @@ -297,12 +299,13 @@ def _terminate_session_by_id(self, session_id: str) -> None: """ request: dict[str, Any] = {"SessionId": session_id} try: - retry_api_call( - self._connection.client.terminate_session, - config=self._retry_config, - logger=_logger, - **request, - ) + with shield_retries(): + retry_api_call( + self._connection.client.terminate_session, + config=self._retry_config, + logger=_logger, + **request, + ) except Exception as e: _logger.exception(f"Failed to terminate session: {session_id}.") raise OperationalError(*e.args) from e @@ -430,6 +433,9 @@ def _start_calculation_execution(self, request: dict[str, Any]) -> str: def _cancel(self, query_id: str) -> None: """Stop a calculation execution with ``StopCalculationExecution``. + The request cleans up, so a stop scope does not stop its retries (see + :func:`~pyathena.util.shield_retries`). + Args: query_id: The calculation execution ID. @@ -438,12 +444,13 @@ def _cancel(self, query_id: str) -> None: """ request: dict[str, Any] = {"CalculationExecutionId": query_id} try: - retry_api_call( - self._connection.client.stop_calculation_execution, - config=self._retry_config, - logger=_logger, - **request, - ) + with shield_retries(): + retry_api_call( + self._connection.client.stop_calculation_execution, + config=self._retry_config, + logger=_logger, + **request, + ) except Exception as e: _logger.exception("Failed to cancel calculation.") raise OperationalError(*e.args) from e diff --git a/tests/pyathena/aio/test_cursor.py b/tests/pyathena/aio/test_cursor.py index c522c5b35..4bb694e89 100644 --- a/tests/pyathena/aio/test_cursor.py +++ b/tests/pyathena/aio/test_cursor.py @@ -402,6 +402,33 @@ def start_query_execution(**kwargs): cancel.assert_awaited_once_with("query_id") assert cursor.query_id == "query_id" + async def test_execute_cancelled_again_while_start_is_retried(self): + """GH-1202: a second cancellation stops the retries of the start request (no AWS).""" + cursor, cancel, _, _ = _starting_cursor() + operation = ThrottledOperation(block=True, failures=1, operation_name="StartQueryExecution") + + def start_query_execution(**kwargs): + operation(**kwargs) + return {"QueryExecutionId": "query_id"} + + cursor._connection.client.start_query_execution.side_effect = start_query_execution + task = asyncio.create_task(cursor.execute("SELECT 1")) + try: + assert await asyncio.to_thread(operation.started.wait, EVENT_TIMEOUT) + task.cancel() + # Let the task handle the cancellation and wait for the start task. + for _ in range(5): + await asyncio.sleep(0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + finally: + operation.release.set() + # A retry without the stop would follow the first attempt at once. + await asyncio.sleep(0.2) + assert operation.calls == 1 + cancel.assert_not_awaited() + async def test_cancel_cancelled_while_retrying(self): """GH-1202: a cancellation does not stop the retries of StopQueryExecution (no AWS).""" cursor, _, _ = _offline_cursor( diff --git a/tests/pyathena/spark/test_spark_cursor.py b/tests/pyathena/spark/test_spark_cursor.py index 39a045d76..cc66789fa 100644 --- a/tests/pyathena/spark/test_spark_cursor.py +++ b/tests/pyathena/spark/test_spark_cursor.py @@ -6,6 +6,7 @@ # SPDX-License-Identifier: MIT import textwrap +import threading import time from concurrent.futures import ThreadPoolExecutor from unittest.mock import MagicMock, patch @@ -15,15 +16,37 @@ from pyathena import NotSupportedError, OperationalError from pyathena.model import AthenaCalculationExecutionStatus, AthenaSessionStatus from pyathena.spark.cursor import SparkCursor +from pyathena.util import RetryConfig, stop_retries_on from tests import ENV from tests.pyathena.util import ( CANCELABLE_SPARK_JOB, + ThrottledOperation, wait_for_spark_job, wait_for_spark_session_state, ) class TestSparkCursor: + @pytest.mark.parametrize( + ("method", "client_method"), + [ + ("_terminate_session_by_id", "terminate_session"), + ("_cancel", "stop_calculation_execution"), + ], + ) + def test_cleanup_in_stopped_scope(self, method, client_method): + """GH-1202: a stop scope does not stop the retries of a cleanup request (no AWS).""" + cursor = SparkCursor.__new__(SparkCursor) # bypass __init__ to avoid AWS calls + cursor._connection = MagicMock() + cursor._retry_config = RetryConfig(attempt=3, multiplier=0) + operation = ThrottledOperation(failures=2) + getattr(cursor._connection.client, client_method).side_effect = operation + stop = threading.Event() + stop.set() + with stop_retries_on(stop): + getattr(cursor, method)("id") + assert operation.calls == 3 + def test_spark_dataframe(self, spark_cursor): spark_cursor.execute( textwrap.dedent( diff --git a/tests/pyathena/test_cursor.py b/tests/pyathena/test_cursor.py index c484786d4..793527f4c 100644 --- a/tests/pyathena/test_cursor.py +++ b/tests/pyathena/test_cursor.py @@ -36,7 +36,7 @@ from pyathena.error import DatabaseError, NotSupportedError, OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.result_set import AthenaDictResultSet -from pyathena.util import RetryConfig +from pyathena.util import RetryConfig, stop_retries_on from tests import ENV from tests.pyathena.conftest import connect from tests.pyathena.tables import TABLES, VIEWS @@ -44,6 +44,7 @@ CONVERTED_VALUES_QUERY, CONVERTED_VALUES_ROW, EVENT_TIMEOUT, + ThrottledOperation, interrupt_start_waits, succeeded_query_execution, throttle_metadata_api, @@ -336,6 +337,19 @@ def execution(schema): == "query_id_other_schema" ) + def test_cancel_in_stopped_scope(self): + """GH-1202: a stop scope does not stop the retries of StopQueryExecution (no AWS).""" + cursor = Cursor.__new__(Cursor) # bypass __init__ to avoid AWS calls + cursor._connection = MagicMock() + cursor._retry_config = RetryConfig(attempt=3, multiplier=0) + operation = ThrottledOperation(failures=2, operation_name="StopQueryExecution") + cursor._connection.client.stop_query_execution.side_effect = operation + stop = threading.Event() + stop.set() + with stop_retries_on(stop): + cursor._cancel("query_id") + assert operation.calls == 3 + def test_cache_size_different_catalog(self): query = "SELECT * FROM one_row" schema = "this_schema" diff --git a/tests/pyathena/test_result_set.py b/tests/pyathena/test_result_set.py index 6d87f73cb..f4acb33ec 100644 --- a/tests/pyathena/test_result_set.py +++ b/tests/pyathena/test_result_set.py @@ -193,6 +193,15 @@ def test_fetch_all_rows_with_close_during_retries(self, close_in): operation.join() assert operation.calls == 1 + @pytest.mark.parametrize("method", ["_fetch_all_rows", "_iter_all_row_pages"]) + def test_fetch_all_rows_with_close_during_request(self, method): + """A whole read does not return the page of a request that close() interrupts.""" + result_set = _fetching_result_set(("1", "2"), ("3",)) + get_query_results = close_during_next_request(result_set) + with pytest.raises(ProgrammingError, match=r"AthenaResultSet is closed\."): + list(getattr(result_set, method)()) + assert get_query_results.call_count == 2 + def test_get_query_results_after_close(self): result_set = _fetching_result_set(("1", "2"), ("3",)) result_set.close() From 6fd371122f5ccf7080c1430a7fda516f6361658e Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 08:22:15 +0900 Subject: [PATCH 8/9] Give the CSV fallback test's bare result set a connection The whole read now checks is_closed after each page, which reads the connection that the test's result set, built without __init__, lacked. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/test_result_set.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/pyathena/test_result_set.py b/tests/pyathena/test_result_set.py index f4acb33ec..240773a87 100644 --- a/tests/pyathena/test_result_set.py +++ b/tests/pyathena/test_result_set.py @@ -241,6 +241,7 @@ def test_fetch_all_rows_as_csv(self): them is data. No AWS calls; GetQueryResults is mocked. """ result_set = AthenaResultSet.__new__(AthenaResultSet) # bypass __init__ + result_set._connection = MagicMock() result_set._query_execution = None result_set._metadata = tuple( {"Name": name, "Type": "varchar", "Precision": 0, "Scale": 0, "Nullable": "UNKNOWN"} From 24b263e0bd6992b00c0d1d91d6e6bae238a32362 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 10:37:27 +0900 Subject: [PATCH 9/9] Rename shield_retries() to unstoppable_retries() "Shield" read as asyncio.shield(), which protects a task from cancellation, and "shield retries" could read as preventing them. The new name states the effect in the block, pairs with stop_retries_on() and the stoppable parameter of async_retry_api_call(), and its docstring now names what still stops the retries: a scope opened in the block, and the stop_when condition of a call. Refs #1202 Co-Authored-By: Claude Opus 5.5 --- docs/api/utilities.rst | 2 +- pyathena/aio/util.py | 12 ++++++------ pyathena/common.py | 6 +++--- pyathena/spark/common.py | 10 +++++----- pyathena/util.py | 17 ++++++++++------- tests/pyathena/test_util.py | 12 ++++++------ 6 files changed, 31 insertions(+), 28 deletions(-) diff --git a/docs/api/utilities.rst b/docs/api/utilities.rst index 93fff21f2..afa1baff8 100644 --- a/docs/api/utilities.rst +++ b/docs/api/utilities.rst @@ -26,7 +26,7 @@ Utility Functions .. autofunction:: pyathena.util.stop_retries_on -.. autofunction:: pyathena.util.shield_retries +.. autofunction:: pyathena.util.unstoppable_retries .. autofunction:: pyathena.util.is_retryable_error diff --git a/pyathena/aio/util.py b/pyathena/aio/util.py index 32740b1d3..5770abec4 100644 --- a/pyathena/aio/util.py +++ b/pyathena/aio/util.py @@ -17,7 +17,7 @@ from collections.abc import Callable from typing import Any, Protocol, TypeVar, cast -from pyathena.util import RetryConfig, retry_api_call, shield_retries, stop_retries_on +from pyathena.util import RetryConfig, retry_api_call, stop_retries_on, unstoppable_retries T = TypeVar("T") @@ -88,9 +88,9 @@ async def async_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 ignore the - stop scopes (see :func:`~pyathena.util.shield_retries`) and run to - their end in the thread, as with ``asyncio.to_thread()``. + 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: @@ -99,7 +99,7 @@ async def async_retry_api_call( def call() -> Any: """Make the call, ignoring the stop scopes unless it is stoppable.""" - with contextlib.nullcontext() if stoppable else shield_retries(): + 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 ) @@ -163,7 +163,7 @@ def construct() -> _ClosableT | None: if not cancelled.is_set(): constructed.append(obj) return obj - with shield_retries(): + with unstoppable_retries(): obj.close() return None diff --git a/pyathena/common.py b/pyathena/common.py index a37c4350b..93974d4cf 100644 --- a/pyathena/common.py +++ b/pyathena/common.py @@ -35,7 +35,7 @@ _get_error_code, _is_throttling_error, retry_api_call, - shield_retries, + unstoppable_retries, ) if TYPE_CHECKING: @@ -1409,7 +1409,7 @@ 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.shield_retries`). + :func:`~pyathena.util.unstoppable_retries`). Args: query_id: The query execution ID. @@ -1419,7 +1419,7 @@ def _cancel(self, query_id: str) -> None: """ request: dict[str, Any] = {"QueryExecutionId": query_id} try: - with shield_retries(): + with unstoppable_retries(): retry_api_call( self._connection.client.stop_query_execution, config=self._retry_config, diff --git a/pyathena/spark/common.py b/pyathena/spark/common.py index daf1b815a..7a49ff932 100644 --- a/pyathena/spark/common.py +++ b/pyathena/spark/common.py @@ -27,7 +27,7 @@ AthenaQueryExecution, AthenaSessionStatus, ) -from pyathena.util import override, parse_output_location, retry_api_call, shield_retries +from pyathena.util import override, parse_output_location, retry_api_call, unstoppable_retries _logger = logging.getLogger(__name__) @@ -289,7 +289,7 @@ def _terminate_session_by_id(self, session_id: str) -> None: including those that override ``_terminate_session`` with a coroutine, so subclasses must not override this method with a coroutine. The request cleans up, so a stop scope does not stop its retries (see - :func:`~pyathena.util.shield_retries`). + :func:`~pyathena.util.unstoppable_retries`). Args: session_id: The session ID. @@ -299,7 +299,7 @@ def _terminate_session_by_id(self, session_id: str) -> None: """ request: dict[str, Any] = {"SessionId": session_id} try: - with shield_retries(): + with unstoppable_retries(): retry_api_call( self._connection.client.terminate_session, config=self._retry_config, @@ -434,7 +434,7 @@ def _cancel(self, query_id: str) -> None: """Stop a calculation execution with ``StopCalculationExecution``. The request cleans up, so a stop scope does not stop its retries (see - :func:`~pyathena.util.shield_retries`). + :func:`~pyathena.util.unstoppable_retries`). Args: query_id: The calculation execution ID. @@ -444,7 +444,7 @@ def _cancel(self, query_id: str) -> None: """ request: dict[str, Any] = {"CalculationExecutionId": query_id} try: - with shield_retries(): + with unstoppable_retries(): retry_api_call( self._connection.client.stop_calculation_execution, config=self._retry_config, diff --git a/pyathena/util.py b/pyathena/util.py index 02c68023d..bb3155727 100644 --- a/pyathena/util.py +++ b/pyathena/util.py @@ -305,13 +305,16 @@ def stop_retries_on(event: threading.Event) -> Iterator[None]: @contextmanager -def shield_retries() -> Iterator[None]: - """Keep the events of the enclosing stop scopes from stopping retries in this block. - - For cleanup that has to run to its end even when the work that it cleans - up was stopped, such as the close of an object whose construction was - cancelled, or a request that stops a query. A scope that - :func:`stop_retries_on` opens in the block stops the retries again. +def unstoppable_retries() -> Iterator[None]: + """Run the retries of AWS API calls in this block to their end. + + The events of the enclosing :func:`stop_retries_on` scopes do not stop + the calls of :func:`retry_api_call` in the block, which retry by their + policy. For cleanup that has to run to its end even when the work that it + cleans up was stopped, such as the close of an object whose construction + was cancelled, or a request that stops a query. A scope that + :func:`stop_retries_on` opens in the block stops the retries again, and so + does the ``stop_when`` condition of a call. Yields: None. diff --git a/tests/pyathena/test_util.py b/tests/pyathena/test_util.py index 307d966ec..c4fb941d0 100644 --- a/tests/pyathena/test_util.py +++ b/tests/pyathena/test_util.py @@ -14,9 +14,9 @@ is_retryable_error, parse_output_location, retry_api_call, - shield_retries, stop_retries_on, strtobool, + unstoppable_retries, ) from tests.pyathena.util import ( EVENT_TIMEOUT, @@ -434,19 +434,19 @@ def test_retry_api_call_stopped_by_enclosing_scope(): assert call.calls == 1 -def test_retry_api_call_shielded(): - """Retries in a shielded block ignore the events of the enclosing scopes.""" +def test_retry_api_call_unstoppable(): + """Retries in an unstoppable block ignore the events of the enclosing scopes.""" event = threading.Event() event.set() operation = ThrottledOperation(failures=2) - with stop_retries_on(event), shield_retries(): + with stop_retries_on(event), unstoppable_retries(): assert retry_api_call(operation, RetryConfig(multiplier=0)) == "success" assert operation.calls == 3 - # A scope opened in the shielded block stops the retries again. + # A scope opened in the unstoppable block stops the retries again. operation = ThrottledOperation(failures=2) with ( stop_retries_on(event), - shield_retries(), + unstoppable_retries(), stop_retries_on(event), pytest.raises(ClientError), ):