Skip to content

Stop the retries of closed result sets and cancelled asyncio calls - #1223

Merged
laughingman7743 merged 9 commits into
masterfrom
fix/1202-retry-stop
Oct 17, 2026
Merged

laughingman7743 merged 9 commits into
masterfrom
fix/1202-retry-stop

Conversation

@laughingman7743

@laughingman7743 laughingman7743 commented Oct 16, 2026 •

Copy link
Copy Markdown
Member

WHAT

Step 1 of #1205 for the Athena API calls and the asyncio entry points: retry_api_call() can now be stopped, and closed result sets and cancelled asyncio calls stop it. The same step for the S3 filesystem (#1203) follows in a separate PR, once #1192, which rewrites the same executor, range-read, and copy code, has settled.

1. A cooperative stop for retry_api_call() (pyathena/util.py)

  • stop_retries_on(event) opens a stop scope: a context variable holds the events of the enclosing scopes, so the threads that run a copy of the context (asyncio.to_thread(), run_in_executor() with copy_context().run) inherit them.
  • stop_when is an explicit condition for a caller that owns the stop, such as a result set.
  • Once an event is set or the condition holds, retry_api_call() makes no further attempt, ends its backoff wait, and raises the error of the last attempt. The wait uses Event.wait() for a single event, and polls every 0.1 seconds otherwise. A first attempt is always made.
  • unstoppable_retries() clears the scopes for cleanup that must run to its end; a scope opened inside it, and a call's stop_when, still stop the retries.
  • Without a scope or condition, tenacity's own sleep is kept, so the existing callers and the tests that patch time.sleep are unchanged.

2. Closed result sets (pyathena/result_set.py, pyathena/aio/result_set.py)

The page requests pass stop_when=lambda: self.is_closed. A close() during the retries of a page fetch ends them, and the fetch returns nothing, as after close(). A whole read through the GetQueryResults fallback (_iter_all_row_pages(), used by the pandas, Arrow, Polars, and S3FS result sets) raises ProgrammingError("AthenaResultSet is closed.") instead of OperationalError after the remaining attempts, and also checks is_closed after each page, so that a close() during a request that succeeds raises instead of returning the page. That error is not logged with _logger.exception().

3. Cancelled asyncio calls (pyathena/aio/util.py, pyathena/aio/common.py)

  • New run_in_thread(): like asyncio.to_thread(), but in a stop scope whose event a cancellation of the caller sets.
  • async_retry_api_call(), which every asyncio Athena and Spark API call uses, runs through it. The requests that clean up, StopQueryExecution, StopCalculationExecution, and TerminateSession, pass the new stoppable=False: their retries run in unstoppable_retries() and keep going after a cancellation, as they did with asyncio.to_thread(). Their synchronous versions (pyathena/common.py, pyathena/spark/common.py) run in unstoppable_retries() too, so that a stop scope of their caller, such as a construction in create_in_thread(), does not stop them.
  • create_in_thread() runs the construction in a scope of its existing cancelled event, and runs the close of an object constructed after the cancellation in unstoppable_retries().
  • The pandas footer reads (_read_parquet_footers()) run in copies of their caller's context, so that a cancelled construction stops them too.

