Re-raise the interrupt after kill_on_interrupt cancellation by laughingman7743 · Pull Request #853 · pyathena-dev/PyAthena

laughingman7743 · 2026-09-26T07:59:18Z

WHAT

BaseCursor._poll() (pyathena/common.py) and AioCursorBase._poll() (pyathena/aio/common.py) now re-raise the original KeyboardInterrupt / asyncio.CancelledError after the kill_on_interrupt cancellation, whatever the query's terminal state is.
The sequence is unchanged up to that point: request StopQueryExecution, then poll until the query reaches a terminal state.
A failure of the stop request or of that wait becomes the interrupt's __cause__.
The cursor does not store the final execution and does not build a result set, even when the query ended SUCCEEDED.
query_id keeps the interrupted query's ID, as it already did on other error paths.
kill_on_interrupt=False is unchanged.
Affected cursors: every SQL cursor that waits in execute(): Cursor, DictCursor, the pandas/arrow/polars/s3fs cursors, and the Aio* cursors.
The thread-pool Async* cursors poll in worker threads, which never receive KeyboardInterrupt, so they see no practical change.
The Spark cursors override _poll() and already follow this contract since Define best-effort Spark calculation cancellation #833.
executemany() stops at the interrupted execution and re-raises, with rowcount == -1 and query_id kept (its existing BaseException handling).
Before, an interrupted execution that ended SUCCEEDED let the loop continue with the next parameters.
dbt-athena is unaffected: its current connection manager uses boto3 directly, and its legacy PyAthena cursor overrides _poll() (dbt-athena/src/dbt/adapters/athena/connections_legacy.py).
Docs: new "Query cancellation on interrupt" section in docs/usage.md and "Task cancellation" in docs/aio.md.
They also state two limits. A second interrupt during the cancellation request or the wait propagates immediately. An asyncio timeout that expires while the query is still being started leaves query_id as None, and the pending start request can still start the query.

Behavior change (release note): with the default kill_on_interrupt=True, an interrupt during execute() now propagates as KeyboardInterrupt / asyncio.CancelledError.
Before, it surfaced as OperationalError (query ended CANCELLED/FAILED) or was lost (query ended SUCCEEDED).
Callers that caught OperationalError after Ctrl-C or task cancellation need to handle the interrupt instead.

WHY

Closes #840.
Swallowing the interrupt lost Ctrl-C when the query finished first.
In asyncio, task.cancel() did not end with a cancelled task, and a timeout from asyncio.wait_for() surfaced as the query's OperationalError or a normal return instead of TimeoutError (the new timeout test fails this way on the original code).
The contract (always re-raise, no stored final execution) was agreed with the maintainer before implementation, matching #833 for Spark.

TEST

Tested commit: f528d50 (the last two commits change only docs and the two test files)

just format, just lint: passed (ruff, format check, mypy, cfn-lint, license headers).
just docs lint: 0 errors.
uv run sphinx-build -q -E docs ~/tmp/pyathena-840-docs: the new {ref} resolves; the only warnings in docs/usage.md / docs/aio.md are the two xref warnings already present on master (managed-query-result-storage, assume-role-provider).
just docs build (sphinx-multiversion) builds committed refs only, so it was not used to check the uncommitted change.
Offline tests, no AWS (--noconftest skips the session setup that creates AWS resources):
uv run --env-file .env pytest --noconftest -p no:xdist tests/pyathena/test_cursor.py tests/pyathena/aio/test_cursor.py -k "interrupt or on_poll_invoked or on_poll_none" -v → 13 passed.
- Sync: after the interrupt, the query reports RUNNING and then CANCELLED/SUCCEEDED. KeyboardInterrupt is re-raised only after the terminal state (asserted through on_poll), query_id is kept, and no result set is built. A cancel or wait failure becomes __cause__. With kill_on_interrupt=False, the interrupt propagates without a stop request.
- asyncio: the same cases via task.cancel() (plus task.cancelled()), and asyncio.wait_for() raises asyncio.TimeoutError after the stop request.
Same 11 interrupt tests on Python 3.10.16 (separate venv outside the repository): 11 passed.
Mutation check: replacing the follow-up wait with a single status request makes all 4 test_execute_kill_on_interrupt cases fail.
Regression check: with the original _poll() restored, 9 of the 11 new tests fail. The SUCCEEDED cases fail with DID NOT RAISE, the CANCELLED/failure cases with OperationalError. The two kill_on_interrupt=False tests pass on both, as expected.
Live check of the best-effort claim (one SELECT 1 with the CI account, script kept outside the repository): StopQueryExecution on a query that has already SUCCEEDED returns HTTP 200 and the state stays SUCCEEDED.
So when the query wins the race, _poll() re-raises the interrupt without a cause, as documented.
Not run locally: the AWS integration suites. They run in CI once the PR is Ready. No live interrupt test was added, to avoid extra Athena queries; the change is in the client-side exception flow only.

