Repository navigation
Stop the retries of closed result sets and cancelled asyncio calls - #1223
Conversation
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>
|
|
||
| def call() -> Any: | ||
| """Make the call, ignoring the stop scopes unless it is stoppable.""" | ||
| with contextlib.nullcontext() if stoppable else shield_retries(): |
There was a problem hiding this comment.
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:
- Cleanup requests lost their retries on cancellation:
StopQueryExecution(pyathena/aio/common.py_cancel),StopCalculationExecutionandTerminateSession(pyathena/aio/spark/cursor.py). Scenario:async with AioSparkCursor()insideasyncio.timeout()with a throttledTerminateSessionleaves the session running. Repaired:async_retry_api_call(stoppable=False), which runs the retries inshield_retries()(this line); teststest_cancel_cancelled_while_retrying,test_async_retry_api_call_not_stoppable[_in_stopped_scope]. - 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, oritertuples(), which ends after an exception), so later fetches lose or misalign rows. Repaired with the maintainer's decision: the fetches stay onasyncio.to_thread(); the cancelled-fetch contract moves to Break result sets after failed or cancelled fetches and drop the cross-thread close guards #1206. - 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.
- Missing docstring on
create_in_thread.construct(); rationale aside indocs/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_whenpolls every 0.1 s during a backoff wait (agreed design);run_in_thread()stays public as the primitive for later steps.
|
|
||
|
|
||
| @contextmanager | ||
| def shield_retries() -> Iterator[None]: |
There was a problem hiding this comment.
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_onis 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, andtest_cancel_cancelled_while_retryingfail without their fixes; the Retries continue after a result set is closed or an asyncio API call is cancelled #1202 reproductions report 0 requests afterclose()and 1 attempt after cancellation atdfbf5d8e. Local runs are Python 3.15 only; the version matrix and the AWS jobs run in CI on Ready. Unchanged docs mentioningasyncio.to_thread()(aio.md, backend pages) describe pre-existing creation/fetch paths and are not made false by this change.
There was a problem hiding this comment.
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 Claudeclaude-fable-5-1(max profile, effort high, plan mode). Both advised dropping "shield", which evokesasyncio.shield()and is ambiguous about what is protected, and keepingstop_retries_on(). Codex proposedignore_retry_stop_events()withstoppablerenamed torespect_retry_stop_events; fable proposedunstoppable_retries()withstoppablekept. 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'sstop_whenstill 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>
| *args: Any, | ||
| stop_on: Callable[[BaseException], bool] | None = None, | ||
| stop_when: Callable[[], bool] | None = None, | ||
| stoppable: bool = True, |
There was a problem hiding this comment.
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):
- A second cancellation stops the retries of the shielded
kill_on_interruptstart (pyathena/aio/util.py:107,aio/common.py:147,aio/spark/cursor.py:158), contradictingdocs/aio.md:172. - Synchronous cleanup requests are not shielded from an enclosing stop scope (
spark/common.py:300,common.py:1418,spark/common.py:441). - 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.
| 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. |
There was a problem hiding this comment.
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.
| logger=_logger, | ||
| **request, | ||
| ) | ||
| with shield_retries(): |
There was a problem hiding this comment.
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.
| 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. |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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>
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()withcopy_context().run) inherit them.stop_whenis an explicit condition for a caller that owns the stop, such as a result set.retry_api_call()makes no further attempt, ends its backoff wait, and raises the error of the last attempt. The wait usesEvent.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'sstop_when, still stop the retries.time.sleepare unchanged.2. Closed result sets (
pyathena/result_set.py,pyathena/aio/result_set.py)The page requests pass
stop_when=lambda: self.is_closed. Aclose()during the retries of a page fetch ends them, and the fetch returns nothing, as afterclose(). A whole read through theGetQueryResultsfallback (_iter_all_row_pages(), used by the pandas, Arrow, Polars, and S3FS result sets) raisesProgrammingError("AthenaResultSet is closed.")instead ofOperationalErrorafter the remaining attempts, and also checksis_closedafter each page, so that aclose()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)run_in_thread(): likeasyncio.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, andTerminateSession, pass the newstoppable=False: their retries run inunstoppable_retries()and keep going after a cancellation, as they did withasyncio.to_thread(). Their synchronous versions (pyathena/common.py,pyathena/spark/common.py) run inunstoppable_retries()too, so that a stop scope of their caller, such as a construction increate_in_thread(), does not stop them.create_in_thread()runs the construction in a scope of its existingcancelledevent, and runs the close of an object constructed after the cancellation inunstoppable_retries()._read_parquet_footers()) run in copies of their caller's context, so that a cancelled construction stops them too.Not stopped:
fetchone()/fetchmany()/fetchall()ofAioPandasCursor,AioArrowCursor,AioPolarsCursor, andAioS3FSCursorstay onasyncio.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.StartQueryExecution/StartCalculationExecutionofkill_on_interruptkeep 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, andasyncio.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.mdstates this.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 increate_in_thread(), whose object is closed on cancellation. The long-lived executors are handled with Cancelled AioS3FileSystem operations keep retrying in executor threads #1203.S3ThreadPoolExecutordoes 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:
GetQueryResultsfallback raisesProgrammingError(closed) whenclose()lands during its requests: instead ofOperationalErrorafter the remaining attempts, or of returning the page of a request that succeeded.close(), or that finds the result set closed just before its request, returns nothing instead of raisingOperationalErrororProgrammingError.AioCursorpages, and result set construction increate_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 ofexecute()while it waits for a shielded start request stops that request's retries.WHY
Part of #1205 (step 1). Closes #1202.
The design (a
ContextVarholding per-call events, set only in the context copied into the worker thread;stop_whenfor close; protected cleanup) was agreed in #1205. The cleanup helper was first namedshield_retries()and renamed tounstoppable_retries()in review, since "shield" read asasyncio.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 docsof the working tree: no warning for the changed files;stop_retries_onandunstoppable_retriesrender in the API reference. (just docs buildbuilds 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_scopeat090655c7: 1327 passed, 12 skipped, 1 failed (test_fetch_all_rows_as_csv, whose result set built with__new__lacked the_connectionthat the new closed check reads). After the test-only fix at the head (6fd37112),tests/pyathena/test_result_set.py: 29 passed.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_blockTestAthenaResultSet.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'stest_fetch_with_close_during_retries: all 14 fail withoutstop_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_threadTestAioCursor.test_cancel_cancelled_while_retrying:StopQueryExecutionkeeps its retries; fails withoutstoppable=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.close()and 1 attempt after the cancellation (7 before).The AWS jobs run when the PR is marked Ready.
🤖 Generated with Claude Code