From 0ebbd4a0ea2c9c677ee700d3b0d1cfcc708e380c Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 17:33:10 +0900 Subject: [PATCH 1/7] Break result sets after a failed fetch and drop the cross-thread guards After a fetch raises, including for an interrupt or a value that fails to convert, later fetches raise ProgrammingError without reading further rows (#1225 option (c)). Before, the pandas and Polars result sets returned no more rows, and the others continued after dropped or skipped rows. The pandas and Polars chunk iterators do the same for their reads, so a failed iter_chunks() read also breaks the fetches that share the iterator. The fetches of AioPandasCursor, AioArrowCursor, AioPolarsCursor, and AioS3FSCursor stop their retries when cancelled and detach the result set, which is closed in a thread once the fetch has ended in its thread. A later fetch therefore raises instead of running alongside the leftover thread, and leaving async with no longer closes the result set during that fetch. PyAthena declares threadsafety = 2, so using one cursor or result set from several threads, including close() and cancel() from another thread, is unsupported. The code that only served that case is removed: the deferred chunk reader close of the pandas and Polars iterators, the is_closed checks after a page request, the buffer replacement and IndexError handling of the bulk take, and the stop of the GetQueryResults retries on close(). The documentation states the limitation, and the cancellation example that called cancel() from a timer thread is removed. Closes #1206 Closes #1225 Co-Authored-By: Claude Opus 5.5 --- docs/aio.md | 6 +- docs/usage.md | 93 +--------- pyathena/aio/common.py | 90 ++++++++-- pyathena/aio/result_set.py | 76 ++++---- pyathena/pandas/result_set.py | 125 ++++++------- pyathena/polars/result_set.py | 90 ++++------ pyathena/result_set.py | 101 +++++------ tests/pyathena/aio/test_common.py | 51 +++++- tests/pyathena/aio/test_result_set.py | 98 ++++++----- tests/pyathena/pandas/test_result_set.py | 157 +++++------------ tests/pyathena/polars/test_result_set.py | 213 +++++------------------ tests/pyathena/test_result_set.py | 136 +++------------ tests/pyathena/util.py | 51 +----- 13 files changed, 477 insertions(+), 810 deletions(-) diff --git a/docs/aio.md b/docs/aio.md index 57d7e66dd..c4f8e9b9a 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -171,7 +171,9 @@ 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. +After a fetch is cancelled, later fetches of its result raise `ProgrammingError`. +The fetch methods of `AioPandasCursor`, `AioArrowCursor`, `AioPolarsCursor`, and `AioS3FSCursor` run in a thread, which keeps running after a cancellation. +A cancellation also stops the retries of their requests, and the cursor drops the result, which is closed once the fetch ends in its thread. (aio-dict-cursor)= @@ -227,7 +229,7 @@ All aio cursors use `await` for fetch operations, so fetching does not block the Parquet) inside `execute()`, wrapped in `asyncio.to_thread()`. When CSV results are read in chunks (`chunksize` for pandas and Polars, or a chunk size chosen by `auto_optimize_chunksize` for pandas), or with `chunksize` on Polars UNLOAD results, fetch calls read S3 lazily instead. - The fetch methods are also wrapped in `asyncio.to_thread()`. + The fetch methods also run in a thread. - `AioS3FSCursor` streams rows from the result file in S3 as they are fetched. With managed query result storage, where the query has no S3 output location, the pandas, Arrow, diff --git a/docs/usage.md b/docs/usage.md index 118bc3901..a9f1d9986 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -130,6 +130,12 @@ for row in cursor: print(row) ``` +After a fetch raises an exception, including `KeyboardInterrupt`, later fetches of the same result raise `ProgrammingError`. +Execute the query again to read its result. + +As `pyathena.threadsafety = 2` declares, threads can share connections, but not cursors. +Calling the methods of one cursor or its result set from several threads, including `close()` and `cancel()`, is not supported. + (usage-query-with-parameters)= ## Query with parameters @@ -399,7 +405,7 @@ Result columns from a passthrough query come back with the source system's types PyAthena provides a callback mechanism that allows you to get immediate access to the query ID as soon as the `start_query_execution` API call is made, before waiting for query completion. -This is useful for monitoring, logging, or cancelling long-running queries from another thread. +This is useful for monitoring, logging, or cancelling long-running queries by their ID. When `cache_size` finds a reusable query, no new query starts, and the callback receives the reused query's ID. The `on_start_query_execution` callback can be configured at both the connection level and @@ -446,91 +452,6 @@ cursor.execute( ) ``` -### Query cancellation example - -A common use case is to cancel long-running analytical queries after a timeout: - -```python -import threading -from pyathena import connect - -def cancel_long_running_query(): - """Example: Cancel a complex analytical query after 10 minutes.""" - timeout_minutes = 10 - - def track_query_start(query_id): - print(f"Long-running analysis started: {query_id}") - - def cancel_on_timeout(cursor): - """Cancel the query that is still running when the timer fires.""" - try: - cursor.cancel() - print(f"Query cancelled after {timeout_minutes} minutes timeout") - except Exception as e: - print(f"Cancellation failed: {e}") - - cursor = connect( - s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/", - region_name="us-west-2", - on_start_query_execution=track_query_start - ).cursor() - - # Complex analytical query that might run for a long time - long_query = """ - WITH daily_metrics AS ( - SELECT - date_trunc('day', timestamp_col) as day, - user_id, - COUNT(*) as events, - AVG(duration) as avg_duration - FROM large_events_table - WHERE timestamp_col >= current_date - interval '1' year - GROUP BY 1, 2 - ), - user_segments AS ( - SELECT - user_id, - CASE - WHEN AVG(events) > 100 THEN 'high_activity' - WHEN AVG(events) > 10 THEN 'medium_activity' - ELSE 'low_activity' - END as segment - FROM daily_metrics - GROUP BY user_id - ) - SELECT - segment, - COUNT(DISTINCT dm.user_id) as users, - AVG(events) as avg_daily_events - FROM daily_metrics dm - JOIN user_segments us ON dm.user_id = us.user_id - GROUP BY segment - ORDER BY avg_daily_events DESC - """ - - # Cancel the query if it is still running after the timeout - timer = threading.Timer(timeout_minutes * 60, cancel_on_timeout, args=(cursor,)) - timer.start() - try: - print("Starting complex analytical query (10-minute timeout)...") - cursor.execute(long_query) - - # Process results - results = cursor.fetchall() - print(f"Analysis completed successfully: {len(results)} segments found") - for row in results: - print(f" {row[0]}: {row[1]} users, {row[2]:.1f} avg events") - - except Exception as e: - print(f"Query failed or was cancelled: {e}") - finally: - # Stop the timer once the query has finished - timer.cancel() - -# Run the example -cancel_long_running_query() -``` - ### Multiple callbacks When both connection-level and execute-level callbacks are specified, diff --git a/pyathena/aio/common.py b/pyathena/aio/common.py index 219a7598b..f29db358f 100644 --- a/pyathena/aio/common.py +++ b/pyathena/aio/common.py @@ -3,7 +3,10 @@ from __future__ import annotations import asyncio +import contextvars import logging +import operator +import threading from collections.abc import Awaitable, Callable, Coroutine from typing import Any, NoReturn, TypeVar, cast @@ -16,13 +19,25 @@ from pyathena.model import AthenaDatabase, AthenaQueryExecution, AthenaTableMetadata from pyathena.options import ExecuteOptions from pyathena.result_set import AthenaResultSet, WithResultSet -from pyathena.util import _is_throttling_error, override +from pyathena.util import _is_throttling_error, override, stop_retries_on _logger = logging.getLogger(__name__) _T = TypeVar("_T") +def _close_result_set(result_set: AthenaResultSet) -> None: + """Close a result set that a cancelled fetch detached, logging an error. + + Args: + result_set: The result set. + """ + try: + result_set.close() + except Exception: + _logger.exception("Failed to close the result set of a cancelled fetch.") + + class AioBaseCursor(BaseCursor): """Async base cursor that overrides I/O methods with async equivalents. @@ -709,8 +724,8 @@ 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 - whose result set fetches asynchronously overrides them. Closing the + the result set's synchronous fetch in a thread (see ``_fetch_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 a result stream can wait for a read in progress. Synchronous iteration @@ -793,8 +808,8 @@ 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 in a thread to avoid blocking the event + loop (see ``_fetch_in_thread()``). Returns: A tuple representing the next row, or None if no more rows. @@ -802,10 +817,7 @@ async def fetchone( Raises: ProgrammingError: If no result set is available. """ - 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 self._fetch_in_thread(operator.methodcaller("fetchone")) @override async def fetchmany( @@ -813,8 +825,8 @@ 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 in a thread to avoid blocking the event + loop (see ``_fetch_in_thread()``). Args: size: Maximum number of rows to fetch. If None or not positive, @@ -826,10 +838,7 @@ async def fetchmany( Raises: ProgrammingError: If no result set is available. """ - 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 self._fetch_in_thread(operator.methodcaller("fetchmany", size)) @override async def fetchall( @@ -837,19 +846,64 @@ 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 in a thread to avoid blocking the event + loop (see ``_fetch_in_thread()``). Returns: List of tuples representing all remaining rows. + Raises: + ProgrammingError: If no result set is available. + """ + return await self._fetch_in_thread(operator.methodcaller("fetchall")) + + async def _fetch_in_thread(self, fetch: Callable[[AthenaResultSet], _T]) -> _T: + """Run a synchronous fetch of the result set in a thread. + + A cancellation stops the retries of the AWS API calls that the fetch makes + (see :func:`~pyathena.util.stop_retries_on`) and detaches the result set from + the cursor, so that later fetches raise ``ProgrammingError``. The fetch can + still be running in its thread, so the result set is closed in a thread once + the fetch has ended there. + + Args: + fetch: The function that fetches from the result set. + + Returns: + The return value of ``fetch``. + Raises: ProgrammingError: If no result set is available. """ 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) + stop = threading.Event() + + def call() -> _T: + """Run the fetch in a scope that the cancellation stops.""" + with stop_retries_on(stop): + return fetch(result_set) + + def close_after_fetch(_: asyncio.Future[_T]) -> None: + """Close the detached result set in a thread once the fetch has ended.""" + try: + loop.run_in_executor(None, _close_result_set, result_set) + except RuntimeError: + # The default executor is shut down, as at the end of asyncio.run(). + _close_result_set(result_set) + + loop = asyncio.get_running_loop() + future = loop.run_in_executor(None, contextvars.copy_context().run, call) + try: + # The shield keeps the future pending until the fetch has ended. + return await asyncio.shield(future) + except asyncio.CancelledError: + stop.set() + if self.result_set is result_set: + self.result_set = None + future.add_done_callback(close_after_fetch) + raise @override def __iter__(self) -> NoReturn: diff --git a/pyathena/aio/result_set.py b/pyathena/aio/result_set.py index f07e87420..40ad51e9b 100644 --- a/pyathena/aio/result_set.py +++ b/pyathena/aio/result_set.py @@ -23,6 +23,8 @@ _logger = logging.getLogger(__name__) +_FETCH_FAILED_MESSAGE = "AthenaAioResultSet cannot be fetched after a failed fetch." + class AthenaAioResultSet(AthenaResultSet): """Async result set that provides async fetch methods. @@ -123,8 +125,7 @@ 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, also during the - request, which then makes no further attempt. + succeeded, or the result set is closed. OperationalError: If the request fails. """ if self.is_closed: @@ -135,13 +136,9 @@ 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: @@ -150,9 +147,7 @@ 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 or the result set is closed. A - ``close()`` during the request stops its retries, and the page is - discarded. + Nothing is fetched when no page remains, as after ``close()``. Raises: ProgrammingError: If the next page is requested and the query ID is @@ -161,20 +156,9 @@ async def _async_fetch(self) -> None: """ if not self._next_token: return - 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: - self._next_token = None - self._rows.clear() + 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) async def _async_pre_fetch(self) -> None: """Fetch the first page of rows along with the result metadata. @@ -203,16 +187,25 @@ async def fetchone( # type: ignore[override] Returns: The next row (a tuple, or a dict for ``AthenaAioDictResultSet``), or None if no more rows. + + Raises: + ProgrammingError: If an earlier fetch failed or was cancelled. """ - if not self._rows: - await self._async_fetch() + if self._fetch_failed: + raise ProgrammingError(_FETCH_FAILED_MESSAGE) + try: if not self._rows: - return None - row = self._rows.popleft() - # As _count_rows() does, inline on the per-row path. - rownumber = self._rownumber - if rownumber is not None: - self._rownumber = rownumber + 1 + await self._async_fetch() + if not self._rows: + return None + row = self._rows.popleft() + # As _count_rows() does, inline on the per-row path. + rownumber = self._rownumber + if rownumber is not None: + self._rownumber = rownumber + 1 + except BaseException: + self._fetch_failed = True + raise return row @override @@ -230,10 +223,18 @@ async def fetchmany( # type: ignore[override] Raises: TypeError: If ``size`` is not an integer and not replaced with ``arraysize``. + ProgrammingError: If an earlier fetch failed or was cancelled. """ if not size or size <= 0: size = self._arraysize - return await self._async_take_rows(operator.index(size)) + size = operator.index(size) + if self._fetch_failed: + raise ProgrammingError(_FETCH_FAILED_MESSAGE) + try: + return await self._async_take_rows(size) + except BaseException: + self._fetch_failed = True + raise @override async def fetchall( # type: ignore[override] @@ -243,8 +244,17 @@ async def fetchall( # type: ignore[override] Returns: The remaining rows. + + Raises: + ProgrammingError: If an earlier fetch failed or was cancelled. """ - return await self._async_take_rows(None) + if self._fetch_failed: + raise ProgrammingError(_FETCH_FAILED_MESSAGE) + try: + return await self._async_take_rows(None) + except BaseException: + self._fetch_failed = True + raise async def _async_take_rows( self, size: int | None diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index ce0cdb12f..4a6e92a89 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -6,11 +6,10 @@ import importlib import logging import sys -import threading from collections import abc from collections.abc import Callable, Iterable, Iterator from concurrent.futures import ThreadPoolExecutor -from contextlib import ExitStack, contextmanager +from contextlib import ExitStack from contextvars import copy_context from functools import partial from io import BufferedReader, BytesIO, IOBase, StringIO, TextIOWrapper @@ -34,7 +33,7 @@ from pyathena.filesystem.s3 import S3FileSystem from pyathena.model import AthenaQueryExecution from pyathena.pandas.reader import _BINARY_NULL, BinaryCSVReader -from pyathena.result_set import AthenaResultSet +from pyathena.result_set import _FETCH_FAILED_MESSAGE, AthenaResultSet from pyathena.util import RetryConfig, override, parse_output_location if TYPE_CHECKING: @@ -47,6 +46,8 @@ _logger = logging.getLogger(__name__) +_READ_FAILED_MESSAGE = "PandasDataFrameIterator cannot be read after a failed read." + def _convert_binary_csv(converter: Callable[[str | None], Any], value: str) -> Any: return converter(None if value == _BINARY_NULL else value) @@ -337,31 +338,8 @@ def __init__( self._reader = reader self._trunc_date = trunc_date self._csv_stream = csv_stream - self._lock = threading.Lock() - self._readers = 0 - self._close_requested = False - - @contextmanager - def _read(self) -> Iterator[TextFileReader | Iterator[DataFrame]]: - """Track a chunk read and run a close that was requested during it. - - Yields: - The reader to read the chunk from. - """ - with self._lock: - self._readers += 1 - reader = iter(()) if self._close_requested else self._reader - try: - yield reader - finally: - with self._lock: - self._readers -= 1 - close = self._close_requested and not self._readers - if close: - try: - self.close() - except Exception: - _logger.exception("Failed to close the chunk reader after a read.") + # Whether a read raised, after which the position of the next chunk is unknown. + self._read_failed = False @override def __next__(self) -> DataFrame: @@ -372,12 +350,35 @@ def __next__(self) -> DataFrame: Raises: StopIteration: When no more chunks are available. + ProgrammingError: If an earlier read failed. """ + return self._read_chunk(next) + + def _read_chunk(self, read: Callable[[Any], DataFrame]) -> DataFrame: + """Read the next chunk with date truncation applied. + + The iterator is closed at the end of the chunks and after an exception, and + later reads after an exception raise ``ProgrammingError``. + + Args: + read: The function that reads the chunk from the reader. + + Returns: + The chunk. + + Raises: + StopIteration: When no more chunks are available. + ProgrammingError: If an earlier read failed. + """ + if self._read_failed: + raise ProgrammingError(_READ_FAILED_MESSAGE) try: - with self._read() as reader: - df = next(reader) - return self._trunc_date(df) + return self._trunc_date(read(self._reader)) + except StopIteration: + self.close() + raise except BaseException: + self._read_failed = True self.close() raise @@ -395,21 +396,11 @@ def __exit__(self, exc_type, exc_value, traceback) -> None: self.close() def close(self) -> None: - """Close the iterator and release resources. - - If another thread is reading a chunk, the read returns that chunk, and the - reader and the CSV stream are closed when it completes. Reads that start after the - close get no chunks, and an error from the deferred close is logged. - """ + """Close the iterator and release resources; later reads get no chunks.""" from pandas.io.parsers import TextFileReader - with self._lock: - if self._readers: - self._close_requested = True - return - self._close_requested = False - reader, self._reader = self._reader, iter(()) - csv_stream, self._csv_stream = self._csv_stream, None + reader, self._reader = self._reader, iter(()) + csv_stream, self._csv_stream = self._csv_stream, None try: if isinstance(reader, (TextFileReader, GeneratorType)): reader.close() @@ -498,16 +489,18 @@ def get_chunk(self, size: int | None = None) -> DataFrame: Returns: DataFrame chunk, with date truncation applied as in iteration. + + Raises: + StopIteration: When no more chunks are available. + ProgrammingError: If an earlier read failed. """ from pandas.io.parsers import TextFileReader - try: - with self._read() as reader: - df = reader.get_chunk(size) if isinstance(reader, TextFileReader) else next(reader) - return self._trunc_date(df) - except BaseException: - self.close() - raise + return self._read_chunk( + lambda reader: reader.get_chunk(size) + if isinstance(reader, TextFileReader) + else next(reader) + ) def as_pandas(self) -> DataFrame: """Collect all remaining chunks into a single DataFrame. @@ -981,25 +974,25 @@ def _trunc_date(self, df: DataFrame) -> DataFrame: def fetchone( self, ) -> tuple[Any | None, ...] | dict[Any, Any | None] | None: - row = next(self._row_iter, None) + if self._fetch_failed: + raise ProgrammingError(_FETCH_FAILED_MESSAGE) + try: + row = next(self._row_iter, None) + except BaseException: + self._fetch_failed = True + raise if row is not None: self._count_rows(1) return row @override def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]: - """Take the next rows from the row iterator in bulk. - - The rows taken before an exception are counted in ``rownumber`` too. - """ - rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = [] - try: - # extend() keeps the rows that it appended before an exception, and - # islice() takes no size above sys.maxsize. - rows.extend(islice(self._row_iter, None if size is None else min(size, sys.maxsize))) - finally: - if rows: - self._count_rows(len(rows)) + """Take the next rows from the row iterator in bulk.""" + # islice() takes no size above sys.maxsize. + rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = list( + islice(self._row_iter, None if size is None else min(size, sys.maxsize)) + ) + self._count_rows(len(rows)) return rows def _read_csv(self) -> TextFileReader | DataFrame: @@ -1569,9 +1562,7 @@ def _iter_parquet_chunks(self, paths: list[str]) -> Iterator[DataFrame]: Args: paths: The files of the data manifest, which ``_read_parquet_footers()`` - read. They are given when the iterator is created, as ``close()`` from - another thread can reset the result set's data manifest before the - first read. + read. Yields: The DataFrames of the result rows. diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 451e5fc6b..94f09095f 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -13,10 +13,8 @@ import importlib.util import logging import sys -import threading from collections import abc from collections.abc import Callable, Iterable, Iterator -from contextlib import contextmanager from io import BytesIO, StringIO from itertools import islice, repeat from multiprocessing import cpu_count @@ -38,7 +36,7 @@ from pyathena.model import AthenaQueryExecution from pyathena.polars.converter import _to_datetimes from pyathena.polars.util import to_column_info, to_microseconds -from pyathena.result_set import AthenaResultSet +from pyathena.result_set import _FETCH_FAILED_MESSAGE, AthenaResultSet from pyathena.util import RetryConfig, override if TYPE_CHECKING: @@ -55,6 +53,7 @@ "later, or install pyarrow and give no read options to execute(), to read them " "correctly." ) +_READ_FAILED_MESSAGE = "PolarsDataFrameIterator cannot be read after a failed read." class PolarsDataFrameIterator(abc.Iterator): # type: ignore[type-arg] @@ -106,31 +105,8 @@ def __init__( self._reader = reader self._converters = converters self._column_names = column_names - self._lock = threading.Lock() - self._readers = 0 - self._close_requested = False - - @contextmanager - def _read(self) -> Iterator[Iterator[pl.DataFrame]]: - """Track a chunk read and run a close that was requested during it. - - Yields: - The reader to read the chunk from. - """ - with self._lock: - self._readers += 1 - reader = iter(()) if self._close_requested else self._reader - try: - yield reader - finally: - with self._lock: - self._readers -= 1 - close = self._close_requested and not self._readers - if close: - try: - self.close() - except Exception: - _logger.exception("Failed to close the chunk reader after a read.") + # Whether a read raised, after which the position of the next chunk is unknown. + self._read_failed = False @override def __next__(self) -> pl.DataFrame: @@ -139,15 +115,24 @@ def __next__(self) -> pl.DataFrame: Returns: The next Polars DataFrame chunk. + The iterator is closed at the end of the chunks and after an exception, and + later reads after an exception raise ``ProgrammingError``. + Raises: StopIteration: When no more chunks are available. + ProgrammingError: If an earlier read failed. """ + if self._read_failed: + raise ProgrammingError(_READ_FAILED_MESSAGE) try: - with self._read() as reader: - return next(reader) + return next(self._reader) except StopIteration: self.close() raise + except BaseException: + self._read_failed = True + self.close() + raise @override def __iter__(self) -> PolarsDataFrameIterator: @@ -163,18 +148,8 @@ def __exit__(self, exc_type, exc_value, traceback) -> None: self.close() def close(self) -> None: - """Close the iterator and release resources. - - If another thread is reading a chunk, the read returns that chunk, and the - reader is closed when it completes. Reads that start after the - close get no chunks, and an error from the deferred close is logged. - """ - with self._lock: - if self._readers: - self._close_requested = True - return - self._close_requested = False - reader, self._reader = self._reader, iter(()) + """Close the iterator and release resources; later reads get no chunks.""" + reader, self._reader = self._reader, iter(()) if isinstance(reader, GeneratorType): reader.close() @@ -632,9 +607,8 @@ def _create_dataframe_iterator(self) -> PolarsDataFrameIterator: files, as ``_as_polars()`` does, before the chunks are read. The chunk readers take the state that they need, such as the output location, - the column names, and the connection, when they are created, so that a - ``close()`` from another thread before their first read does not take it away. - The credentials are still frozen at the first read. + the column names, and the connection, when they are created. The credentials + are frozen at the first read. Returns: PolarsDataFrameIterator that reads each chunk lazily. @@ -656,25 +630,25 @@ def _create_dataframe_iterator(self) -> PolarsDataFrameIterator: def fetchone( self, ) -> tuple[Any | None, ...] | dict[Any, Any | None] | None: - row = next(self._row_iter, None) + if self._fetch_failed: + raise ProgrammingError(_FETCH_FAILED_MESSAGE) + try: + row = next(self._row_iter, None) + except BaseException: + self._fetch_failed = True + raise if row is not None: self._count_rows(1) return row @override def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]: - """Take the next rows from the row iterator in bulk. - - The rows taken before an exception are counted in ``rownumber`` too. - """ - rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = [] - try: - # extend() keeps the rows that it appended before an exception, and - # islice() takes no size above sys.maxsize. - rows.extend(islice(self._row_iter, None if size is None else min(size, sys.maxsize))) - finally: - if rows: - self._count_rows(len(rows)) + """Take the next rows from the row iterator in bulk.""" + # islice() takes no size above sys.maxsize. + rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = list( + islice(self._row_iter, None if size is None else min(size, sys.maxsize)) + ) + self._count_rows(len(rows)) return rows def _is_csv_readable(self) -> bool: diff --git a/pyathena/result_set.py b/pyathena/result_set.py index 4328bd36c..8d402a687 100644 --- a/pyathena/result_set.py +++ b/pyathena/result_set.py @@ -26,6 +26,8 @@ _logger = logging.getLogger(__name__) +_FETCH_FAILED_MESSAGE = "AthenaResultSet cannot be fetched after a failed fetch." + class AthenaResultSet(CursorIterator): """Result set for Athena query execution using the GetQueryResults API. @@ -115,6 +117,8 @@ def __init__( collections.deque() ) self._next_token: str | None = None + # Whether a fetch raised, after which the position of the next row is unknown. + self._fetch_failed = False if self.state == AthenaQueryExecution.STATE_SUCCEEDED: self._rownumber = 0 @@ -438,8 +442,7 @@ 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, also during the - request, which then makes no further attempt. + succeeded, or the result set is closed. OperationalError: If the request fails. """ if self.is_closed: @@ -450,13 +453,9 @@ 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: @@ -466,8 +465,7 @@ 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 or the result set is closed. A ``close()`` during the request - stops its retries, and the page is discarded. + page remains, as after ``close()``. Raises: ProgrammingError: If the next page is requested and the query ID is @@ -476,20 +474,9 @@ def _fetch(self) -> None: """ if not self._next_token: return - 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: - self._next_token = None - self._rows.clear() + response = self._get_query_results(self._arraysize, self._next_token) + rows, self._next_token = self._parse_result_rows(response) + self._process_rows(rows) def _pre_fetch(self) -> None: """Fetch the first page of rows along with the result metadata. @@ -510,15 +497,21 @@ def _pre_fetch(self) -> None: def fetchone( self, ) -> tuple[Any | None, ...] | dict[Any, Any | None] | None: - if not self._rows: - self._fetch() + if self._fetch_failed: + raise ProgrammingError(_FETCH_FAILED_MESSAGE) + try: if not self._rows: - return None - row = self._rows.popleft() - # As _count_rows() does, inline on the per-row path. - rownumber = self._rownumber - if rownumber is not None: - self._rownumber = rownumber + 1 + self._fetch() + if not self._rows: + return None + row = self._rows.popleft() + # As _count_rows() does, inline on the per-row path. + rownumber = self._rownumber + if rownumber is not None: + self._rownumber = rownumber + 1 + except BaseException: + self._fetch_failed = True + raise return row @override @@ -536,10 +529,18 @@ def fetchmany( Raises: TypeError: If ``size`` is not an integer and not replaced with ``arraysize``. + ProgrammingError: If an earlier fetch failed. """ if not size or size <= 0: size = self._arraysize - return self._take_rows(operator.index(size)) + size = operator.index(size) + if self._fetch_failed: + raise ProgrammingError(_FETCH_FAILED_MESSAGE) + try: + return self._take_rows(size) + except BaseException: + self._fetch_failed = True + raise @override def fetchall( @@ -549,8 +550,17 @@ def fetchall( Returns: The remaining rows. + + Raises: + ProgrammingError: If an earlier fetch failed. """ - return self._take_rows(None) + if self._fetch_failed: + raise ProgrammingError(_FETCH_FAILED_MESSAGE) + try: + return self._take_rows(None) + except BaseException: + self._fetch_failed = True + raise def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]: """Take the next rows and count them in ``rownumber``. @@ -586,26 +596,15 @@ def _take_buffered_rows( start = len(rows) limit = None if size is None else size - start if limit is not None and limit < len(self._rows): - # A row is never None, so popleft() runs until the limit or an empty - # buffer, and extend() keeps the rows that it appended before that. - # try costs nothing without an exception, unlike contextlib.suppress(). - try: # noqa: SIM105 - rows.extend(islice(iter(self._rows.popleft, None), limit)) - except IndexError: - pass + # A row is never None, so popleft() runs until the limit. + rows.extend(islice(iter(self._rows.popleft, None), limit)) else: - # The buffer is replaced, not cleared after it is read, so that a close() - # in between empties only the new buffer. - buffered, self._rows = self._rows, collections.deque() - rows.extend(buffered) + rows.extend(self._rows) + self._rows.clear() self._count_rows(len(rows) - start) def _count_rows(self, count: int) -> None: - """Count the rows that a fetch took in ``rownumber``. - - ``rownumber`` stays None for a result set that another thread closed before - the update. It is read once, so that a ``close()`` during the update raises no - TypeError, though it may leave the count in place. + """Count the rows that a fetch took in ``rownumber``, unless it is None. Args: count: The number of rows taken. @@ -763,8 +762,7 @@ def _iter_all_row_pages(self) -> Iterator[tuple[list[dict[str, Any]], int]]: 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. + ProgrammingError: If the result set is closed. OperationalError: If a request fails. """ _logger.warning( @@ -777,9 +775,6 @@ 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/tests/pyathena/aio/test_common.py b/tests/pyathena/aio/test_common.py index aabe6f675..da9816739 100644 --- a/tests/pyathena/aio/test_common.py +++ b/tests/pyathena/aio/test_common.py @@ -17,7 +17,8 @@ from pyathena.aio.polars.cursor import AioPolarsCursor from pyathena.aio.s3fs.cursor import AioS3FSCursor from pyathena.error import OperationalError, ProgrammingError -from tests.pyathena.util import EVENT_TIMEOUT +from pyathena.util import RetryConfig +from tests.pyathena.util import EVENT_TIMEOUT, ThrottledOperation, retried _ALL_CURSOR_CLASSES = [ AioCursor, @@ -105,6 +106,54 @@ async def test_fetch_without_result_set(self, cursor_class): with pytest.raises(ProgrammingError, match=r"No result set\."): await fetch() + @pytest.mark.parametrize( + ("method", "args"), [("fetchone", ()), ("fetchmany", (2,)), ("fetchall", ())] + ) + @pytest.mark.parametrize( + "cursor_class", + [AioArrowCursor, AioPandasCursor, AioPolarsCursor, AioS3FSCursor], + ) + async def test_cancelled_fetch_detaches_result_set(self, cursor_class, method, args): + """GH-1206: a cancelled fetch stops its retries and detaches the result set. + + The result set is closed in a thread once the fetch has ended in its thread, + and later fetches raise ProgrammingError. Without the stop, the retries would + wait at least EVENT_TIMEOUT seconds before the next attempt. + """ + cursor = cursor_class( + connection=MagicMock(), converter=None, formatter=None, retry_config=None + ) + operation = ThrottledOperation(block=True) + closes = [] + closed = threading.Event() + + def close(): + closes.append(threading.get_ident()) + closed.set() + + result_set = MagicMock(is_closed=False) + result_set.close.side_effect = close + getattr(result_set, method).side_effect = retried( + operation, RetryConfig(attempt=2, multiplier=EVENT_TIMEOUT) + ) + cursor.result_set = result_set + task = asyncio.create_task(getattr(cursor, method)(*args)) + try: + assert await asyncio.to_thread(operation.started.wait, EVENT_TIMEOUT) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert cursor.result_set is None + for fetch in (cursor.fetchone, cursor.fetchmany, cursor.fetchall): + with pytest.raises(ProgrammingError, match=r"No result set\."): + await fetch() + assert not closes + finally: + operation.release.set() + assert await asyncio.to_thread(closed.wait, EVENT_TIMEOUT) + assert operation.calls == 1 + assert closes != [threading.get_ident()] + @pytest.mark.parametrize("cursor_class", _ALL_CURSOR_CLASSES) async def test_async_iteration(self, cursor_class): cursor = cursor_class( diff --git a/tests/pyathena/aio/test_result_set.py b/tests/pyathena/aio/test_result_set.py index 6b22afcb9..793c55c20 100644 --- a/tests/pyathena/aio/test_result_set.py +++ b/tests/pyathena/aio/test_result_set.py @@ -4,7 +4,8 @@ # See LICENSE or https://opensource.org/licenses/MIT. # # SPDX-License-Identifier: MIT -import time +import asyncio +import threading from unittest.mock import MagicMock import pytest @@ -14,12 +15,7 @@ from pyathena.error import OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.util import RetryConfig -from tests.pyathena.util import ( - EVENT_TIMEOUT, - close_during_next_request, - close_during_next_retries, - get_query_results_responses, -) +from tests.pyathena.util import EVENT_TIMEOUT, get_query_results_responses async def _create_result_set(result_set_class, pages=(("1", "2"),)): @@ -50,6 +46,17 @@ async def _create_result_set(result_set_class, pages=(("1", "2"),)): ) +async def _assert_fetches_fail(result_set): + """Assert that each fetch of a result set raises for an earlier failed fetch. + + Args: + result_set: The asyncio result set. + """ + for fetch in (result_set.fetchone, result_set.fetchmany, result_set.fetchall): + with pytest.raises(ProgrammingError, match=r"cannot be fetched after a failed fetch\."): + await fetch() + + class TestAthenaAioResultSet: @pytest.mark.parametrize("result_set_class", [AthenaAioResultSet, AthenaAioDictResultSet]) async def test_sync_iteration_raises(self, result_set_class): @@ -105,15 +112,43 @@ async def test_fetchone_across_pages(self): assert result_set.rownumber == 3 assert result_set.connection.client.get_query_results.call_count == 2 - @pytest.mark.parametrize(("method", "args"), [("fetchmany", (3,)), ("fetchall", ())]) - async def test_fetch_counts_rows_taken_before_error(self, method, args): - """The rows taken before a failed page request are counted in rownumber.""" - result_set = await _create_result_set( - AthenaAioResultSet, [("1", "2"), RuntimeError("failed")] - ) - with pytest.raises(OperationalError): - await getattr(result_set, method)(*args) - assert result_set.rownumber == 2 + @pytest.mark.parametrize( + ("page", "expected"), [(RuntimeError("failed"), OperationalError), (("x",), ValueError)] + ) + @pytest.mark.parametrize("method", ["fetchone", "fetchmany", "fetchall"]) + async def test_fetch_after_failed_fetch(self, method, page, expected): + """GH-1225: after a fetch raises, later fetches raise without a request.""" + result_set = await _create_result_set(AthenaAioResultSet, [("1", "2"), page, ("3",)]) + assert await result_set.fetchmany(2) == [(1,), (2,)] + with pytest.raises(expected): + await getattr(result_set, method)() + await _assert_fetches_fail(result_set) + assert result_set.connection.client.get_query_results.call_count == 2 + + @pytest.mark.parametrize("method", ["fetchone", "fetchmany", "fetchall"]) + async def test_fetch_after_cancelled_fetch(self, method): + """GH-1225: after a fetch is cancelled during its page request, later fetches raise.""" + result_set = await _create_result_set(AthenaAioResultSet, [("1", "2"), ("3",)]) + assert await result_set.fetchmany(2) == [(1,), (2,)] + requested, release = threading.Event(), threading.Event() + responses = result_set.connection.client.get_query_results.side_effect + + def respond(**kwargs): + requested.set() + assert release.wait(EVENT_TIMEOUT) + return next(responses) + + result_set.connection.client.get_query_results.side_effect = respond + task = asyncio.create_task(getattr(result_set, method)()) + try: + assert await asyncio.to_thread(requested.wait, EVENT_TIMEOUT) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + await _assert_fetches_fail(result_set) + finally: + release.set() + assert result_set.connection.client.get_query_results.call_count == 2 async def test_fetch_after_close(self): result_set = await _create_result_set(AthenaAioResultSet, [("1", "2"), ("3",)]) @@ -123,39 +158,8 @@ 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() with pytest.raises(ProgrammingError, match=r"AthenaAioResultSet is closed\."): await result_set._async_get_query_results(2, "token1") - - @pytest.mark.parametrize( - ("method", "expected"), [("fetchone", None), ("fetchmany", []), ("fetchall", [])] - ) - async def test_fetch_with_close_during_request(self, method, expected): - """A close() during a page request discards the page and fetches nothing more.""" - result_set = await _create_result_set(AthenaAioResultSet, [("1", "2"), ("3", "4"), ("5",)]) - assert await result_set.fetchmany(2) == [(1,), (2,)] - get_query_results = close_during_next_request(result_set) - assert await getattr(result_set, method)() == expected - assert await result_set.fetchone() is None - assert await result_set.fetchmany() == [] - assert await result_set.fetchall() == [] - assert result_set.rownumber is None - assert get_query_results.call_count == 2 diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index 5440941f1..59c91e9dd 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -5,13 +5,11 @@ # # SPDX-License-Identifier: MIT -import collections import csv import inspect import io import threading from collections import defaultdict -from concurrent.futures import ThreadPoolExecutor from unittest.mock import ANY, MagicMock, PropertyMock, patch, sentinel import numpy as np @@ -148,40 +146,25 @@ def test_as_pandas_single_dataframe(self): [PandasDataFrameIterator.__next__, PandasDataFrameIterator.get_chunk], ids=["next", "get_chunk"], ) - @pytest.mark.parametrize("close_error", [None, RuntimeError("close failed")]) - def test_close_during_read(self, caplog, read, close_error): - """A close during a read in another thread is deferred until the read completes.""" - reading, release = threading.Event(), threading.Event() + @pytest.mark.parametrize("error", [OperationalError("failed"), KeyboardInterrupt()]) + def test_read_after_failed_read(self, read, error): + """GH-1225: after a read raises, the iterator is closed and later reads raise.""" def chunks(): - try: - yield pd.DataFrame({"n": [1]}) - reading.set() - release.wait(10) - yield pd.DataFrame({"n": [2]}) - yield pd.DataFrame({"n": [3]}) - finally: - if close_error: - raise close_error + yield pd.DataFrame({"n": [1]}) + raise error csv_stream = io.StringIO() reader = chunks() df_iter = PandasDataFrameIterator(reader, _no_trunc_date, csv_stream) - next(df_iter) - with ThreadPoolExecutor(1) as executor: - future = executor.submit(read, df_iter) - assert reading.wait(10) - df_iter.close() - with pytest.raises(StopIteration): - next(df_iter) - assert inspect.getgeneratorstate(reader) == inspect.GEN_RUNNING - assert not csv_stream.closed - release.set() - assert future.result(10)["n"].tolist() == [2] + assert read(df_iter)["n"].tolist() == [1] + with pytest.raises(type(error)): + read(df_iter) assert inspect.getgeneratorstate(reader) == inspect.GEN_CLOSED - assert ("Failed to close the chunk reader" in caplog.text) == bool(close_error) assert csv_stream.closed - assert list(df_iter) == [] + for later in (next, PandasDataFrameIterator.get_chunk): + with pytest.raises(ProgrammingError, match=r"cannot be read after a failed read\."): + later(df_iter) _FS = sentinel.pyathena_fs @@ -251,20 +234,17 @@ def test_fetch_duplicate_column_names(self): assert result_set.fetchmany() == [(3, "b", 4)] @pytest.mark.parametrize( - ("seconds", "method", "rownumber"), + ("seconds", "method"), [ - ([0, 10**13, 0], "fetchone", 1), - ([0, 10**13, 0], "fetchmany", 1), - ([0, 10**13, 0], "fetchall", 1), - ([0, 0, 10**13], "fetchmany", 2), - ([0, 0, 10**13], "fetchall", 2), + ([0, 10**13, 0], "fetchone"), + ([0, 10**13, 0], "fetchmany"), + ([0, 10**13, 0], "fetchall"), + ([0, 0, 10**13], "fetchmany"), + ([0, 0, 10**13], "fetchall"), ], ) - def test_fetch_after_conversion_error(self, seconds, method, rownumber): - """After a value fails to convert, the fetch methods return no more rows. - - ``rownumber`` counts the rows taken before the error, as fetching row by row did. - """ + def test_fetch_after_conversion_error(self, seconds, method): + """GH-1225: after a value fails to convert, the fetch methods raise.""" # A timestamp that does not fit a Python datetime fails while its row is built. timestamps = pa.array(seconds, type=pa.timestamp("s")) chunks = [ @@ -277,11 +257,21 @@ def test_fetch_after_conversion_error(self, seconds, method, rownumber): assert result_set.fetchone() == (0, pd.Timestamp(0, unit="s")) with pytest.raises(OverflowError): getattr(result_set, method)() - assert result_set.rownumber == rownumber - assert result_set.fetchone() is None - assert result_set.fetchmany() == [] - assert result_set.fetchall() == [] - assert result_set.rownumber == rownumber + _assert_fetches_fail(result_set) + + def test_fetch_after_failed_chunk_read(self): + """GH-1225: after iter_chunks() fails to read a chunk, the fetch methods raise.""" + + def chunks(): + yield pd.DataFrame({"n": [0]}) + raise OperationalError("failed") + + result_set = _fetching_result_set(chunks(), arraysize=3) + with pytest.raises(OperationalError): + list(result_set.iter_chunks()) + with pytest.raises(ProgrammingError, match=r"cannot be read after a failed read\."): + result_set.fetchone() + _assert_fetches_fail(result_set) def test_fetch_after_chunks_run_out(self): """The rows of a chunk that a fetch started stay after iter_chunks() runs out.""" @@ -292,59 +282,6 @@ def test_fetch_after_chunks_run_out(self): assert result_set.fetchall() == [(i,) for i in range(1, size)] assert result_set.rownumber == size - def test_close_during_fetchall(self): - """A fetchall() in another thread does not set rownumber after close().""" - requested = threading.Event() - closed = threading.Event() - - class Chunks: - def __init__(self): - self._chunks = iter( - [pd.DataFrame({"n": [0, 1]}), pd.DataFrame({"n": [2, 3]}, index=[2, 3])] - ) - self._calls = 0 - - def __iter__(self): - return self - - def __next__(self): - self._calls += 1 - if self._calls == 2: - requested.set() - assert closed.wait(10) - return next(self._chunks) - - result_set = _fetching_result_set(Chunks(), arraysize=3) - with ThreadPoolExecutor(max_workers=1) as executor: - future = executor.submit(result_set.fetchall) - assert requested.wait(10) - result_set.close() - closed.set() - # The fetch still returns the chunk that it was reading. - assert future.result(10) == [(0,), (1,), (2,), (3,)] - assert result_set.rownumber is None - - def test_close_during_fetchall_generator(self): - """A close() during a fetch from a generator, such as a chunked UNLOAD, does not raise.""" - reading, release = threading.Event(), threading.Event() - - def chunks(): - yield pd.DataFrame({"n": [0, 1]}) - reading.set() - release.wait(10) - yield pd.DataFrame({"n": [2, 3]}, index=[2, 3]) - yield pd.DataFrame({"n": [4]}, index=[4]) - - reader = chunks() - result_set = _fetching_result_set(reader, arraysize=3) - with ThreadPoolExecutor(max_workers=1) as executor: - future = executor.submit(result_set.fetchall) - assert reading.wait(10) - result_set.close() - release.set() - assert future.result(10) == [(0,), (1,), (2,), (3,)] - assert inspect.getgeneratorstate(reader) == inspect.GEN_CLOSED - def test_fetchmany_size_above_maxsize(self): result_set = _fetching_result_set(pd.DataFrame({"n": [1, 2]}), arraysize=3) assert result_set.fetchmany(10**100) == [(1,), (2,)] @@ -600,25 +537,6 @@ def test_iter_parquet_chunks_nested_values(self, tmp_path): assert chunk.dtypes.tolist() == whole.dtypes.tolist() assert_frame_equal(pd.concat(chunks), whole, check_exact=True) - def test_iter_parquet_chunks_after_close(self, tmp_path): - """Chunks of an UNLOAD result are read after close(), which resets the data manifest. - - A fetch in another thread, such as one of AioPandasCursor that keeps running - after its task is cancelled, can start its first read as close() lands. - """ - result_set, paths = _unload_result_set(tmp_path, [pa.table({"i": [1, 2]})], chunksize=1) - with ( - patch.object(AthenaPandasResultSet, "_read_data_manifest", return_value=paths), - patch.object( - AthenaPandasResultSet, "is_unload", new_callable=PropertyMock, return_value=True - ), - ): - reader = result_set._as_pandas() - result_set._rows = collections.deque() - result_set._df_iter = PandasDataFrameIterator(pd.DataFrame(), _no_trunc_date) - 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. @@ -1360,6 +1278,13 @@ def test_read_csv_with_pyarrow_multiline_values_across_blocks(): assert df["v"].tolist() == [v for _, v in rows] +def _assert_fetches_fail(result_set): + """Assert that each fetch of a result set raises for an earlier failed fetch.""" + for fetch in (result_set.fetchone, result_set.fetchmany, result_set.fetchall): + with pytest.raises(ProgrammingError, match=r"cannot be fetched after a failed fetch\."): + fetch() + + def _fetching_result_set(result, arraysize): """Create a result set of a succeeded query whose read gives ``result``, without AWS. diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 8684aa6c7..95a68d1d5 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -6,12 +6,9 @@ # SPDX-License-Identifier: MIT import inspect -import threading from collections.abc import Iterator -from concurrent.futures import ThreadPoolExecutor from datetime import UTC, date, datetime, time, timedelta from decimal import Decimal -from io import BytesIO from unittest.mock import MagicMock, PropertyMock, patch import polars as pl @@ -19,7 +16,7 @@ from fsspec.implementations.local import LocalFileSystem from pyathena.converter import _to_date, _to_default -from pyathena.error import OperationalError +from pyathena.error import OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.polars.converter import DefaultPolarsTypeConverter from pyathena.polars.result_set import AthenaPolarsResultSet, PolarsDataFrameIterator @@ -61,6 +58,13 @@ def _read_all(result: pl.DataFrame | Iterator[pl.DataFrame]) -> pl.DataFrame: return result if isinstance(result, pl.DataFrame) else pl.concat(list(result)) +def _assert_fetches_fail(result_set): + """Assert that each fetch of a result set raises for an earlier failed fetch.""" + for fetch in (result_set.fetchone, result_set.fetchmany, result_set.fetchall): + with pytest.raises(ProgrammingError, match=r"cannot be fetched after a failed fetch\."): + fetch() + + def _fetching_result_set(chunks, names, arraysize, converters=None) -> AthenaPolarsResultSet: """Create a result set of a succeeded query whose chunks are ``chunks``, without AWS. @@ -139,20 +143,17 @@ def test_fetchall_across_chunks(self): assert result_set.rownumber == 4 @pytest.mark.parametrize( - ("dates", "method", "rownumber"), + ("dates", "method"), [ - (["2024-01-01", "bad", "2024-01-03"], "fetchone", 1), - (["2024-01-01", "bad", "2024-01-03"], "fetchmany", 1), - (["2024-01-01", "bad", "2024-01-03"], "fetchall", 1), - (["2024-01-01", "2024-01-02", "bad"], "fetchmany", 2), - (["2024-01-01", "2024-01-02", "bad"], "fetchall", 2), + (["2024-01-01", "bad", "2024-01-03"], "fetchone"), + (["2024-01-01", "bad", "2024-01-03"], "fetchmany"), + (["2024-01-01", "bad", "2024-01-03"], "fetchall"), + (["2024-01-01", "2024-01-02", "bad"], "fetchmany"), + (["2024-01-01", "2024-01-02", "bad"], "fetchall"), ], ) - def test_fetch_after_conversion_error(self, dates, method, rownumber): - """After a value fails to convert, the fetch methods return no more rows. - - ``rownumber`` counts the rows taken before the error, as fetching row by row did. - """ + def test_fetch_after_conversion_error(self, dates, method): + """GH-1225: after a value fails to convert, the fetch methods raise.""" chunks = [ pl.DataFrame({"n": [0, 1, 2], "d": dates}), pl.DataFrame({"n": [3], "d": ["2024-01-04"]}), @@ -163,11 +164,21 @@ def test_fetch_after_conversion_error(self, dates, method, rownumber): assert result_set.fetchone() == (0, date(2024, 1, 1)) with pytest.raises(ValueError, match="bad"): getattr(result_set, method)() - assert result_set.rownumber == rownumber - assert result_set.fetchone() is None - assert result_set.fetchmany() == [] - assert result_set.fetchall() == [] - assert result_set.rownumber == rownumber + _assert_fetches_fail(result_set) + + def test_fetch_after_failed_chunk_read(self): + """GH-1225: after iter_chunks() fails to read a chunk, the fetch methods raise.""" + + def chunks(): + yield pl.DataFrame({"n": [0]}) + raise OperationalError("failed") + + result_set = _fetching_result_set(chunks(), ["n"], arraysize=3) + with pytest.raises(OperationalError): + list(result_set.iter_chunks()) + with pytest.raises(ProgrammingError, match=r"cannot be read after a failed read\."): + result_set.fetchone() + _assert_fetches_fail(result_set) def test_fetch_after_chunks_run_out(self): """The rows of a chunk that a fetch started stay after iter_chunks() runs out.""" @@ -178,133 +189,6 @@ def test_fetch_after_chunks_run_out(self): assert result_set.fetchall() == [(i,) for i in range(1, size)] assert result_set.rownumber == size - def test_close_during_fetchall(self): - """A fetchall() in another thread does not set rownumber after close().""" - requested = threading.Event() - closed = threading.Event() - - class Chunks: - def __init__(self): - self._chunks = iter([pl.DataFrame({"n": [0, 1]}), pl.DataFrame({"n": [2, 3]})]) - self._calls = 0 - - def __iter__(self): - return self - - def __next__(self): - self._calls += 1 - if self._calls == 2: - requested.set() - assert closed.wait(10) - return next(self._chunks) - - result_set = _fetching_result_set(Chunks(), ["n"], arraysize=3) - with ThreadPoolExecutor(max_workers=1) as executor: - future = executor.submit(result_set.fetchall) - assert requested.wait(10) - result_set.close() - closed.set() - # The fetch still returns the chunk that it was reading. - assert future.result(10) == [(0,), (1,), (2,), (3,)] - assert result_set.rownumber is None - - def test_close_during_fetchall_generator(self): - """A close() during a fetch from a generator, such as a chunked UNLOAD, does not raise.""" - reading, release = threading.Event(), threading.Event() - - def chunks(): - yield pl.DataFrame({"n": [0, 1]}) - reading.set() - release.wait(10) - yield pl.DataFrame({"n": [2, 3]}) - yield pl.DataFrame({"n": [4]}) - - reader = chunks() - result_set = _fetching_result_set(reader, ["n"], arraysize=3) - with ThreadPoolExecutor(max_workers=1) as executor: - future = executor.submit(result_set.fetchall) - assert reading.wait(10) - result_set.close() - release.set() - assert future.result(10) == [(0,), (1,), (2,), (3,)] - assert inspect.getgeneratorstate(reader) == inspect.GEN_CLOSED - - @pytest.mark.parametrize( - ("unload", "kwargs", "credentials"), - [(False, {}, (0, 1)), (True, {}, (1, 1)), (True, {"low_memory": True}, (1, 2))], - ids=["csv", "parquet-chunk-reader", "scan-parquet"], - ) - def test_close_during_first_read(self, tmp_path, unload, kwargs, credentials): - """A close() as a fetch starts its first read leaves the fetch its chunk. - - A fetch in another thread, such as one of AioPolarsCursor that keeps running - after its task is cancelled, can start its first read as close() lands. The - credentials of the Polars scans are frozen at the first read, so ``credentials`` - counts them after the result set is created and after the read; the schema of - an UNLOAD result is read when it is created. - """ - if unload: - pl.DataFrame({"a": [1, 2, 3]}).write_parquet(tmp_path / "0.parquet") - manifest = str(tmp_path / "0.parquet") - output_location = "s3://bucket/result.csv" - else: - (tmp_path / "result.csv").write_text('"a"\n"1"\n"2"\n"3"\n') - manifest = "" - output_location = str(tmp_path / "result.csv") - connection = MagicMock(region_name=None) - connection.session.get_credentials.return_value = None - connection.s3_client.get_object.side_effect = lambda **_: { - "Body": BytesIO(manifest.encode()) - } - query_execution = MagicMock( - state=AthenaQueryExecution.STATE_SUCCEEDED, - output_location=output_location, - data_manifest_location="s3://bucket/manifest.csv", - query="UNLOAD (SELECT a FROM t) TO 's3://bucket/unload/' WITH (format = 'PARQUET')" - if unload - else "SELECT a FROM t", - ) - metadata = ( - {"Name": "a", "Type": "integer", "Precision": 10, "Scale": 0, "Nullable": "UNKNOWN"}, - ) - with ( - patch.object( - AthenaPolarsResultSet, - "_pre_fetch", - autospec=True, - side_effect=lambda rs: setattr(rs, "_metadata", metadata), - ), - patch.object(AthenaPolarsResultSet, "_get_content_length", return_value=16), - patch("pyathena.polars.result_set.S3FileSystem", lambda **_: LocalFileSystem()), - ): - result_set = AthenaPolarsResultSet( - connection=connection, - converter=DefaultPolarsTypeConverter(), - query_execution=query_execution, - arraysize=3, - retry_config=RetryConfig(), - unload=unload, - chunksize=2, - **kwargs, - ) - assert connection.session.get_credentials.call_count == credentials[0] - reading, release = threading.Event(), threading.Event() - reader = result_set._df_iter._reader - - def blocked(): - reading.set() - assert release.wait(10) - yield from reader - - result_set._df_iter._reader = blocked() - with ThreadPoolExecutor(max_workers=1) as executor: - future = executor.submit(result_set.fetchall) - assert reading.wait(10) - result_set.close() - release.set() - assert future.result(10) == [(1,), (2,)] - assert connection.session.get_credentials.call_count == credentials[1] - def test_fetchmany_size_above_maxsize(self): """A size above sys.maxsize returns the remaining rows.""" result_set = _fetching_result_set([pl.DataFrame({"n": [1, 2]})], ["n"], arraysize=3) @@ -791,34 +675,19 @@ def test_close_stops_iteration(self, reader): df_iter.close() assert list(df_iter) == [] - @pytest.mark.parametrize("close_error", [None, RuntimeError("close failed")]) - def test_close_during_read(self, caplog, close_error): - """A close during a read in another thread is deferred until the read completes.""" - reading, release = threading.Event(), threading.Event() + @pytest.mark.parametrize("error", [OperationalError("failed"), KeyboardInterrupt()]) + def test_read_after_failed_read(self, error): + """GH-1225: after a read raises, the iterator is closed and later reads raise.""" def chunks(): - try: - yield pl.DataFrame({"a": [1]}) - reading.set() - release.wait(10) - yield pl.DataFrame({"a": [2]}) - yield pl.DataFrame({"a": [3]}) - finally: - if close_error: - raise close_error + yield pl.DataFrame({"a": [1]}) + raise error reader = chunks() df_iter = PolarsDataFrameIterator(reader, {}, ["a"]) - next(df_iter) - with ThreadPoolExecutor(1) as executor: - future = executor.submit(next, df_iter) - assert reading.wait(10) - df_iter.close() - with pytest.raises(StopIteration): - next(df_iter) - assert inspect.getgeneratorstate(reader) == inspect.GEN_RUNNING - release.set() - assert future.result(10)["a"].to_list() == [2] + assert next(df_iter)["a"].to_list() == [1] + with pytest.raises(type(error)): + next(df_iter) assert inspect.getgeneratorstate(reader) == inspect.GEN_CLOSED - assert ("Failed to close the chunk reader" in caplog.text) == bool(close_error) - assert list(df_iter) == [] + with pytest.raises(ProgrammingError, match=r"cannot be read after a failed read\."): + next(df_iter) diff --git a/tests/pyathena/test_result_set.py b/tests/pyathena/test_result_set.py index d68d69522..c4b73a07e 100644 --- a/tests/pyathena/test_result_set.py +++ b/tests/pyathena/test_result_set.py @@ -4,9 +4,7 @@ # See LICENSE or https://opensource.org/licenses/MIT. # # SPDX-License-Identifier: MIT -import collections import sys -import time from unittest.mock import MagicMock, patch import pytest @@ -17,13 +15,8 @@ from pyathena.error import OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.result_set import AthenaResultSet, WithFetch, WithResultSet -from pyathena.util import RetryConfig, override -from tests.pyathena.util import ( - EVENT_TIMEOUT, - close_during_next_request, - close_during_next_retries, - get_query_results_responses, -) +from pyathena.util import RetryConfig +from tests.pyathena.util import get_query_results_responses def _row(*values): @@ -53,31 +46,6 @@ def _fetching_result_set(*pages): ) -class _ClosingDeque(collections.deque): - """A buffer that closes its result set when it is read, as another thread may.""" - - def __init__(self, rows, result_set): - """Initialize the buffer. - - Args: - rows: The buffered rows. - result_set: The result set to close. - """ - super().__init__(rows) - self._result_set = result_set - - @override - def popleft(self): - row = super().popleft() - self._result_set.close() - return row - - @override - def __iter__(self): - self._result_set.close() - return super().__iter__() - - class TestAthenaResultSet: def test_fetchmany_across_pages(self): result_set = _fetching_result_set(("1", "2"), ("3", "4"), ("5",)) @@ -110,90 +78,36 @@ def test_fetchmany_size_not_integer(self): result_set.fetchmany(1.5) assert result_set.fetchall() == [(1,), (2,)] - @pytest.mark.parametrize(("method", "args"), [("fetchmany", (3,)), ("fetchall", ())]) - def test_fetch_counts_rows_taken_before_error(self, method, args): - """The rows taken before a failed page request are counted in rownumber.""" - result_set = _fetching_result_set(("1", "2"), RuntimeError("failed")) - with pytest.raises(OperationalError): - getattr(result_set, method)(*args) - assert result_set.rownumber == 2 - - def test_fetch_after_close(self): - result_set = _fetching_result_set(("1", "2"), ("3",)) - result_set.close() - assert result_set.fetchmany() == [] - assert result_set.fetchall() == [] - assert result_set.rownumber is None - - @pytest.mark.parametrize( - ("method", "args", "expected"), - [("fetchmany", (2,), [(1,)]), ("fetchall", (), [(1,), (2,), (3,)])], - ) - def test_fetch_with_close_during_take(self, method, args, expected): - """A close() while the buffered rows are taken keeps the rows taken before it.""" - result_set = _fetching_result_set(("1", "2", "3")) - result_set._rows = _ClosingDeque(result_set._rows, result_set) - assert getattr(result_set, method)(*args) == expected - assert result_set.rownumber is None - - @pytest.mark.parametrize( - ("method", "expected"), [("fetchone", None), ("fetchmany", []), ("fetchall", [])] - ) - def test_fetch_with_close_during_request(self, method, expected): - """A close() during a page request discards the page and fetches nothing more.""" - result_set = _fetching_result_set(("1", "2"), ("3", "4"), ("5",)) - assert result_set.fetchmany(2) == [(1,), (2,)] - get_query_results = close_during_next_request(result_set) - assert getattr(result_set, method)() == expected - assert result_set.fetchone() is None - assert result_set.fetchmany() == [] - assert result_set.fetchall() == [] - 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", [])] + ("page", "expected"), + [ + (RuntimeError("failed"), OperationalError), + (KeyboardInterrupt(), KeyboardInterrupt), + (("x",), ValueError), + ], ) - 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. + @pytest.mark.parametrize("method", ["fetchone", "fetchmany", "fetchall"]) + def test_fetch_after_failed_fetch(self, method, page, expected): + """GH-1225: after a fetch raises, later fetches raise without a request. - The fetch returns nothing, as after a close(), instead of retrying the - throttled request with the policy's attempts and backoff. + The page request fails, is interrupted, or returns a value that fails to + convert. """ - result_set = _fetching_result_set(("1", "2"), ("3",)) + result_set = _fetching_result_set(("1", "2"), page, ("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 + with pytest.raises(expected): + getattr(result_set, method)() + for fetch in (result_set.fetchone, result_set.fetchmany, result_set.fetchall): + with pytest.raises(ProgrammingError, match=r"cannot be fetched after a failed fetch\."): + fetch() + assert result_set.connection.client.get_query_results.call_count == 2 - @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.""" + def test_fetch_after_close(self): 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 + result_set.close() + assert result_set.fetchmany() == [] + assert result_set.fetchall() == [] + assert result_set.rownumber is None def test_get_query_results_after_close(self): result_set = _fetching_result_set(("1", "2"), ("3",)) diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 95d4a053c..3f6e2bebd 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -26,7 +26,7 @@ from pyathena.filesystem.s3_core import S3Core from pyathena.glue import GlueMetadataClient from pyathena.model import AthenaCalculationExecutionStatus, AthenaQueryExecution -from pyathena.util import RetryConfig, retry_api_call +from pyathena.util import retry_api_call _queries = Environment( loader=FileSystemLoader(Path(__file__).parents[1].resolve() / "resources" / "queries") @@ -41,8 +41,9 @@ def get_query_results_responses(*pages): """Build stubbed ``GetQueryResults`` responses of an integer column ``a``. Args: - *pages: The values of the column, one sequence per page, or an exception that - the request for the page raises. + *pages: The values of the column, one sequence per page, or an exception, + including a ``BaseException`` such as ``KeyboardInterrupt``, that the + request for the page raises. Returns: The responses, the first with the result set metadata and each but the last @@ -50,7 +51,7 @@ def get_query_results_responses(*pages): """ responses = [] for i, page in enumerate(pages): - if isinstance(page, Exception): + if isinstance(page, BaseException): responses.append(page) continue response = {"ResultSet": {"Rows": [{"Data": [{"VarCharValue": v}]} for v in page]}} @@ -64,27 +65,6 @@ def get_query_results_responses(*pages): return responses -def close_during_next_request(result_set): - """Close the result set during its next ``GetQueryResults`` request, as another thread may. - - Args: - result_set: A result set whose ``GetQueryResults`` responses are stubbed with - the pages of :func:`get_query_results_responses`, none of them an exception. - - Returns: - The stubbed ``GetQueryResults`` mock. - """ - get_query_results = result_set.connection.client.get_query_results - responses = get_query_results.side_effect - - def close_and_respond(**kwargs): - result_set.close() - return next(responses) - - get_query_results.side_effect = close_and_respond - return get_query_results - - def throttling_error(operation_name="GetQueryResults"): """Build the ``ThrottlingException`` error of an AWS API operation. @@ -192,27 +172,6 @@ def join(self): 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. From 7ba9fbfbe30a2ca809144247a12ea52b8b2fdb17 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 17:39:08 +0900 Subject: [PATCH 2/7] Remove the stop_when condition of retry_api_call() The result sets were its only callers, to stop their retries on a close() from another thread, which is no longer supported. The stop scopes of stop_retries_on() remain for asyncio cancellations. Refs #1206 Co-Authored-By: Claude Opus 5.5 --- pyathena/aio/util.py | 6 +----- pyathena/util.py | 18 +++++------------- tests/pyathena/aio/test_util.py | 11 ----------- tests/pyathena/test_util.py | 16 ---------------- 4 files changed, 6 insertions(+), 45 deletions(-) diff --git a/pyathena/aio/util.py b/pyathena/aio/util.py index 5770abec4..43735bf10 100644 --- a/pyathena/aio/util.py +++ b/pyathena/aio/util.py @@ -69,7 +69,6 @@ async def async_retry_api_call( 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: @@ -85,7 +84,6 @@ 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 @@ -100,9 +98,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 unstoppable_retries(): - return retry_api_call( - func, config, logger, *args, stop_on=stop_on, stop_when=stop_when, **kwargs - ) + return retry_api_call(func, config, logger, *args, stop_on=stop_on, **kwargs) return await (run_in_thread(call) if stoppable else asyncio.to_thread(call)) diff --git a/pyathena/util.py b/pyathena/util.py index bb3155727..296255a2f 100644 --- a/pyathena/util.py +++ b/pyathena/util.py @@ -313,8 +313,7 @@ def unstoppable_retries() -> Iterator[None]: 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. + :func:`stop_retries_on` opens in the block stops the retries again. Yields: None. @@ -332,7 +331,6 @@ def retry_api_call( 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. @@ -348,10 +346,6 @@ 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: @@ -379,13 +373,13 @@ def retry_api_call( 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()) + """Whether an event of the enclosing scopes is set.""" + return any(event.is_set() for event in events) # 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 + wait_for_event = len(events) == 1 def record_failure(retry_state: RetryCallState) -> None: """Keep the error of the failed attempt, which a stopped wait raises. @@ -419,9 +413,7 @@ def sleep(seconds: float) -> None: 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 {} - ) + stop_kwargs: dict[str, Any] = {"before_sleep": record_failure, "sleep": sleep} if events else {} retry = tenacity.Retrying( retry=retry_if_exception( diff --git a/tests/pyathena/aio/test_util.py b/tests/pyathena/aio/test_util.py index 35764cb62..3de600b9c 100644 --- a/tests/pyathena/aio/test_util.py +++ b/tests/pyathena/aio/test_util.py @@ -12,7 +12,6 @@ from concurrent.futures import ThreadPoolExecutor import pytest -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, stop_retries_on @@ -239,16 +238,6 @@ async def test_async_retry_api_call_cancelled(stop_in): 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) diff --git a/tests/pyathena/test_util.py b/tests/pyathena/test_util.py index c4fb941d0..f3c93759e 100644 --- a/tests/pyathena/test_util.py +++ b/tests/pyathena/test_util.py @@ -394,22 +394,6 @@ def test_retry_api_call_stopped_by_event(stop_in): 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() From 6b9c77a03d56eaa88bf0c23e36f882aeb34548fe Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 18:00:29 +0900 Subject: [PATCH 3/7] Close the result set of a cancelled fetch without the event loop Add pyathena.aio.util.run_in_thread_or_close(), which the asyncio fetches use instead of their own shielded executor call. As run_in_thread() does, a cancellation skips a fetch that has not started and stops the retries of a running one. The fetch's thread closes the result set when the fetch ends, and otherwise a new thread closes it, as create_in_thread() does, so the close neither runs on the event loop nor waits for a closed loop. The future is awaited directly, since asyncio.shield() and asyncio.wait() add an event loop iteration, about 20 microseconds, to every fetch. Restore the counting of the rows that the pandas and Polars bulk takes took before an exception, as the base and asyncio result sets count them, and correct docstrings and documentation found in review. Refs #1206 Co-Authored-By: Claude Opus 5.5 --- docs/usage.md | 3 +- pyathena/aio/common.py | 49 +++-------------- pyathena/aio/s3fs/cursor.py | 4 +- pyathena/aio/util.py | 78 ++++++++++++++++++++++++++ pyathena/pandas/result_set.py | 18 ++++-- pyathena/polars/result_set.py | 24 +++++--- tests/pyathena/aio/test_util.py | 98 ++++++++++++++++++++++++++++++++- 7 files changed, 213 insertions(+), 61 deletions(-) diff --git a/docs/usage.md b/docs/usage.md index a9f1d9986..a52ca948b 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -131,6 +131,7 @@ for row in cursor: ``` After a fetch raises an exception, including `KeyboardInterrupt`, later fetches of the same result raise `ProgrammingError`. +An invalid argument, such as a `fetchmany()` size that is not an integer, raises its error without that effect. Execute the query again to read its result. As `pyathena.threadsafety = 2` declares, threads can share connections, but not cursors. @@ -405,7 +406,7 @@ Result columns from a passthrough query come back with the source system's types PyAthena provides a callback mechanism that allows you to get immediate access to the query ID as soon as the `start_query_execution` API call is made, before waiting for query completion. -This is useful for monitoring, logging, or cancelling long-running queries by their ID. +This is useful for monitoring, logging, or stopping long-running queries by their ID, for example with `AsyncCursor.cancel(query_id)`. When `cache_size` finds a reusable query, no new query starts, and the callback receives the reused query's ID. The `on_start_query_execution` callback can be configured at both the connection level and diff --git a/pyathena/aio/common.py b/pyathena/aio/common.py index f29db358f..0ea27c71a 100644 --- a/pyathena/aio/common.py +++ b/pyathena/aio/common.py @@ -3,41 +3,27 @@ from __future__ import annotations import asyncio -import contextvars import logging import operator -import threading from collections.abc import Awaitable, Callable, Coroutine from typing import Any, NoReturn, TypeVar, cast 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_or_close from pyathena.common import BaseCursor, CursorIterator from pyathena.error import DatabaseError, OperationalError, ProgrammingError from pyathena.glue import GlueMetadataClient from pyathena.model import AthenaDatabase, AthenaQueryExecution, AthenaTableMetadata from pyathena.options import ExecuteOptions from pyathena.result_set import AthenaResultSet, WithResultSet -from pyathena.util import _is_throttling_error, override, stop_retries_on +from pyathena.util import _is_throttling_error, override _logger = logging.getLogger(__name__) _T = TypeVar("_T") -def _close_result_set(result_set: AthenaResultSet) -> None: - """Close a result set that a cancelled fetch detached, logging an error. - - Args: - result_set: The result set. - """ - try: - result_set.close() - except Exception: - _logger.exception("Failed to close the result set of a cancelled fetch.") - - class AioBaseCursor(BaseCursor): """Async base cursor that overrides I/O methods with async equivalents. @@ -728,7 +714,7 @@ class WithAsyncFetch(WithResultSet, AioBaseCursor, CursorIterator): 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 - a result stream can wait for a read in progress. Synchronous iteration + a result stream can wait for its read-ahead requests in progress. Synchronous iteration raises ``TypeError``. Subclasses override ``execute()`` and optionally ``__init__`` and @@ -861,10 +847,9 @@ async def _fetch_in_thread(self, fetch: Callable[[AthenaResultSet], _T]) -> _T: """Run a synchronous fetch of the result set in a thread. A cancellation stops the retries of the AWS API calls that the fetch makes - (see :func:`~pyathena.util.stop_retries_on`) and detaches the result set from - the cursor, so that later fetches raise ``ProgrammingError``. The fetch can - still be running in its thread, so the result set is closed in a thread once - the fetch has ended there. + and detaches the result set from the cursor, so that later fetches raise + ``ProgrammingError``. The result set is closed once the fetch is no longer + running (see :func:`~pyathena.aio.util.run_in_thread_or_close`). Args: fetch: The function that fetches from the result set. @@ -878,31 +863,11 @@ async def _fetch_in_thread(self, fetch: Callable[[AthenaResultSet], _T]) -> _T: if not self.has_result_set: raise ProgrammingError("No result set.") result_set = cast(AthenaResultSet, self.result_set) - stop = threading.Event() - - def call() -> _T: - """Run the fetch in a scope that the cancellation stops.""" - with stop_retries_on(stop): - return fetch(result_set) - - def close_after_fetch(_: asyncio.Future[_T]) -> None: - """Close the detached result set in a thread once the fetch has ended.""" - try: - loop.run_in_executor(None, _close_result_set, result_set) - except RuntimeError: - # The default executor is shut down, as at the end of asyncio.run(). - _close_result_set(result_set) - - loop = asyncio.get_running_loop() - future = loop.run_in_executor(None, contextvars.copy_context().run, call) try: - # The shield keeps the future pending until the fetch has ended. - return await asyncio.shield(future) + return await run_in_thread_or_close(result_set, fetch, result_set) except asyncio.CancelledError: - stop.set() if self.result_set is result_set: self.result_set = None - future.add_done_callback(close_after_fetch) raise @override diff --git a/pyathena/aio/s3fs/cursor.py b/pyathena/aio/s3fs/cursor.py index d1653658b..f690dc085 100644 --- a/pyathena/aio/s3fs/cursor.py +++ b/pyathena/aio/s3fs/cursor.py @@ -32,8 +32,8 @@ 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. + ``asyncio.to_thread``. Fetch operations run in a 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 43735bf10..dc75b2bd0 100644 --- a/pyathena/aio/util.py +++ b/pyathena/aio/util.py @@ -19,6 +19,8 @@ from pyathena.util import RetryConfig, retry_api_call, stop_retries_on, unstoppable_retries +_logger = logging.getLogger(__name__) + T = TypeVar("T") @@ -179,3 +181,79 @@ def construct() -> _ClosableT | None: threading.Thread(target=constructed[0].close, name="create_in_thread-close").start() raise return cast(_ClosableT, future.result()) + + +async def run_in_thread_or_close( + closable: _Closable, func: Callable[..., T], /, *args: Any, **kwargs: Any +) -> T: + """Run a blocking function that uses a closable object in the default executor. + + Like :func:`run_in_thread`, the function runs in the event loop's default + executor with a copy of the caller's context, a cancellation of the caller + skips a function that has not started, and it stops the retries of the AWS + API calls that a running function makes. A cancellation also closes the + object once the function is no longer running: a running function closes it + in its thread when it ends, and otherwise a new thread closes it. As with + :func:`create_in_thread`, the close needs neither the event loop nor its + executor. + + Args: + closable: The object that the function uses. + func: The blocking function. + *args: Positional arguments for ``func``. + **kwargs: Keyword arguments for ``func``. + + Returns: + The return value of ``func``. + """ + lock = threading.Lock() + cancelled = threading.Event() + started = ended = False + + def call() -> T | None: + """Run the function unless the caller was cancelled, closing the object if it was then. + + Returns: + The return value of the function, or None if the caller was cancelled + before it started. + """ + nonlocal started, ended + with lock: + if cancelled.is_set(): + return None + started = True + try: + with stop_retries_on(cancelled): + return func(*args, **kwargs) + finally: + with lock: + ended = True + close = cancelled.is_set() + if close: + _close_logging_errors(closable) + + future = asyncio.get_running_loop().run_in_executor(None, contextvars.copy_context().run, call) + try: + return cast(T, await future) + except asyncio.CancelledError: + with lock: + cancelled.set() + close = not started or ended + if close: + threading.Thread( + target=_close_logging_errors, args=(closable,), name="run_in_thread_or_close-close" + ).start() + raise + + +def _close_logging_errors(closable: _Closable) -> None: + """Close an object for a cancelled caller, logging an error that nobody else sees. + + Args: + closable: The object to close. + """ + try: + with unstoppable_retries(): + closable.close() + except Exception: + _logger.exception(f"Failed to close {closable!r} after a cancellation.") diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 4a6e92a89..d6eead076 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -987,12 +987,18 @@ def fetchone( @override def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]: - """Take the next rows from the row iterator in bulk.""" - # islice() takes no size above sys.maxsize. - rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = list( - islice(self._row_iter, None if size is None else min(size, sys.maxsize)) - ) - self._count_rows(len(rows)) + """Take the next rows from the row iterator in bulk. + + The rows taken before an exception are counted in ``rownumber`` too. + """ + rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = [] + try: + # extend() keeps the rows that it appended before an exception, and + # islice() takes no size above sys.maxsize. + rows.extend(islice(self._row_iter, None if size is None else min(size, sys.maxsize))) + finally: + if rows: + self._count_rows(len(rows)) return rows def _read_csv(self) -> TextFileReader | DataFrame: diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 94f09095f..90bcb1bb5 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -112,12 +112,12 @@ def __init__( def __next__(self) -> pl.DataFrame: """Get the next DataFrame chunk. - Returns: - The next Polars DataFrame chunk. - The iterator is closed at the end of the chunks and after an exception, and later reads after an exception raise ``ProgrammingError``. + Returns: + The next Polars DataFrame chunk. + Raises: StopIteration: When no more chunks are available. ProgrammingError: If an earlier read failed. @@ -643,12 +643,18 @@ def fetchone( @override def _take_rows(self, size: int | None) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]: - """Take the next rows from the row iterator in bulk.""" - # islice() takes no size above sys.maxsize. - rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = list( - islice(self._row_iter, None if size is None else min(size, sys.maxsize)) - ) - self._count_rows(len(rows)) + """Take the next rows from the row iterator in bulk. + + The rows taken before an exception are counted in ``rownumber`` too. + """ + rows: list[tuple[Any | None, ...] | dict[Any, Any | None]] = [] + try: + # extend() keeps the rows that it appended before an exception, and + # islice() takes no size above sys.maxsize. + rows.extend(islice(self._row_iter, None if size is None else min(size, sys.maxsize))) + finally: + if rows: + self._count_rows(len(rows)) return rows def _is_csv_readable(self) -> bool: diff --git a/tests/pyathena/aio/test_util.py b/tests/pyathena/aio/test_util.py index 3de600b9c..98f09ea19 100644 --- a/tests/pyathena/aio/test_util.py +++ b/tests/pyathena/aio/test_util.py @@ -13,7 +13,12 @@ import pytest -from pyathena.aio.util import async_retry_api_call, create_in_thread, run_in_thread +from pyathena.aio.util import ( + async_retry_api_call, + create_in_thread, + run_in_thread, + run_in_thread_or_close, +) from pyathena.util import RetryConfig, retry_api_call, stop_retries_on from tests.pyathena.util import EVENT_TIMEOUT, ThrottledOperation, wait_until @@ -204,6 +209,97 @@ def factory(): executor.shutdown(wait=True) +async def test_run_in_thread_or_close(): + """The function runs in another thread with the caller's context, and nothing is closed.""" + _VALUE.set("caller") + resource = _Resource() + thread, value = await run_in_thread_or_close( + resource, lambda: (threading.get_ident(), _VALUE.get()) + ) + assert thread != threading.get_ident() + assert value == "caller" + assert not resource.closed.is_set() + + +async def test_run_in_thread_or_close_cancelled_during_function(): + """GH-1206: the object is closed by the function's thread once the function ends.""" + started, release = threading.Event(), threading.Event() + resource = _Resource() + threads = [] + + def func(): + threads.append(threading.get_ident()) + started.set() + assert release.wait(EVENT_TIMEOUT) + + task = asyncio.create_task(run_in_thread_or_close(resource, func)) + try: + assert await asyncio.to_thread(started.wait, EVENT_TIMEOUT) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert not resource.closed.is_set() + finally: + release.set() + assert await asyncio.to_thread(resource.closed.wait, EVENT_TIMEOUT) + assert resource.closed_in == threads[0] + + +async def test_run_in_thread_or_close_cancelled_after_function(): + """GH-1206: an object whose function ended before the cancellation is closed in a new thread.""" + resource = _Resource() + threads = [] + started, release, returned = threading.Event(), threading.Event(), threading.Event() + + def func(): + threads.append(threading.get_ident()) + started.set() + assert release.wait(EVENT_TIMEOUT) + returned.set() + + task = asyncio.create_task(run_in_thread_or_close(resource, func)) + assert await asyncio.to_thread(started.wait, EVENT_TIMEOUT) + # Blocking the event loop keeps the task from resuming after the function ends. + release.set() + assert returned.wait(EVENT_TIMEOUT) + time.sleep(0.1) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + + assert await asyncio.to_thread(resource.closed.wait, EVENT_TIMEOUT) + assert resource.closed_in not in (None, threads[0], threading.get_ident()) + + +async def test_run_in_thread_or_close_cancelled_before_function(): + """GH-1206: a function that has not started is skipped, and a new thread closes the object.""" + loop = asyncio.get_running_loop() + # One busy worker keeps the function queued until the caller is cancelled. + executor = ThreadPoolExecutor(max_workers=1) + loop.set_default_executor(executor) + busy = threading.Event() + calls = [] + resource = _Resource() + try: + blocker = loop.run_in_executor(None, busy.wait, EVENT_TIMEOUT) + task = asyncio.create_task(run_in_thread_or_close(resource, lambda: calls.append(1))) + await asyncio.sleep(0) + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + # The only worker of the default executor is busy, so this waits on the loop. + assert resource.closed.wait(EVENT_TIMEOUT) + assert resource.closed_in != threading.get_ident() + busy.set() + await blocker + # A function that the executor did not drop returns without running. + await loop.run_in_executor(None, lambda: None) + assert calls == [] + finally: + busy.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") From f5b27b2f1e7b57037f1e700bcda3ebe49ef552bc Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 18:05:40 +0900 Subject: [PATCH 4/7] Simplify the failed-fetch and cancelled-close code - Use one failed-fetch message for all result sets and share the test helper that asserts it. - Track whether the function of run_in_thread_or_close() is running with one flag, and close the objects of create_in_thread() with the same logging helper. - Read the data manifest in the pandas Parquet chunk reader again; the paths parameter only served a close() from another thread. - Update rownumber in _count_rows() directly. Refs #1206 Co-Authored-By: Claude Opus 5.5 --- pyathena/aio/result_set.py | 4 +--- pyathena/aio/util.py | 17 +++++++++-------- pyathena/pandas/result_set.py | 11 ++++------- pyathena/result_set.py | 7 +++---- tests/pyathena/pandas/test_result_set.py | 13 +++---------- tests/pyathena/polars/test_result_set.py | 12 +++--------- tests/pyathena/test_result_set.py | 6 ++---- tests/pyathena/util.py | 13 +++++++++++++ 8 files changed, 38 insertions(+), 45 deletions(-) diff --git a/pyathena/aio/result_set.py b/pyathena/aio/result_set.py index 40ad51e9b..b5e0eb0f1 100644 --- a/pyathena/aio/result_set.py +++ b/pyathena/aio/result_set.py @@ -15,7 +15,7 @@ from pyathena.converter import Converter from pyathena.error import OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution -from pyathena.result_set import AthenaDictResultSet, AthenaResultSet +from pyathena.result_set import _FETCH_FAILED_MESSAGE, AthenaDictResultSet, AthenaResultSet from pyathena.util import RetryConfig, override if TYPE_CHECKING: @@ -23,8 +23,6 @@ _logger = logging.getLogger(__name__) -_FETCH_FAILED_MESSAGE = "AthenaAioResultSet cannot be fetched after a failed fetch." - class AthenaAioResultSet(AthenaResultSet): """Async result set that provides async fetch methods. diff --git a/pyathena/aio/util.py b/pyathena/aio/util.py index dc75b2bd0..21a3ae1ad 100644 --- a/pyathena/aio/util.py +++ b/pyathena/aio/util.py @@ -161,8 +161,7 @@ def construct() -> _ClosableT | None: if not cancelled.is_set(): constructed.append(obj) return obj - with unstoppable_retries(): - obj.close() + _close_logging_errors(obj) return None future = asyncio.get_running_loop().run_in_executor( @@ -178,7 +177,9 @@ def construct() -> _ClosableT | None: # Retrieve a later construction error, which nobody else awaits. future.add_done_callback(lambda f: f.cancelled() or f.exception()) if constructed: - threading.Thread(target=constructed[0].close, name="create_in_thread-close").start() + threading.Thread( + target=_close_logging_errors, args=(constructed[0],), name="create_in_thread-close" + ).start() raise return cast(_ClosableT, future.result()) @@ -208,7 +209,7 @@ async def run_in_thread_or_close( """ lock = threading.Lock() cancelled = threading.Event() - started = ended = False + running = False def call() -> T | None: """Run the function unless the caller was cancelled, closing the object if it was then. @@ -217,17 +218,17 @@ def call() -> T | None: The return value of the function, or None if the caller was cancelled before it started. """ - nonlocal started, ended + nonlocal running with lock: if cancelled.is_set(): return None - started = True + running = True try: with stop_retries_on(cancelled): return func(*args, **kwargs) finally: with lock: - ended = True + running = False close = cancelled.is_set() if close: _close_logging_errors(closable) @@ -238,7 +239,7 @@ def call() -> T | None: except asyncio.CancelledError: with lock: cancelled.set() - close = not started or ended + close = not running if close: threading.Thread( target=_close_logging_errors, args=(closable,), name="run_in_thread_or_close-close" diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index d6eead076..3793d832e 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -1561,14 +1561,11 @@ def _reads_parquet_in_chunks(self) -> bool: and len(set(columns)) == len(columns) ) - def _iter_parquet_chunks(self, paths: list[str]) -> Iterator[DataFrame]: + def _iter_parquet_chunks(self) -> Iterator[DataFrame]: """Read the UNLOAD result in DataFrames of ``chunksize`` rows. See ``_convert_parquet_chunk()`` for the conversion of each chunk. - - Args: - paths: The files of the data manifest, which ``_read_parquet_footers()`` - read. + ``_read_parquet_footers()`` must have read the data manifest first. Yields: The DataFrames of the result rows. @@ -1580,7 +1577,7 @@ def _iter_parquet_chunks(self, paths: list[str]) -> Iterator[DataFrame]: try: with ParquetChunkReader( self._fs, - paths, + self._data_manifest, cast(int, self._chunksize), columns=self._kwargs.get("columns"), cache_type=self._cache_type, @@ -1733,7 +1730,7 @@ def _as_pandas(self) -> TextFileReader | DataFrame | Iterator[DataFrame]: if engine != "pyarrow": raise ProgrammingError("Engine must be `pyarrow`.") self._metadata = to_column_info(schema) - return self._iter_parquet_chunks(self._data_manifest) + return self._iter_parquet_chunks() # Columns that are not top-level columns of the result are read whole, # where pandas selects nested fields or raises for missing columns. if self.is_unload: diff --git a/pyathena/result_set.py b/pyathena/result_set.py index 8d402a687..818ac1d8f 100644 --- a/pyathena/result_set.py +++ b/pyathena/result_set.py @@ -26,7 +26,7 @@ _logger = logging.getLogger(__name__) -_FETCH_FAILED_MESSAGE = "AthenaResultSet cannot be fetched after a failed fetch." +_FETCH_FAILED_MESSAGE = "The result set cannot be fetched after a failed fetch." class AthenaResultSet(CursorIterator): @@ -609,9 +609,8 @@ def _count_rows(self, count: int) -> None: Args: count: The number of rows taken. """ - rownumber = self._rownumber - if rownumber is not None: - self._rownumber = rownumber + count + if self._rownumber is not None: + self._rownumber += count def _process_metadata(self, response: dict[str, Any]) -> None: result_set = response.get("ResultSet") diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index 59c91e9dd..7edfacbcd 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -32,7 +32,7 @@ _read_csv_with_pyarrow, ) from pyathena.util import RetryConfig, retry_api_call, stop_retries_on -from tests.pyathena.util import ThrottledOperation +from tests.pyathena.util import ThrottledOperation, assert_fetches_fail def _typed_values(rows): @@ -257,7 +257,7 @@ def test_fetch_after_conversion_error(self, seconds, method): assert result_set.fetchone() == (0, pd.Timestamp(0, unit="s")) with pytest.raises(OverflowError): getattr(result_set, method)() - _assert_fetches_fail(result_set) + assert_fetches_fail(result_set) def test_fetch_after_failed_chunk_read(self): """GH-1225: after iter_chunks() fails to read a chunk, the fetch methods raise.""" @@ -271,7 +271,7 @@ def chunks(): list(result_set.iter_chunks()) with pytest.raises(ProgrammingError, match=r"cannot be read after a failed read\."): result_set.fetchone() - _assert_fetches_fail(result_set) + assert_fetches_fail(result_set) def test_fetch_after_chunks_run_out(self): """The rows of a chunk that a fetch started stay after iter_chunks() runs out.""" @@ -1278,13 +1278,6 @@ def test_read_csv_with_pyarrow_multiline_values_across_blocks(): assert df["v"].tolist() == [v for _, v in rows] -def _assert_fetches_fail(result_set): - """Assert that each fetch of a result set raises for an earlier failed fetch.""" - for fetch in (result_set.fetchone, result_set.fetchmany, result_set.fetchall): - with pytest.raises(ProgrammingError, match=r"cannot be fetched after a failed fetch\."): - fetch() - - def _fetching_result_set(result, arraysize): """Create a result set of a succeeded query whose read gives ``result``, without AWS. diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 95a68d1d5..734f64637 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -21,6 +21,7 @@ from pyathena.polars.converter import DefaultPolarsTypeConverter from pyathena.polars.result_set import AthenaPolarsResultSet, PolarsDataFrameIterator from pyathena.util import RetryConfig +from tests.pyathena.util import assert_fetches_fail _ROWS_BEFORE_FAILURE = 300_000 # A timestamp column of a CSV result file with 12 fractional digits and NULL, as text and @@ -58,13 +59,6 @@ def _read_all(result: pl.DataFrame | Iterator[pl.DataFrame]) -> pl.DataFrame: return result if isinstance(result, pl.DataFrame) else pl.concat(list(result)) -def _assert_fetches_fail(result_set): - """Assert that each fetch of a result set raises for an earlier failed fetch.""" - for fetch in (result_set.fetchone, result_set.fetchmany, result_set.fetchall): - with pytest.raises(ProgrammingError, match=r"cannot be fetched after a failed fetch\."): - fetch() - - def _fetching_result_set(chunks, names, arraysize, converters=None) -> AthenaPolarsResultSet: """Create a result set of a succeeded query whose chunks are ``chunks``, without AWS. @@ -164,7 +158,7 @@ def test_fetch_after_conversion_error(self, dates, method): assert result_set.fetchone() == (0, date(2024, 1, 1)) with pytest.raises(ValueError, match="bad"): getattr(result_set, method)() - _assert_fetches_fail(result_set) + assert_fetches_fail(result_set) def test_fetch_after_failed_chunk_read(self): """GH-1225: after iter_chunks() fails to read a chunk, the fetch methods raise.""" @@ -178,7 +172,7 @@ def chunks(): list(result_set.iter_chunks()) with pytest.raises(ProgrammingError, match=r"cannot be read after a failed read\."): result_set.fetchone() - _assert_fetches_fail(result_set) + assert_fetches_fail(result_set) def test_fetch_after_chunks_run_out(self): """The rows of a chunk that a fetch started stay after iter_chunks() runs out.""" diff --git a/tests/pyathena/test_result_set.py b/tests/pyathena/test_result_set.py index c4b73a07e..0612754f1 100644 --- a/tests/pyathena/test_result_set.py +++ b/tests/pyathena/test_result_set.py @@ -16,7 +16,7 @@ from pyathena.model import AthenaQueryExecution from pyathena.result_set import AthenaResultSet, WithFetch, WithResultSet from pyathena.util import RetryConfig -from tests.pyathena.util import get_query_results_responses +from tests.pyathena.util import assert_fetches_fail, get_query_results_responses def _row(*values): @@ -97,9 +97,7 @@ def test_fetch_after_failed_fetch(self, method, page, expected): assert result_set.fetchmany(2) == [(1,), (2,)] with pytest.raises(expected): getattr(result_set, method)() - for fetch in (result_set.fetchone, result_set.fetchmany, result_set.fetchall): - with pytest.raises(ProgrammingError, match=r"cannot be fetched after a failed fetch\."): - fetch() + assert_fetches_fail(result_set) assert result_set.connection.client.get_query_results.call_count == 2 def test_fetch_after_close(self): diff --git a/tests/pyathena/util.py b/tests/pyathena/util.py index 3f6e2bebd..8d3a08412 100644 --- a/tests/pyathena/util.py +++ b/tests/pyathena/util.py @@ -15,12 +15,14 @@ from unittest.mock import patch from zoneinfo import ZoneInfo +import pytest from botocore.config import Config from botocore.exceptions import ClientError from botocore.response import StreamingBody from jinja2 import Environment, FileSystemLoader from sqlalchemy import types +from pyathena.error import ProgrammingError from pyathena.filesystem.s3 import S3FileSystem from pyathena.filesystem.s3_async import AioS3FileSystem from pyathena.filesystem.s3_core import S3Core @@ -65,6 +67,17 @@ def get_query_results_responses(*pages): return responses +def assert_fetches_fail(result_set): + """Assert that each fetch of a result set raises for an earlier failed fetch. + + Args: + result_set: The synchronous result set. + """ + for fetch in (result_set.fetchone, result_set.fetchmany, result_set.fetchall): + with pytest.raises(ProgrammingError, match=r"cannot be fetched after a failed fetch\."): + fetch() + + def throttling_error(operation_name="GetQueryResults"): """Build the ``ThrottlingException`` error of an AWS API operation. From 416ab98bb105ee1d8ceab41c4aeccc46ecd78018 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 18:06:38 +0900 Subject: [PATCH 5/7] Describe how the asyncio cursors run fetches in threads The pandas, Polars, Arrow, and S3FS pages said that their asyncio cursors use asyncio.to_thread() for result set creation and fetches, which they no longer do, and the asyncio page now says what a cancellation does to a fetch that has not started. Refs #1206 Co-Authored-By: Claude Opus 5.5 --- docs/aio.md | 5 +++-- docs/arrow.md | 4 ++-- docs/pandas.md | 4 ++-- docs/polars.md | 4 ++-- docs/s3fs.md | 4 ++-- 5 files changed, 11 insertions(+), 10 deletions(-) diff --git a/docs/aio.md b/docs/aio.md index c4f8e9b9a..0339686b1 100644 --- a/docs/aio.md +++ b/docs/aio.md @@ -172,8 +172,9 @@ 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. After a fetch is cancelled, later fetches of its result raise `ProgrammingError`. -The fetch methods of `AioPandasCursor`, `AioArrowCursor`, `AioPolarsCursor`, and `AioS3FSCursor` run in a thread, which keeps running after a cancellation. -A cancellation also stops the retries of their requests, and the cursor drops the result, which is closed once the fetch ends in its thread. +The fetch methods of `AioPandasCursor`, `AioArrowCursor`, `AioPolarsCursor`, and `AioS3FSCursor` run in a thread. +A cancellation skips a fetch that has not started, and a started fetch keeps running in its thread without the retries of its requests. +The cursor drops the result, which is closed once its fetch is no longer running. (aio-dict-cursor)= diff --git a/docs/arrow.md b/docs/arrow.md index 649290dba..76b3e9bde 100644 --- a/docs/arrow.md +++ b/docs/arrow.md @@ -485,8 +485,8 @@ cursor = connect( ## AioArrowCursor AioArrowCursor is a native asyncio cursor that returns results as Apache Arrow Tables. -Unlike AsyncArrowCursor which uses `concurrent.futures`, this cursor uses -`asyncio.to_thread()` for both result set creation and fetch operations, +Unlike AsyncArrowCursor which uses `concurrent.futures`, this cursor runs +both result set creation and fetch operations in threads, keeping the event loop free. ```python diff --git a/docs/pandas.md b/docs/pandas.md index 67b8fd3ca..88a16b404 100644 --- a/docs/pandas.md +++ b/docs/pandas.md @@ -825,8 +825,8 @@ cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/", ## AioPandasCursor AioPandasCursor is a native asyncio cursor that returns results as pandas DataFrames. -Unlike AsyncPandasCursor which uses `concurrent.futures`, this cursor uses -`asyncio.to_thread()` for both result set creation and fetch operations, +Unlike AsyncPandasCursor which uses `concurrent.futures`, this cursor runs +both result set creation and fetch operations in threads, keeping the event loop free. ```python diff --git a/docs/polars.md b/docs/polars.md index 762fad79c..78de012ee 100644 --- a/docs/polars.md +++ b/docs/polars.md @@ -618,8 +618,8 @@ for chunk in result_set.iter_chunks(): ## AioPolarsCursor AioPolarsCursor is a native asyncio cursor that returns results as Polars DataFrames. -Unlike AsyncPolarsCursor which uses `concurrent.futures`, this cursor uses -`asyncio.to_thread()` for both result set creation and fetch operations, +Unlike AsyncPolarsCursor which uses `concurrent.futures`, this cursor runs +both result set creation and fetch operations in threads, keeping the event loop free. ```python diff --git a/docs/s3fs.md b/docs/s3fs.md index 5ecbd98d3..24769f469 100644 --- a/docs/s3fs.md +++ b/docs/s3fs.md @@ -410,8 +410,8 @@ cursor.cancel(query_id) ## AioS3FSCursor AioS3FSCursor is a native asyncio cursor that uses the same lightweight CSV parsing as S3FSCursor. -Unlike AsyncS3FSCursor which uses `concurrent.futures`, this cursor uses -`asyncio.to_thread()` for both result set creation and fetch operations, +Unlike AsyncS3FSCursor which uses `concurrent.futures`, this cursor runs +both result set creation and fetch operations in threads, keeping the event loop free. Since `AthenaS3FSResultSet` lazily streams rows from S3 via a CSV reader, From 1822a288a350728721ff9ff16d4980145f4762c2 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 18:25:27 +0900 Subject: [PATCH 6/7] Break the pandas and Polars result sets on an interrupt while counting Their fetchone() counted the row it took outside the exception handler, so an interrupt at that point left later fetches enabled; the base result set counts inside it. The test of a cancellation after the function ended now waits for the call to return on a single-worker executor instead of sleeping. Refs #1206 Co-Authored-By: Claude Opus 5.5 --- pyathena/pandas/result_set.py | 4 +-- pyathena/polars/result_set.py | 4 +-- tests/pyathena/aio/test_util.py | 35 ++++++++++++++---------- tests/pyathena/pandas/test_result_set.py | 10 +++++++ tests/pyathena/polars/test_result_set.py | 10 +++++++ 5 files changed, 45 insertions(+), 18 deletions(-) diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 3793d832e..30348f9c6 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -978,11 +978,11 @@ def fetchone( raise ProgrammingError(_FETCH_FAILED_MESSAGE) try: row = next(self._row_iter, None) + if row is not None: + self._count_rows(1) except BaseException: self._fetch_failed = True raise - if row is not None: - self._count_rows(1) return row @override diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 90bcb1bb5..bd1e10505 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -634,11 +634,11 @@ def fetchone( raise ProgrammingError(_FETCH_FAILED_MESSAGE) try: row = next(self._row_iter, None) + if row is not None: + self._count_rows(1) except BaseException: self._fetch_failed = True raise - if row is not None: - self._count_rows(1) return row @override diff --git a/tests/pyathena/aio/test_util.py b/tests/pyathena/aio/test_util.py index 98f09ea19..66a8957a0 100644 --- a/tests/pyathena/aio/test_util.py +++ b/tests/pyathena/aio/test_util.py @@ -247,28 +247,35 @@ def func(): async def test_run_in_thread_or_close_cancelled_after_function(): """GH-1206: an object whose function ended before the cancellation is closed in a new thread.""" + loop = asyncio.get_running_loop() + # With one worker, a later job starts only after the function's call has returned. + executor = ThreadPoolExecutor(max_workers=1) + loop.set_default_executor(executor) resource = _Resource() threads = [] - started, release, returned = threading.Event(), threading.Event(), threading.Event() + started, release = threading.Event(), threading.Event() def func(): threads.append(threading.get_ident()) started.set() assert release.wait(EVENT_TIMEOUT) - returned.set() - - task = asyncio.create_task(run_in_thread_or_close(resource, func)) - assert await asyncio.to_thread(started.wait, EVENT_TIMEOUT) - # Blocking the event loop keeps the task from resuming after the function ends. - release.set() - assert returned.wait(EVENT_TIMEOUT) - time.sleep(0.1) - task.cancel() - with pytest.raises(asyncio.CancelledError): - await task - assert await asyncio.to_thread(resource.closed.wait, EVENT_TIMEOUT) - assert resource.closed_in not in (None, threads[0], threading.get_ident()) + try: + task = asyncio.create_task(run_in_thread_or_close(resource, func)) + await asyncio.sleep(0) + # From here, blocking the event loop keeps the task from resuming after the + # call returns, which the next job of the single worker waits for. + assert started.wait(EVENT_TIMEOUT) + release.set() + assert executor.submit(lambda: None).result(EVENT_TIMEOUT) is None + task.cancel() + with pytest.raises(asyncio.CancelledError): + await task + assert resource.closed.wait(EVENT_TIMEOUT) + assert resource.closed_in not in (None, threads[0], threading.get_ident()) + finally: + release.set() + executor.shutdown(wait=True) async def test_run_in_thread_or_close_cancelled_before_function(): diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index 7edfacbcd..e1f3b47de 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -259,6 +259,16 @@ def test_fetch_after_conversion_error(self, seconds, method): getattr(result_set, method)() assert_fetches_fail(result_set) + def test_fetchone_interrupted_while_counting(self): + """GH-1225: an interrupt after fetchone() took its row breaks the result set.""" + result_set = _fetching_result_set(pd.DataFrame({"n": [0, 1]}), arraysize=3) + with ( + patch.object(result_set, "_count_rows", side_effect=KeyboardInterrupt), + pytest.raises(KeyboardInterrupt), + ): + result_set.fetchone() + assert_fetches_fail(result_set) + def test_fetch_after_failed_chunk_read(self): """GH-1225: after iter_chunks() fails to read a chunk, the fetch methods raise.""" diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 734f64637..e90f1f79f 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -160,6 +160,16 @@ def test_fetch_after_conversion_error(self, dates, method): getattr(result_set, method)() assert_fetches_fail(result_set) + def test_fetchone_interrupted_while_counting(self): + """GH-1225: an interrupt after fetchone() took its row breaks the result set.""" + result_set = _fetching_result_set([pl.DataFrame({"n": [0, 1]})], ["n"], arraysize=3) + with ( + patch.object(result_set, "_count_rows", side_effect=KeyboardInterrupt), + pytest.raises(KeyboardInterrupt), + ): + result_set.fetchone() + assert_fetches_fail(result_set) + def test_fetch_after_failed_chunk_read(self): """GH-1225: after iter_chunks() fails to read a chunk, the fetch methods raise.""" From 4677cc77df83f8eeb8589836c087dcafdaaa60cd Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 17 Oct 2026 18:44:49 +0900 Subject: [PATCH 7/7] Expect ProgrammingError after a failed binary CSV chunk read The PandasCursor test of a failed chunk read expected later reads to stop and fetchone() to return None, the behavior that #1225 replaced. AWS CI failed on it. Refs #1206 Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/pandas/test_cursor.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/tests/pyathena/pandas/test_cursor.py b/tests/pyathena/pandas/test_cursor.py index da7b142ef..1058386de 100644 --- a/tests/pyathena/pandas/test_cursor.py +++ b/tests/pyathena/pandas/test_cursor.py @@ -160,11 +160,10 @@ def test_binary_csv_stream_close_on_read_error(self, pandas_cursor, engine, read read_chunk() assert stream.closed assert source.closed - with pytest.raises(StopIteration): - next(chunks) - with pytest.raises(StopIteration): - chunks.get_chunk() - assert pandas_cursor.fetchone() is None + # GH-1225: reads and fetches after the failed read raise. + for later in (lambda: next(chunks), chunks.get_chunk, pandas_cursor.fetchone): + with pytest.raises(ProgrammingError, match=r"after a failed read\."): + later() def test_binary_converter_override(self, pandas_cursor): pandas_cursor.execute(