🤖 Generated with Claude Code

With kill_on_interrupt enabled, the shared _poll() of the SQL cursors requested StopQueryExecution after a KeyboardInterrupt or task cancellation, waited for a terminal state, and then returned the final execution instead of re-raising. Callers saw an OperationalError for a CANCELLED or FAILED query, or a normal return when the query SUCCEEDED first, and asyncio.wait_for()/asyncio.timeout() could not turn the cancellation into a timeout. Re-raise the original interrupt after the stop request and the wait, as #833 does for the Spark cursors, with a failure to stop or wait as its cause. Closes #840 Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

laughingman7743 · 2026-09-26T08:00:31Z

-        return query_execution
+                self.__poll(query_id)
+            except Exception as e:
+                raise interrupt from e


Self-review round 1: implementation behavior. Result: CLEAN

Scope: c01c56f73c7dbf973fe73b52b6093f0a32952321..2de91e768196dc061251f7609ae0d9b9be2905cd (full diff: pyathena/common.py, pyathena/aio/common.py, both test files, docs/usage.md, docs/aio.md).

Covered:

Callers of the shared _poll(): execute() of Cursor/DictCursor, pandas/arrow/polars/s3fs (sync), and AioCursor plus the aio pandas/arrow/polars/s3fs cursors. The thread-pool Async* cursors call _poll() from executor threads (pyathena/async_cursor.py:153, pyathena/pandas/async_cursor.py:131, etc.), which never receive KeyboardInterrupt, so they are unaffected. The Spark cursors override _poll(). SQLAlchemy does not call _poll() directly.

Cursor state after an interrupt: _reset_state() already cleared result_set and rowcount, query_id is set before polling, and no result set is built on the interrupt path, so nothing is left open.

executemany(): both pyathena/result_set.py:1064 and pyathena/aio/common.py:651 catch BaseException, close, and re-raise, so an interrupt stops the loop with rowcount == -1 and query_id kept. Before this change, an interrupted execution that ended SUCCEEDED let the loop continue with the next parameters.