Not stopped:

  • The fetchone() / fetchmany() / fetchall() of AioPandasCursor, AioArrowCursor, AioPolarsCursor, and AioS3FSCursor stay on asyncio.to_thread(). Their retries are mostly S3 reads of a reader that later fetches continue (a CSV stream, or the row iteration of pandas and Polars, which ends after an exception). A stopped read would leave it part-way, so later fetches would lose or misalign rows. How a cancelled fetch ends is left to the operation guard of Break result sets after failed or cancelled fetches and drop the cross-thread close guards #1206.
  • The shielded StartQueryExecution / StartCalculationExecution of kill_on_interrupt keep their retries after a cancellation without extra code: PyAthena sets the events only in the context copies of the threads, never in a task's context, and asyncio.shield() keeps the start task from being cancelled. A second cancellation during that wait cancels the start task, which now also stops its retries (before, the thread kept retrying and could start a query that nobody stopped); docs/aio.md states this.
  • Read-ahead work that outlives its call (Arrow's Parquet prefetch, ReadAheadStream) runs in plain thread pools that do not copy the submitter's context, so a stop does not reach it. On builds where a new thread inherits the context of its creator (free-threaded builds, or -X thread_inherit_context=1), a worker thread started within a scope inherits it; in this PR the scopes cover only single API calls and constructions in create_in_thread(), whose object is closed on cancellation. The long-lived executors are handled with Cancelled AioS3FileSystem operations keep retrying in executor threads #1203.
  • The sync S3ThreadPoolExecutor does not copy the context yet (Cancelled AioS3FileSystem operations keep retrying in executor threads #1203), so the multi-range reads of a large S3 object do not see the stop in this PR.

Behavior changes for the release notes:

  • A whole read of a closed result set through the GetQueryResults fallback raises ProgrammingError (closed) when close() lands during its requests: instead of OperationalError after the remaining attempts, or of returning the page of a request that succeeded.
  • A page fetch whose request fails after a close(), or that finds the result set closed just before its request, returns nothing instead of raising OperationalError or ProgrammingError.
  • Cancelling an asyncio call stops the retries of its Athena requests in the thread (status polling, metadata, AioCursor pages, and result set construction in create_in_thread()); a request already sent still runs to its end. The fetches of the other asyncio cursors and the cleanup requests are unchanged. A second cancellation of execute() while it waits for a shielded start request stops that request's retries.

WHY

Part of #1205 (step 1). Closes #1202.

The design (a ContextVar holding per-call events, set only in the context copied into the worker thread; stop_when for close; protected cleanup) was agreed in #1205. The cleanup helper was first named shield_retries() and renamed to unstoppable_retries() in review, since "shield" read as asyncio.shield(). #1203, the S3 filesystem part of the same step, was split off because PR #1192 rewrites the same code (S3Executor.submit_to(), S3AioExecutor.submit(), _fetch_range(), the async multipart copy) and already cancels the queued range reads and part copies.

TEST

Tested on Python 3.15 (macOS).

  • just lint, just docs lint: passed at the head. sphinx-build docs of the working tree: no warning for the changed files; stop_retries_on and unstoppable_retries render in the API reference. (just docs build builds only the committed versions.)
  • uv run --env-file .env pytest -n 4 tests/pyathena/test_util.py tests/pyathena/test_result_set.py tests/pyathena/aio/ tests/pyathena/test_cursor.py tests/pyathena/pandas/ tests/pyathena/arrow/test_result_set.py tests/pyathena/s3fs/ tests/pyathena/spark/test_spark_cursor.py::TestSparkCursor::test_cleanup_in_stopped_scope at 090655c7: 1327 passed, 12 skipped, 1 failed (test_fetch_all_rows_as_csv, whose result set built with __new__ lacked the _connection that the new closed check reads). After the test-only fix at the head (6fd37112), tests/pyathena/test_result_set.py: 29 passed.
  • Not run locally: the Polars, Spark (beyond the offline tests above), SQLAlchemy, and filesystem suites, and Python versions other than 3.15; CI runs them when the PR is Ready.
  • New regression tests, using stubbed throttling responses (no AWS) and a backoff of at least EVENT_TIMEOUT (10) seconds that only the stop can shorten:
    • test_retry_api_call_stopped_by_event[request|backoff], ..._stopped_by_stop_when[...], ..._stopped_before_call, ..._stopped_by_enclosing_scope, ..._unstoppable, ..._scope_ends_with_block
    • TestAthenaResultSet.test_fetch_with_close_during_retries[fetchone|fetchmany|fetchall × request|backoff], test_fetch_all_rows_with_close_during_retries[request|backoff], and the asyncio result set's test_fetch_with_close_during_retries: all 14 fail without stop_when.
    • test_async_retry_api_call_cancelled[request|backoff], test_async_retry_api_call_not_stoppable, ..._not_stoppable_in_stopped_scope, test_create_in_thread_cancelled_stops_retries[request|backoff], test_create_in_thread_cleanup_not_stopped, test_run_in_thread
    • TestAioCursor.test_cancel_cancelled_while_retrying: StopQueryExecution keeps its retries; fails without stoppable=False.
    • TestAioCursor.test_execute_cancelled_while_start_is_retried: the shielded start keeps its retries; test_execute_cancelled_again_while_start_is_retried: a second cancellation stops them.
    • TestCursor.test_cancel_in_stopped_scope, TestSparkCursor.test_cleanup_in_stopped_scope[_terminate_session_by_id|_cancel], TestAthenaResultSet.test_fetch_all_rows_with_close_during_request[...]: all fail without their fixes.
    • TestAthenaPandasResultSet.test_read_parquet_footers_stopped: fails without the context copies.
  • The reproductions of Retries continue after a result set is closed or an asyncio API call is cancelled #1202 report 0 requests after close() and 1 attempt after the cancellation (7 before).

The AWS jobs run when the PR is marked Ready.

🤖 Generated with Claude Code

laughingman7743 and others added 6 commits October 17, 2026 01:18
A stop scope (stop_retries_on()) carries a threading.Event in a context
variable, which threads that run a copy of the context inherit, and
stop_when is an explicit condition. Once either holds, retry_api_call()
makes no further attempt, ends its backoff wait, and raises the error of
the last attempt. shield_retries() keeps the scopes around a cleanup from
stopping it. Without a scope or condition, tenacity's own sleep is kept.

Refs #1202

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The result sets pass stop_when=is_closed to their GetQueryResults
requests. A close() during the retries of a page fetch now ends them, and
the fetch returns nothing, as after close(); a whole read through the
GetQueryResults fallback raises ProgrammingError for the closed result set
instead of OperationalError after the remaining attempts.

Refs #1202

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
run_in_thread() runs a blocking function as asyncio.to_thread() does, in a
stop scope whose event a cancellation of the caller sets.
async_retry_api_call() and the to_thread() fetches of the asyncio cursors
use it, and create_in_thread() stops a running construction with its
existing cancellation event, while the close of an object constructed
after the cancellation is shielded. The shielded start requests of
kill_on_interrupt keep their retries, since the event exists only in the
context copies of the threads. The pandas footer reads run in a copy of
the context of their caller, so that a cancelled construction stops them
too.

Closes #1202

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
StopQueryExecution, StopCalculationExecution, and TerminateSession clean
up, so async_retry_api_call(stoppable=False) keeps their retries after a
cancellation, as asyncio.to_thread() did. The fetches of the asyncio
pandas, Arrow, Polars, and S3FS cursors go back to asyncio.to_thread():
their retries are mostly S3 reads of a reader that later fetches continue,
which a stopped read would leave part-way; #1206 defines a cancelled fetch.

Also documents construct() and reuses the shared throttling helper.

Refs #1202

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
async_retry_api_call(stoppable=False) now runs its retries in
shield_retries(), so a stop scope of the caller does not stop a cleanup
request either. retry_api_call() defines its stop hooks at function level,
and the pandas footer reads pass each copied context with a lambda.

The tests share one throttling stub that stops its retries
(StoppingOperation), reuse ThrottledOperation, and wait with wait_until().

Refs #1202

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Comment thread pyathena/aio/util.py Outdated

def call() -> Any:
"""Make the call, ignoring the stop scopes unless it is stoppable."""
with contextlib.nullcontext() if stoppable else shield_retries():

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Self-review round 1 (behavior and implementation): FINDINGS, repaired

Range: git diff b2d7aa24b00a204e9e150952c708e3bb1126f9e2..dfbf5d8e49cd9d5ef09c2180833d62a502ce8cde (first pass at 4b584efed99e35377f8d9edd7b8d632c4aa48611; repairs in 78f65558, 27641102, dfbf5d8e).
Covered: pyathena/util.py (stop scopes, stop_when, interruptible sleep, tenacity hooks), pyathena/aio/util.py, pyathena/aio/common.py, pyathena/aio/spark/cursor.py, both result sets (_get_query_results, _fetch, _iter_all_row_pages, whole reads), the pandas footer reads, docs, and the tests; callers traced into the pandas/Arrow/Polars/S3FS result sets (their _fetch() overrides do not use _get_query_results). Passes: manual, /code-review (high), /simplify.

Findings and repairs:

  1. Cleanup requests lost their retries on cancellation: StopQueryExecution (pyathena/aio/common.py _cancel), StopCalculationExecution and TerminateSession (pyathena/aio/spark/cursor.py). Scenario: async with AioSparkCursor() inside asyncio.timeout() with a throttled TerminateSession leaves the session running. Repaired: async_retry_api_call(stoppable=False), which runs the retries in shield_retries() (this line); tests test_cancel_cancelled_while_retrying, test_async_retry_api_call_not_stoppable[_in_stopped_scope].
  2. Stopping the to_thread() fetches of the asyncio pandas/Arrow/Polars/S3FS cursors stops S3 reads of a reader that later fetches continue (a CSV stream mid-record, or itertuples(), which ends after an exception), so later fetches lose or misalign rows. Repaired with the maintainer's decision: the fetches stay on asyncio.to_thread(); the cancelled-fetch contract moves to Break result sets after failed or cancelled fetches and drop the cross-thread close guards #1206.
  3. On free-threaded builds a thread created within a scope inherits it; moot for this PR once (2) is applied (scopes cover single API calls and constructions closed on cancellation). Recorded in the PR body; the long-lived executors are handled with Cancelled AioS3FileSystem operations keep retrying in executor threads #1203.
  4. Missing docstring on create_in_thread.construct(); rationale aside in docs/aio.md; duplicated throttling stubs and polling loops in the tests. Repaired.

Deferred with reasons:

  • A construction stopped in create_in_thread() fails like one whose retries ran out; release of what a constructor opened on failure is Close objects built for cancelled asyncio constructions #1208 (step 4).
  • stop_when polls every 0.1 s during a backoff wait (agreed design); run_in_thread() stays public as the primitive for later steps.

Comment thread pyathena/util.py Outdated


@contextmanager
def shield_retries() -> Iterator[None]:

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Self-review round 2 (claims, callers, operations): FINDINGS, repaired

Range: git diff b2d7aa24b00a204e9e150952c708e3bb1126f9e2..dfbf5d8e49cd9d5ef09c2180833d62a502ce8cde.
Claims checked against code: every asyncio Athena/Spark API call goes through async_retry_api_call() (Glue's to_thread request has botocore retries only); without a scope or stop_when, tenacity's own sleep is kept (stop_kwargs empty), so callers and time.sleep patches are unchanged; a first attempt is always made; the fallback whole reads of pandas/Arrow/Polars/S3FS all go through _iter_all_row_pages(); the closed ProgrammingError is not logged; the footer reads open and read through PyAthena's S3FileSystem in the copied context; retry_api_call/async_retry_api_call add keyword-only arguments after *args, which no AWS parameter (CapWords) can collide with; RetryCallState, sleep= and before_sleep= exist in tenacity 8.0.1; parenthesized with needs Python 3.10 (floor 3.11).
Operator view: the stop only removes attempts and shortens waits; botocore's own retries inside an attempt are unchanged, as documented.

Corrections:

  • shield_retries() named a multipart abort, which this PR does not touch, as its example; now names this PR's cleanups (this line).
  • PR body: "events never in a task's context" narrowed to PyAthena's own code (stop_retries_on is public); "read-ahead does not inherit the scope" qualified for builds where new threads inherit their creator's context.
    Evidence: the new result-set tests (14), the pandas footer test, and test_cancel_cancelled_while_retrying fail without their fixes; the Retries continue after a result set is closed or an asyncio API call is cancelled #1202 reproductions report 0 requests after close() and 1 attempt after cancellation at dfbf5d8e. Local runs are Python 3.15 only; the version matrix and the AWS jobs run in CI on Ready. Unchanged docs mentioning asyncio.to_thread() (aio.md, backend pages) describe pre-existing creation/fetch paths and are not made false by this change.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Repair: shield_retries() renamed to unstoppable_retries() in 24b263e0, at the maintainer's request.

  • Consultation (advice only): the same neutral package went to Codex (codex exec -s read-only) and Claude claude-fable-5-1 (max profile, effort high, plan mode). Both advised dropping "shield", which evokes asyncio.shield() and is ambiguous about what is protected, and keeping stop_retries_on(). Codex proposed ignore_retry_stop_events() with stoppable renamed to respect_retry_stop_events; fable proposed unstoppable_retries() with stoppable kept. The maintainer chose the latter.
  • Docstring: the events of enclosing stop_retries_on() scopes are ignored in the block; a scope opened in the block and a call's stop_when still stop the retries.

Self-review of the repair, both perspectives: CLEAN. Behavior: a mechanical rename across pyathena/util.py, common.py, aio/util.py, spark/common.py, docs/api/utilities.rst, and tests/pyathena/test_util.py, with no remaining reference. The function is unreleased, so no alias. Claims: the docstrings and the PR body use the new name. The affected tests (150) pass, and sphinx-build renders the new entry. Relayed Codex follow-up (read-only, session 01a14782-4e38-7892-afb9-cf4ccbdc00f9, git diff 6fd37112..24b263e0): CLEAN, no reference left, docstrings accurate, no behavior change. Static review.

…page

The synchronous StopQueryExecution, StopCalculationExecution, and
TerminateSession requests run in shield_retries(), as the asyncio ones do,
so that a stop scope of their caller, such as a construction in
create_in_thread(), does not stop them. A whole read through the
GetQueryResults fallback checks is_closed after each page, so that a
close() during a request that succeeds raises ProgrammingError instead of
returning the page.

The docs state that a second cancellation stops the retries of a shielded
start request.

Refs #1202

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Comment thread pyathena/aio/util.py
*args: Any,
stop_on: Callable[[BaseException], bool] | None = None,
stop_when: Callable[[], bool] | None = None,
stoppable: bool = True,

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Independent review (relayed Codex result)

Reviewer: OpenAI Codex CLI 0.162.1 (model gpt-6-astra, reasoning effort max), codex exec -s read-only, run in a detached snapshot of the reviewed commit with the literal diff and requirements only (no PR text, commit messages, or author findings); no edits, builds, tests, GitHub, or memory access. Static review.

Initial review: git diff b2d7aa24b00a204e9e150952c708e3bb1126f9e2..dfbf5d8e49cd9d5ef09c2180833d62a502ce8cde, session 01a146ef-7e48-7b63-a343-0d5a2a0449df: FINDINGS (3, each P2, recorded inline on their repairs):

  1. A second cancellation stops the retries of the shielded kill_on_interrupt start (pyathena/aio/util.py:107, aio/common.py:147, aio/spark/cursor.py:158), contradicting docs/aio.md:172.
  2. Synchronous cleanup requests are not shielded from an enclosing stop scope (spark/common.py:300, common.py:1418, spark/common.py:441).
  3. Pre-existing: a whole read through the fallback returns the page of a successful request that close() interrupted (result_set.py:781), which the new docstrings promise to raise for.

Follow-up review of the repair git diff dfbf5d8e49cd9d5ef09c2180833d62a502ce8cde..090655c79ec7798072747337779c1e04e3103e97 (range-diff: earlier commits unchanged), session 01a146f7-ff34-7032-893a-53b5ef5600b1: CLEAN — covered the sync, future-based, and asyncio cancellation paths, Spark session start/termination and calculation cancellation, cancelled create_in_thread() constructions, the whole-read callers of all four result sets, docs, and tests. Out of scope as requested: WithAsyncFetch fetches and pyathena/filesystem/.

Self-review of the repair (both perspectives, narrow; git range-diff b2d7aa24b00a204e9e150952c708e3bb1126f9e2..dfbf5d8e49cd9d5ef09c2180833d62a502ce8cde b2d7aa24b00a204e9e150952c708e3bb1126f9e2..090655c79ec7798072747337779c1e04e3103e97): CLEAN. Behavior: the sync _cancel()/_terminate_session_by_id() keep their signatures and error translation; _iter_all_row_pages() raises before the page is parsed or yielded, so no caller converts it; the second-cancellation behavior is now tested. Claims: docs/aio.md and the PR body state the first-versus-second cancellation contract, the sync shielding, and the release-note change for whole reads. The five new tests for (2) and (3) fail without their fixes. Snapshot and PR worktree verified unchanged at 090655c7.

Comment thread docs/aio.md
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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Codex finding 1 (P2): second cancellation stops the start's retries. Verified. Decision: the behavior is kept and the claim corrected — before, the thread kept retrying after a second cancellation and could start a query that nobody stopped. This line now states that a first cancellation keeps the retries and a second one stops them; test_execute_cancelled_again_while_start_is_retried covers it. Follow-up review: CLEAN.

Comment thread pyathena/spark/common.py Outdated
logger=_logger,
**request,
)
with shield_retries():

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Codex finding 2 (P2): synchronous cleanup not shielded. Verified: no PyAthena path runs these inside a stop scope today, but create_in_thread() and run_in_thread() are public, and #1208 plans to construct Spark cursors through create_in_thread(). Repaired: StopQueryExecution (pyathena/common.py:1422), TerminateSession (here) and StopCalculationExecution (spark/common.py:447) run in shield_retries(); TestCursor.test_cancel_in_stopped_scope and TestSparkCursor.test_cleanup_in_stopped_scope[...] fail without it. Follow-up review: CLEAN.

Comment thread pyathena/result_set.py
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.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Codex finding 3 (P2, pre-existing): fallback page returned after close(). Verified; folded in because the new docstrings promise the exception. Repaired: is_closed is checked after each page here; test_fetch_all_rows_with_close_during_request[_fetch_all_rows|_iter_all_row_pages] fails without it. Release note added to the PR body. Follow-up review: CLEAN.

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Repair follow-up: the local run at 090655c7 found test_fetch_all_rows_as_csv failing with the new is_closed check, because the test builds its result set with __new__ and had no _connection. 6fd37112 gives it a MagicMock() connection (test-only); tests/pyathena/test_result_set.py passes (29). Relayed Codex follow-up (read-only, session 01a14706-2bb8-76b2-a175-67ca616c32ea, git diff 090655c7..6fd37112): CLEAN — the real page iterator and CSV serializer still run and the exact-bytes assertion is unchanged. Static review only.

The whole read now checks is_closed after each page, which reads the
connection that the test's result set, built without __init__, lacked.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
"Shield" read as asyncio.shield(), which protects a task from
cancellation, and "shield retries" could read as preventing them. The
new name states the effect in the block, pairs with stop_retries_on() and
the stoppable parameter of async_retry_api_call(), and its docstring now
names what still stops the retries: a scope opened in the block, and the
stop_when condition of a call.

Refs #1202

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

1 participant