Retries continue after a result set is closed or an asyncio API call is cancelled #1202

Part of #1205 (step 1).

Problem

retry_api_call() keeps retrying a throttled request after the work it serves is abandoned. The retries continue for up to RetryConfig.attempt attempts and the exponential waits between them (8 attempts and about 127 seconds plus jitter with the defaults).

1. Page requests of a closed result set. If close() runs while AthenaResultSet._fetch() or AthenaAioResultSet._async_fetch() is retrying a throttled GetQueryResults request, the remaining attempts are still sent. The fetch then raises OperationalError, although a closed result set otherwise fetches nothing.
2. Whole reads through the GetQueryResults fallback. _fetch_all_rows() and _fetch_all_rows_as_csv() are used by the pandas, Arrow, Polars and S3FS result sets when the output location is unavailable, for example with managed query result storage. They behave the same way after close() and then raise OperationalError instead of reporting that the result set is closed.
3. Cancelled asyncio API calls. async_retry_api_call() runs the whole retry loop in asyncio.to_thread(). When the awaiting task is cancelled, for example by asyncio.wait_for(), the thread keeps sending the remaining attempts and occupies a default-executor thread for the backoff waits. This affects every asyncio Athena and Spark API call, including GetQueryResults, GetQueryExecution, and the metadata calls.

Expected:

• No further attempt is made, and no backoff wait continues, once the result set is closed or the awaiting task is cancelled.
• A closed result set's page fetch returns None / [].
• A whole read of a closed result set raises ProgrammingError saying that the result set is closed.

Not affected:

• Retries of StartQueryExecution in an asyncio cursor with kill_on_interrupt. That request runs in a task shielded from cancellation so that the started query can be stopped, and the request is not cancelled itself.
• The sync cursor's status polling. Those retries obtain the state that the poll needs.

Reproduction

No AWS calls. GetQueryResults is stubbed to raise ThrottlingException, and RetryConfig(multiplier=0) removes the waits so the attempts can be counted.

import asyncio
import time
from unittest.mock import MagicMock

from botocore.exceptions import ClientError

from pyathena.aio.util import async_retry_api_call
from pyathena.converter import DefaultTypeConverter
from pyathena.model import AthenaQueryExecution
from pyathena.result_set import AthenaResultSet
from pyathena.util import RetryConfig

THROTTLE = ClientError({"Error": {"Code": "ThrottlingException", "Message": "x"}}, "GetQueryResults")
META = {"ColumnInfo": [{"Name": "a", "Type": "integer"}]}
calls = []


def get_query_results(**kwargs):
    calls.append(kwargs.get("NextToken"))
    if len(calls) == 1:
        rows = [{"Data": [{"VarCharValue": "1"}]}]
        return {"ResultSet": {"Rows": rows, "ResultSetMetadata": META}, "NextToken": "t1"}
    if len(calls) == 2:
        result_set.close()  # another thread closes the result set during the request
    raise THROTTLE


connection = MagicMock()
connection.client.get_query_results.side_effect = get_query_results
result_set = AthenaResultSet(
    connection=connection,
    converter=DefaultTypeConverter(),
    query_execution=MagicMock(state=AthenaQueryExecution.STATE_SUCCEEDED, query_id="q"),
    arraysize=1,
    retry_config=RetryConfig(multiplier=0),
)
print(result_set.fetchone())  # (1,)
try:
    result_set.fetchone()
except Exception as e:
    print(type(e).__name__)  # OperationalError
print(len(calls) - 2)  # 7 requests after close()


async def cancelled_call():
    attempts = []

    def operation(**kwargs):
        attempts.append(1)
        time.sleep(0.02)
        raise THROTTLE

    task = asyncio.create_task(async_retry_api_call(operation, RetryConfig(multiplier=0)))
    await asyncio.sleep(0.01)
    task.cancel()
    await asyncio.gather(task, return_exceptions=True)
    print(len(attempts))  # 1 at cancellation
    await asyncio.sleep(0.5)
    print(len(attempts))  # 8: the thread kept retrying


asyncio.run(cancelled_call())

The asyncio result set (AthenaAioResultSet) behaves like case 1, and _fetch_all_rows() behaves like case 2. Both show 7 requests after close(), followed by OperationalError.

Environment

• PyAthena master at b2d7aa2 (the retry behavior is the same in v4.1.0)
• Python 3.15, tenacity as resolved in uv.lock

Proposed fix (optional)

• Let retry_api_call() take a stop condition. Once the condition holds, it makes no further attempt, ends the backoff wait early, and raises the last error.
• The result sets pass is_closed as the condition. A page fetch of a closed result set returns nothing, and a request that fails or stops on a closed result set raises ProgrammingError for the closed result set.
• async_retry_api_call() sets the condition when the awaiting task is cancelled, and combines it with any condition the caller passes.

Validate with unit tests for retry_api_call() / async_retry_api_call() and for the sync and asyncio result sets, using stubbed throttling responses.

The same behavior in AioS3FileSystem is tracked in #1203, because its S3 retries run through its own executor.

Scope extended by #1205

This issue is step 1 of #1205. Besides the cases above, the cooperative stop also covers the following:

• create_in_thread() (pyathena/aio/util.py:72), cancelled. Construction keeps running, including GetQueryResults retries (7 more requests) and the pandas Parquet footer reads. The aio pandas, Arrow, Polars, and S3FS cursors use it.
• The asyncio to_thread() fetches (pyathena/aio/common.py:807, 831, 851), cancelled. Their retries keep running.

Design (consulted with Claude fable-5-1 and Codex; see #1205):

• A ContextVar holds a per-call threading.Event. It is set only in the context copied into the worker thread, never in the task's context.
• The asyncio entry points set the event on cancellation.
• retry_api_call() honors it, as well as an explicit stop_when for a closed result set. It makes no further attempt, ends the backoff wait early, and raises the last error.
• Protected cleanup and the shielded starts mask the signal.