Exception flow: the bare raise after the inner try/except Exception re-raises the outer interrupt. A second interrupt or cancellation during the wait is a BaseException, so it propagates instead of becoming a cause (same as Define best-effort Spark calculation cancellation #833).

Tests: the new tests fail on the original _poll() for the defect itself. The SUCCEEDED cases fail with DID NOT RAISE, because the fake result-set class lets the old code return normally.

No actionable findings. Round two will add the executemany() consequence to the PR description.

laughingman7743 · 2026-09-26T08:00:31Z

+        assert cursor.query_id == "query_id"
+        assert cursor.result_set is None
+
+    async def test_execute_kill_on_interrupt_timeout(self):


Self-review round 1: test determinism

The first status request blocks on an event that is never set, so the 0.01 s timeout always fires while polling.
The re-poll after the stop request returns at once, so asyncio.wait_for() sees CancelledError and raises asyncio.TimeoutError on both the 3.10 wait_for implementation and the timeout()-based one from 3.12.
asyncio.TimeoutError is used instead of the builtin TimeoutError, because the two are distinct on 3.10.

laughingman7743 · 2026-09-26T08:02:25Z

+
+With `kill_on_interrupt` enabled, which is the default, a `KeyboardInterrupt` while `execute()` waits for the query
+requests cancellation, waits until the query reaches a terminal state, and then propagates.
+Cancellation is a best-effort request, so the query can still end as `SUCCEEDED` or `FAILED`.


Self-review round 2: claims, callers, and operations. Result: FINDINGS (PR description only, repaired)

Scope: c01c56f73c7dbf973fe73b52b6093f0a32952321..2de91e768196dc061251f7609ae0d9b9be2905cd, all claims in the PR description, commit message, _poll() docstrings, docs/usage.md, and docs/aio.md.

Claims checked:

Best-effort stop can still end SUCCEEDED: measured with one SELECT 1 on the CI account. StopQueryExecution on an already SUCCEEDED query returns HTTP 200 and the state stays SUCCEEDED. So the race re-raises the interrupt without a cause, and this sentence and the _poll() docstrings hold.

"If the cancellation request fails, ... as its cause": covered by the failure[cancel] tests on both bases.

kill_on_interrupt=False keeps the query running: no stop request is made (test_execute_without_kill_on_interrupt).

Timeouts: the commit message claim holds. On the original code, asyncio.wait_for() surfaces OperationalError (the new timeout test fails this way), and CPython's Timeout.__aexit__ converts only CancelledError into TimeoutError.

Existing callers: dbt-athena's current connection manager uses boto3 directly, and its legacy PyAthena cursor overrides _poll() (dbt-athena/src/dbt/adapters/athena/connections_legacy.py:167), so it is unaffected. The kill_on_interrupt parameter descriptions in pyathena/connection.py:228 and the cursor docstrings remain accurate.

Findings, repaired in the PR description:

The WHY claimed that TaskGroup could not handle the cancellation. A TaskGroup still raises its ExceptionGroup when a sibling fails, so this is narrowed to the verified asyncio.wait_for() effect.

The description omitted the executemany() consequence: it now stops at the interrupted execution, where an interrupted execution that ended SUCCEEDED used to let the loop continue. This is added, along with the dbt-athena compatibility note and the live stop measurement.

Deferred (pre-existing, out of scope): the thread-pool Async* cursor docstrings (e.g. pyathena/arrow/async_cursor.py:91) say kill_on_interrupt cancels on keyboard interrupt, but their polling runs in executor threads, which never receive KeyboardInterrupt. This PR does not change that path.

Limit the documented asyncio timeout behavior to a timeout while the query is being waited for, note that a repeated interrupt or task cancellation skips the wait, and make the tests assert that the interrupt is re-raised only after a non-terminal poll reaches a terminal state. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

laughingman7743 · 2026-09-26T08:09:30Z

+            if not self._kill_on_interrupt:
+                raise
+            _logger.warning("Query canceled by user.")
+            try:


Independent review (relayed): Codex CLI 0.157.0, model gpt-6-sol, reasoning effort high, session 01a0dcbd-22a4-7b63-b2a6-f2802f6e8c27. Result: FINDINGS

Scope: c01c56f73c7dbf973fe73b52b6093f0a32952321..2de91e768196dc061251f7609ae0d9b9be2905cd. The reviewer ran in a --sandbox read-only detached snapshot of the head with no .env. The prompt contained only the diff range, a file list, and the review questions; it had no PR number, description, commit message, or prior findings. Static review: the reviewer ran no tests. The snapshot and the PR worktree were unchanged afterwards.

Covered (reviewer's words): "synchronous and native asyncio execute() and executemany() paths through the default, dict, pandas, Arrow, Polars, and S3FS cursors; cursor state and exception chaining; the thread-backed async and Spark overrides; the new tests and both documentation examples."

[P2] Pre-existing, exposed by the new contract: "A second KeyboardInterrupt or task cancellation during the stop request or follow-up poll escapes the handler because both are BaseException subclasses, outside except Exception. With the query still running, execute() exits before observing a terminal state. ... The behavior predates the diff, while the new documentation states the wait without this qualification." (also pyathena/aio/common.py:155)

Disposition: documented; behavior kept (pre-existing). Verified: a second KeyboardInterrupt or task cancellation raised during _cancel() or the follow-up poll is a BaseException, so except Exception does not catch it and it propagates with the first interrupt as its context. This predates the diff, matches the Spark contract from #833, and gives users a way to stop waiting on a query that does not stop. da16975 documents it: docs/usage.md says "A second KeyboardInterrupt during that wait propagates without waiting for the terminal state.", and docs/aio.md has the equivalent sentence for a repeated task cancellation.

Independent follow-up review (relayed): Codex CLI, model gpt-6-sol, reasoning effort high, session 01a0dcc3-f3ef-7833-a415-86e914ba8fa2. Scope: 2de91e768196dc061251f7609ae0d9b9be2905cd..da169757662dda486e41349789ce7a77e67ddc91 (same base c01c56f73c7dbf973fe73b52b6093f0a32952321), read-only snapshot, static review.

Covered surfaces: the four files changed in 2de91e7..da16975, traced through pyathena/aio/cursor.py, pyathena/aio/common.py, pyathena/aio/util.py, pyathena/cursor.py, and pyathena/common.py. This was a static, read-only review; I did not run tests.

FINDINGS

P2 — docs/aio.md:142, docs/usage.md:511: The new qualification covers a second interruption while polling for the terminal state, but says nothing about one during the stop request. Both handlers call _cancel before that follow-up poll, and their except Exception blocks do not catch a second CancelledError or KeyboardInterrupt. For example, a second task cancellation while the stop request is awaiting its worker thread exits the handler without confirming that the stop request ran or that the query became terminal. The documentation still implies the first interruption completes those steps unless interrupted “during that wait.”

Prior items

Resolved. docs/aio.md:145 limits the timeout claim to polling and explains that cancellation during startup can leave query_id unset while the worker-thread request continues.

Not an actual defect. In this test, _execute is an immediately completing AsyncMock; the first status request then enters an Event.wait() and suspends. On Python 3.10, wait_for() schedules the child task before its timeout can cancel it; on 3.11–3.14, it awaits the coroutine inside the timeout context. The event loop therefore reaches that blocked status request before delivering the 10 ms timeout. A separate barrier is unnecessary for this fixture.

Resolved. Both tests now return RUNNING before a terminal state and assert that the polling callback saw both states before the original interruption propagated.

Unresolved for the stop-request window; resolved for follow-up polling. The source behavior described in the finding remains possible, and the added wording qualifies only “that wait.”

Fixed in f528d50. The repeated-interrupt note now covers the cancellation request as well as the follow-up wait, in both docs/usage.md and docs/aio.md, and says the query can keep running.

Independent follow-up review (relayed): Codex CLI, model gpt-6-sol, reasoning effort high, session 01a0dcc8-cb13-7712-8373-34cdcfc3f7e1. Scope: da169757662dda486e41349789ce7a77e67ddc91..f528d50bc242ccacb06631081a7f21bda31d70a2, read-only snapshot, static review.

Covered surfaces: docs/aio.md “Task cancellation” and docs/usage.md “Query cancellation on interrupt,” checked against the five named source files. The new wording covers a second interruption during both the cancellation request and the follow-up poll.

CLEAN. No actionable inaccuracies found in either section. This was a read-only source review; no tests were run.

laughingman7743 · 2026-09-26T08:09:30Z

+If the cancellation request fails, `asyncio.CancelledError` is raised with the error as its cause.
+With `kill_on_interrupt=False`, `asyncio.CancelledError` is raised immediately and the query keeps running.
+
+A timeout from `asyncio.wait_for()` therefore cancels the query and raises `asyncio.TimeoutError`:


Independent review (relayed, Codex gpt-6-sol): [P2] introduced

"The wait_for() example says a timeout cancels the query. If the timeout occurs while start_query_execution is running in a worker thread, execute() has not assigned query_id; cancellation cannot stop the request, and Athena may start a query after the task exits. The example can print None while that query keeps running."

Disposition: fixed in da16975. Verified: aio StartQueryExecution runs through asyncio.to_thread() (pyathena/aio/util.py:42), and query_id is assigned only when _execute() returns. The docs now limit the cancellation to a timeout that expires while execute() waits for the query, and state that a timeout during the start leaves query_id as None while the pending start request can still start the query. execute() has no await between the query_id assignment and _poll(), so there is no third window. The example prints query_id as a value that may be None.

laughingman7743 · 2026-09-26T08:09:30Z

+            kill_on_interrupt=True, final_state=AthenaQueryExecution.STATE_CANCELLED
+        )
+        with pytest.raises(asyncio.TimeoutError):
+            await asyncio.wait_for(cursor.execute("SELECT 1"), timeout=0.01)


Independent review (relayed, Codex gpt-6-sol): [P2] introduced

"The 10 ms timeout test has no barrier confirming that the task reached polling. Under a scheduling delay, it can time out before a query ID is assigned; cancel is then never awaited and the assertion fails intermittently."

Disposition: rejected with evidence. No await in execute() yields to the event loop before the first status request blocks. _execute is an AsyncMock, which completes without suspending, and _call_on_start_query_execution is synchronous, so query_id is always assigned before the loop can run the timeout callback. On 3.12+, wait_for() awaits the coroutine inside the calling task under timeouts.timeout(), so the timer can fire only at the first suspension, the blocked poll. On 3.10/3.11, wait_for() wraps the coroutine in a task whose first step call_soon places in _ready before _run_once moves any expired timer into _ready, so that step runs first even after a scheduling delay. The test passed on 3.13.1 and 3.10.16 locally.

laughingman7743 · 2026-09-26T08:09:30Z

+    def test_execute_kill_on_interrupt(self, final_state):
+        """An interrupt cancels the query, waits for it, and is re-raised (no AWS)."""
+        cursor, cancel = _offline_cursor(kill_on_interrupt=True)
+        cursor._get_query_execution = MagicMock(


Independent review (relayed, Codex gpt-6-sol): [P3] introduced

"The follow-up status is immediately terminal in the new tests. Given a query that stays RUNNING after the stop request, an implementation that polls once and raises before terminal state could still satisfy these tests. They exercise the prior implementation's failing path, but do not assert the promised wait through a nonterminal state." (also tests/pyathena/aio/test_cursor.py:197)

Disposition: fixed in da16975. Both test_execute_kill_on_interrupt tests now report RUNNING after the stop request, then the terminal state. They assert through the on_poll hook that the polled states were [RUNNING, final_state] before the interrupt or cancellation was re-raised. Mutation check: replacing the follow-up __poll() with a single _get_query_execution() call in both _poll() implementations makes all 4 of these tests fail. With the real implementation, 13 targeted tests pass on 3.13.1, and 11 interrupt tests pass on 3.10.16.

A second interrupt or task cancellation also escapes while the cancellation request is in progress, not only during the follow-up wait. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

laughingman7743 commented Sep 26, 2026

View reviewed changes

Cover the cancellation request in the repeated-interrupt note

f528d50

A second interrupt or task cancellation also escapes while the cancellation request is in progress, not only during the follow-up wait. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

laughingman7743 marked this pull request as ready for review September 26, 2026 08:34

Navigation Menu

Re-raise the interrupt after kill_on_interrupt cancellation - #853

Re-raise the interrupt after kill_on_interrupt cancellation#853
laughingman7743 wants to merge 3 commits into
masterfrom
fix/840-sql-cursor-interrupt

laughingman7743 commented Sep 26, 2026 •

edited

Loading

Uh oh!

laughingman7743 Sep 26, 2026

Uh oh!

laughingman7743 Sep 26, 2026

Uh oh!

laughingman7743 Sep 26, 2026

Uh oh!

laughingman7743 Sep 26, 2026

Uh oh!

laughingman7743 Sep 26, 2026

Uh oh!

laughingman7743 Sep 26, 2026

Uh oh!

laughingman7743 Sep 26, 2026

Uh oh!

laughingman7743 Sep 26, 2026

Uh oh!

laughingman7743 Sep 26, 2026

Uh oh!

laughingman7743 Sep 26, 2026

Uh oh!

laughingman7743 Sep 26, 2026

Uh oh!

laughingman7743 Sep 26, 2026

Uh oh!

laughingman7743 Sep 26, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

laughingman7743 commented Sep 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

WHAT

WHY

TEST

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

laughingman7743 commented Sep 26, 2026 •

edited

Loading