diff --git a/docs/aio.md b/docs/aio.md index 9841119b7..37a43752a 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -166,6 +166,13 @@ async with await aio_connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/", print(f"Query timed out: {cursor.query_id}") ``` +Cancelling a task also stops the retries of the Athena API requests that the asyncio cursors make in a thread for it, such as the status polling of `execute()` and the page requests of `AioCursor`. +No further attempt is made, and the wait before the next attempt ends. +A request that has already been sent still runs to its end in its thread. +The requests that stop a query or a calculation and the request that terminates a Spark session keep their retries. +The start request of `execute()` with `kill_on_interrupt` enabled keeps its retries after a cancellation, and stops them when the task is cancelled again. +The fetch methods of `AioPandasCursor`, `AioArrowCursor`, `AioPolarsCursor`, and `AioS3FSCursor` run to their end in their thread. + (aio-dict-cursor)= ## AioDictCursor diff --git a/docs/api/utilities.rst b/docs/api/utilities.rst index bb90d570a..afa1baff8 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.unstoppable_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..219a7598b 100644 --- a/pyathena/aio/common.py +++ b/pyathena/aio/common.py @@ -267,6 +267,7 @@ async def _cancel(self, query_id: str) -> None: # type: ignore[override] self._connection.client.stop_query_execution, config=self._retry_config, logger=_logger, + stoppable=False, **request, ) except Exception as e: 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/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 b52c6d741..5770abec4 100644 --- a/pyathena/aio/util.py +++ b/pyathena/aio/util.py @@ -10,13 +10,16 @@ from __future__ import annotations import asyncio +import contextlib import contextvars import logging import threading from collections.abc import Callable from typing import Any, Protocol, TypeVar, cast -from pyathena.util import RetryConfig, retry_api_call +from pyathena.util import RetryConfig, retry_api_call, stop_retries_on, unstoppable_retries + +T = TypeVar("T") class _Closable(Protocol): @@ -26,18 +29,55 @@ def close(self) -> Any: ... _ClosableT = TypeVar("_ClosableT", bound=_Closable) +async def run_in_thread(func: Callable[..., T], /, *args: Any, **kwargs: Any) -> T: + """Run a blocking function in the default executor, stopping its retries on cancellation. + + Like ``asyncio.to_thread()``, the function runs in the event loop's default + executor with a copy of the caller's context, and a cancellation of the + caller skips a function that has not started. A cancellation also stops + the retries of the AWS API calls that the running function makes (see + :func:`~pyathena.util.stop_retries_on`). A request that has been sent + still runs to its end, and the function returns or raises in its + thread, where nobody awaits it. + + Args: + func: The blocking function. + *args: Positional arguments for ``func``. + **kwargs: Keyword arguments for ``func``. + + Returns: + The return value of ``func``. + """ + stop = threading.Event() + + def call() -> T: + """Run the function in a scope that the cancellation stops.""" + with stop_retries_on(stop): + return func(*args, **kwargs) + + future = asyncio.get_running_loop().run_in_executor(None, contextvars.copy_context().run, call) + try: + return await future + except asyncio.CancelledError: + stop.set() + raise + + async def async_retry_api_call( func: Callable[..., Any], config: RetryConfig, logger: logging.Logger | None = None, *args: Any, stop_on: Callable[[BaseException], bool] | None = None, + stop_when: Callable[[], bool] | None = None, + stoppable: bool = True, **kwargs: Any, ) -> Any: """Execute a function with retry logic in a thread to avoid blocking the event loop. - Wraps ``retry_api_call`` with ``asyncio.to_thread()`` so that blocking - boto3 calls do not block the asyncio event loop. + Runs ``retry_api_call`` with :func:`run_in_thread` so that blocking boto3 + calls do not block the asyncio event loop. Cancelling the caller makes no + further attempt and ends the backoff wait, unless ``stoppable`` is false. Args: func: The AWS API function to call. @@ -45,14 +85,26 @@ 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``. + stoppable: Whether a cancellation of the caller stops the retries. A + request that cleans up, such as one that stops a query or + terminates a session, passes False: its retries then run to their + end in the thread (see :func:`~pyathena.util.unstoppable_retries`), + as with ``asyncio.to_thread()``. **kwargs: Keyword arguments to pass to the function. Returns: The result of the successful function call. """ - return await asyncio.to_thread( - retry_api_call, func, config, logger, *args, stop_on=stop_on, **kwargs - ) + + def call() -> Any: + """Make the call, ignoring the stop scopes unless it is stoppable.""" + with contextlib.nullcontext() if stoppable else unstoppable_retries(): + return retry_api_call( + func, config, logger, *args, stop_on=stop_on, stop_when=stop_when, **kwargs + ) + + return await (run_in_thread(call) if stoppable else asyncio.to_thread(call)) async def close_in_thread(close: Callable[[], Any]) -> None: @@ -76,7 +128,8 @@ async def create_in_thread( Like ``asyncio.to_thread()``, the constructor runs in the event loop's default executor with the caller's context. If the caller is cancelled, a - construction that has not started is skipped, an object constructed after + construction that has not started is skipped, the retries of a running + one are stopped as by :func:`run_in_thread`, an object constructed after the cancellation is closed by the constructing thread, and an object constructed just before it is closed in a new thread: nothing else can close it, and a close by the garbage collector could block the event @@ -96,14 +149,22 @@ async def create_in_thread( constructed: list[_ClosableT] = [] def construct() -> _ClosableT | None: + """Construct the object in a scope that the cancellation stops. + + Returns: + The object, or None if the caller was cancelled, after which an + object constructed here is closed here. + """ if cancelled.is_set(): return None - obj = factory(*args, **kwargs) + with stop_retries_on(cancelled): + obj = factory(*args, **kwargs) with lock: if not cancelled.is_set(): constructed.append(obj) return obj - obj.close() + with unstoppable_retries(): + obj.close() return None future = asyncio.get_running_loop().run_in_executor( diff --git a/pyathena/common.py b/pyathena/common.py index ae09d313e..93974d4cf 100644 --- a/pyathena/common.py +++ b/pyathena/common.py @@ -35,6 +35,7 @@ _get_error_code, _is_throttling_error, retry_api_call, + unstoppable_retries, ) if TYPE_CHECKING: @@ -1407,6 +1408,9 @@ def close(self) -> None: def _cancel(self, query_id: str) -> None: """Stop a query execution with ``StopQueryExecution``. + The request cleans up, so a stop scope does not stop its retries (see + :func:`~pyathena.util.unstoppable_retries`). + Args: query_id: The query execution ID. @@ -1415,12 +1419,13 @@ def _cancel(self, query_id: str) -> None: """ request: dict[str, Any] = {"QueryExecutionId": query_id} try: - retry_api_call( - self._connection.client.stop_query_execution, - config=self._retry_config, - logger=_logger, - **request, - ) + with unstoppable_retries(): + retry_api_call( + self._connection.client.stop_query_execution, + config=self._retry_config, + logger=_logger, + **request, + ) except Exception as e: _logger.exception("Failed to cancel query.") raise OperationalError(*e.args) from e diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 2762b9b86..290c8a9a2 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -10,6 +10,7 @@ from collections.abc import Callable, Iterable, Iterator from concurrent.futures import ThreadPoolExecutor from contextlib import ExitStack, contextmanager +from contextvars import copy_context from functools import partial from io import BufferedReader, BytesIO, IOBase, StringIO, TextIOWrapper from itertools import chain, islice @@ -1625,11 +1626,20 @@ def _read_parquet_footers(self) -> list[FileMetaData]: self.connection.s3_client.meta.config.max_pool_connections, len(self._data_manifest), ) + # Each read runs with a copy of this context, so that a stop of the retries in + # this thread (see stop_retries_on()) stops those of the reads too. + contexts = [copy_context() for _ in self._data_manifest] try: # When the failure of a read reaches map()'s iterator, it cancels the reads # that have not started. with ThreadPoolExecutor(max_workers=workers) as executor: - footers = list(executor.map(self._read_parquet_footer, self._data_manifest)) + footers = list( + executor.map( + lambda context, path: context.run(self._read_parquet_footer, path), + contexts, + self._data_manifest, + ) + ) except Exception as e: _logger.exception(f"Failed to read {self._unload_location}.") raise OperationalError(*e.args) from e diff --git a/pyathena/result_set.py b/pyathena/result_set.py index c0a7cf9e3..81b9eadd1 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). " @@ -763,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 @@ -785,7 +804,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 +835,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/pyathena/spark/common.py b/pyathena/spark/common.py index b3341d645..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 +from pyathena.util import override, parse_output_location, retry_api_call, unstoppable_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.unstoppable_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 unstoppable_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.unstoppable_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 unstoppable_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/pyathena/util.py b/pyathena/util.py index a88dc6bf5..bb3155727 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,69 @@ 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 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. + """ + 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 +348,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 +376,58 @@ 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()) + + # The error of the attempt before a backoff wait, raised if the wait is stopped. + failures: list[BaseException] = [] + # 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( - 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 +440,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/aio/test_cursor.py b/tests/pyathena/aio/test_cursor.py index c20e81747..4bb694e89 100644 --- a/tests/pyathena/aio/test_cursor.py +++ b/tests/pyathena/aio/test_cursor.py @@ -20,8 +20,10 @@ from tests.pyathena.aio.conftest import _aio_connect from tests.pyathena.util import ( EVENT_TIMEOUT, + ThrottledOperation, succeeded_query_execution, throttle_metadata_api, + wait_until, ) @@ -369,6 +371,83 @@ 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_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( + 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() + 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).""" cursor, cancel, started, release = _starting_cursor() diff --git a/tests/pyathena/aio/test_result_set.py b/tests/pyathena/aio/test_result_set.py index dc0292ba8..6b22afcb9 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,22 @@ 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,)] + 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 + operation.join() + assert await result_set.fetchone() is None + assert operation.calls == 1 + 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/aio/test_util.py b/tests/pyathena/aio/test_util.py index 666c5a8b4..35764cb62 100644 --- a/tests/pyathena/aio/test_util.py +++ b/tests/pyathena/aio/test_util.py @@ -8,16 +8,39 @@ 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, stop_retries_on +from tests.pyathena.util import EVENT_TIMEOUT, ThrottledOperation, wait_until _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. + 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) + 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 +203,133 @@ 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_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") + 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) + 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) 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 83df71187..240773a87 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,50 @@ 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,)] + operation = close_during_next_retries(result_set, close_in) + started = time.monotonic() + assert getattr(result_set, method)() == expected + assert time.monotonic() - started < EVENT_TIMEOUT + operation.join() + assert result_set.fetchone() is None + assert operation.calls == 1 + + @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",)) + 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 + 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() @@ -191,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"} diff --git a/tests/pyathena/test_util.py b/tests/pyathena/test_util.py index 1c5eed6fa..c4fb941d0 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,7 +14,15 @@ is_retryable_error, parse_output_location, retry_api_call, + stop_retries_on, strtobool, + unstoppable_retries, +) +from tests.pyathena.util import ( + EVENT_TIMEOUT, + StoppingOperation, + ThrottledOperation, + throttling_error, ) @@ -202,13 +212,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 @@ -235,7 +238,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(): @@ -255,7 +258,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 @@ -372,3 +375,90 @@ def call(): assert caught.value is error assert calls == expected_calls + + +@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 = 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)) + 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 = StoppingOperation(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 = StoppingOperation(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 = StoppingOperation(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_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), unstoppable_retries(): + assert retry_api_call(operation, RetryConfig(multiplier=0)) == "success" + assert operation.calls == 3 + # A scope opened in the unstoppable block stops the retries again. + operation = ThrottledOperation(failures=2) + with ( + stop_retries_on(event), + unstoppable_retries(), + stop_retries_on(event), + pytest.raises(ClientError), + ): + retry_api_call(operation, RetryConfig(multiplier=0)) + assert operation.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() + 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 cf227edb4..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 @@ -25,12 +26,17 @@ 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") ) +# 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``. @@ -79,6 +85,125 @@ 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" + + +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. + + 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" or "backoff"; see :class:`StoppingOperation`. + + Returns: + 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) + operation = StoppingOperation(result_set.close, close_in) + result_set.connection.client.get_query_results.side_effect = operation + return operation + + +async def wait_until(predicate, timeout=EVENT_TIMEOUT): + """Wait, without blocking the event loop, until a condition holds. + + 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): """Return a callable that records the peak number of its overlapping calls. @@ -308,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."""