Cancel a Spark calculation interrupted while it is being started

[laughingman7743]

WHAT

• An interrupt while a Spark cursor is still starting a calculation now cancels that calculation, following the Define best-effort Spark calculation cancellation #833 contract.
  • Sync (SparkCursor, and AsyncSparkCursor whose execute() starts the calculation on the caller's thread): with kill_on_interrupt, SparkBaseCursor._calculate() sends StartCalculationExecution from a short-lived helper thread and waits for it on the calling thread. On KeyboardInterrupt, it waits for the request to finish, requests StopCalculationExecution for the calculation it started, waits for a terminal state, stores calculation_id and the final execution, and re-raises the interrupt.
  • asyncio (AioSparkCursor): with kill_on_interrupt, the start request is wrapped in a task and awaited through asyncio.shield(). On asyncio.CancelledError, the same recovery runs and the cancellation is re-raised, so task.cancelled() is true and asyncio.wait_for() raises TimeoutError.
  • A failure to start, cancel, or wait becomes the interrupt's __cause__. A second interrupt during the recovery propagates at once without cancelling.
  • kill_on_interrupt=False is unchanged: the request runs on the calling thread / directly in the task, and the interrupt propagates immediately.
• The Spark cursors send a generated ClientRequestToken (uuid4, one per execute()) when the caller passes no client_request_token; a caller-supplied token is sent as is. botocore does not auto-generate this token for StartCalculationExecution (it does for StartQueryExecution), so a retried start after a lost response is a new request; on the session made busy by the first attempt it fails with InvalidRequestException (measured for a different token on a busy session), leaving the first calculation running with no ID on the cursor.
• The polling-phase recovery from Define best-effort Spark calculation cancellation #833 is unchanged; it now shares a small __cancel_and_wait() helper with the start phase.
• Shared code is not changed: the new logic lives in SparkBaseCursor (overriding BaseCursor._calculate) and AioSparkCursor; _build_start_calculation_execution_request is reused as is.
• Docs: docs/spark.md cancellation sections describe the start-phase behavior, the generated token, that a caller-supplied token must be unique per calculation, and that a stop request sent right after a calculation starts can occasionally have no effect.

Behavior change (release note):

• Spark cursors send a generated ClientRequestToken unless client_request_token is given.
• With kill_on_interrupt (default), an interrupt while the calculation is being started waits for the start request, then stops that calculation before the interrupt propagates. If the helper thread has not begun the request when the interrupt is handled, the request is abandoned and never sent.

WHY

Closes #841. Parent: #791.

Measured with the CI account on pyathena-spark (DPU 1/2/1, one session at a time, all sessions terminated afterwards):

• Replaying StartCalculationExecution with the same token, while the calculation runs or after it finished, returned the same calculation ID without a BUSY error and without a second calculation. A start with a different token on the busy session failed with InvalidRequestException.
• Replaying the token with a different CodeBlock also returned the original calculation without an error, so recovery by replaying a token could start code the user interrupted if the original request never arrived; recovery by listing calculations could stop another borrower's calculation on a shared session. This PR therefore waits for the in-flight request instead.
• A stop request right after the start (CREATED) cancelled the calculation within 0.4 s in 5 of 6 attempts; in 1 of 6 it had no effect and the calculation ended COMPLETED after about 95 s. Cancellation stays best-effort.

TEST

Tested commits: 73a738e for the local live run and 6d0e5e9 for the offline tests (before rebasing). The later commits change only docstrings, docs/spark.md, and the name and comment of the wait interval constant. The current head is 1466e7b (rebased onto 9a1a72d): just lint and just docs lint passed, and the offline Spark tests passed (113 passed, 2 skipped). The AWS suites run in CI.

• just format, just lint: passed. just docs lint: 0 errors. uv run sphinx-build -q -E docs <tmp>: no warnings from docs/spark.md.
• Offline (no AWS), uv run --env-file .env pytest --noconftest -p no:xdist tests/pyathena/spark tests/pyathena/aio/spark -k "not (spark_dataframe or spark_sql or failed or test_cancel or executemany or context_manager)": 83 passed. The existing live aio test test_context_manager also ran against Athena with the generated token on 73a738e and passed.
  • Token generated once and reused across a throttled retry, a new token per call, and a caller token preserved (sync × both cursors × kill_on_interrupt; asyncio).
  • Sync: an interrupt raised while the helper thread is being started (helper not yet running / running later) sends no request and does not stop anything (fails 2 of 4 cases on c0188ce, the head before that repair).
  • Sync: KeyboardInterrupt injected into the wait while start_calculation_execution is blocked on an event; the request is sent once, the stop request uses the returned ID, calculation_id and the final execution are stored, and the same interrupt is re-raised (final state CANCELED and COMPLETED). Start / cancel / wait failures become __cause__; a second interrupt propagates with the first as __context__ and no stop; without kill_on_interrupt no helper thread is created.
  • asyncio: task.cancel() while the start is blocked in the worker thread; the task stays pending until the request is released, then the stop is sent and task.cancelled() is true. asyncio.wait_for() raises TimeoutError after the stop. Failures become __cause__; a second cancellation propagates without a stop; without kill_on_interrupt the cancellation propagates at once.
• Regression check: with the new tests on origin/master (0ba0875) code, 28 of the 33 new tests fail; the 5 that pass cover preserved behavior (caller token, second cancellation, no kill_on_interrupt).
• Live: uv run --env-file .env pytest -n 1 tests/pyathena/spark tests/pyathena/aio/spark → 95 passed, 2 skipped, 286 s; afterwards ListSessions showed no non-terminated sessions. One unrelated Test run was in progress at the time.
• Real SIGINT (not a patched wait) delivered while the start request was blocked, on Python 3.10.16 and 3.13.1 (probe outside the repository, mocked client): the wait was interrupted, the stop request used the returned ID, and KeyboardInterrupt was re-raised.
• Not tested live: an interrupt landing inside the ~0.3 s start window against Athena; it is covered by the deterministic offline tests and the probe above.

🤖 Generated with Claude Code

[laughingman7743]

Self-review round one — implementation behavior (base 0ba0875c72c702e0b893843d3c50aa98b4a61b8f, head 73a738e4fd01b00e3a1f947959c730867dcf0767; the later head c0188ced93137cd106a416bd092affaf13c4a552 only adds a docs sentence)

Covered: SparkBaseCursor._calculate / __start_calculation / __wait_for_start / __cancel_and_wait, AioSparkCursor._calculate / __start_calculation / __cancel_and_wait, the unchanged _poll recovery, callers SparkCursor.execute, AsyncSparkCursor.execute (starts on the caller's thread), AioSparkCursor.execute, and the new offline tests.

Result: CLEAN (no actionable findings).

Checks:

• Interruptibility: the helper-thread wait loops on concurrent.futures.wait(timeout=0.1), so a KeyboardInterrupt is raised promptly even where untimed lock waits are not interruptible. Verified with a real SIGINT (not the patched wait used by the tests) on Python 3.10.16 and 3.13.1: the stop request used the returned ID and the interrupt was re-raised.
• No thread leak on the normal path: the daemon helper ends after setting the future (asserted by joining pyathena-spark-start threads in the token test). Without kill_on_interrupt, no thread is created (asserted).
• A start failure without an interrupt still propagates as DatabaseError through future.result().
• asyncio: without kill_on_interrupt, the start is awaited directly (no shield), so a cancelled task cannot leave an unretrieved task exception. With it, a second cancellation cancels the awaited start task and propagates.
• Tests exercise the failing path: 28 of 33 new tests fail on the base code.

Recorded limitations (not introduced here): a KeyboardInterrupt raised inside Thread.start() itself, before the try, propagates without recovery as all interrupts during the start did before. The recovery wait is bounded by the connection's RetryConfig for the start request; a second interrupt ends it at once.

[laughingman7743]

Self-review round two — claims, callers, operations (head c0188ced93137cd106a416bd092affaf13c4a552)

Result: FINDINGS, repaired in c0188ce and the PR body.

Claims checked:

• botocore 1.43.102 marks ClientRequestToken as idempotencyToken for StartQueryExecution but not for StartCalculationExecution / StartSession — confirmed from the service model.
• Token replay returns the same calculation, including with a different CodeBlock; a different token on a busy session fails with InvalidRequestException; a stop right after the start was ignored in 1 of 6 attempts — from the Spark: an interrupt during StartCalculationExecution leaves the calculation running #841 measurement.
• 28/33 failing on base, 80 offline passed, 95 passed / 2 skipped live with no sessions left — rerun results for this head.

Findings and repairs:

1. The PR body said a retried start after a lost response "previously failed"; that case was inferred from the busy-session measurement, not observed. Reworded to state the measured part.
2. Docs did not say that a caller-supplied client_request_token reused with different code silently returns the earlier calculation (measured). Added to this paragraph.

Caller compatibility: _calculate's signature and return are unchanged; execute(client_request_token=...) still sends the caller's token; kill_on_interrupt=False keeps the old control flow. AWS operator: one extra thread per execute() with kill_on_interrupt; retries of the start request now reuse one token, so they cannot create a second calculation.

[laughingman7743]

Independent review (relayed) — Codex CLI 0.157.0, codex exec --sandbox read-only, model_reasoning_effort=high, run on a detached snapshot of c0188ced93137cd106a416bd092affaf13c4a552 against merge-base 0ba0875c72c702e0b893843d3c50aa98b4a61b8f, without the PR description or prior findings. Static review; the reviewer ran no tests.

Coverage reported: the full diff (Spark cursor changes, offline tests, docs/spark.md), the execute() callers, the request builder, and the retry helpers; confirmed pyathena/common.py and pyathena/aio/common.py are unchanged.

Verdict: FINDINGS (1)

• P2 — pyathena/spark/common.py:407 (at c0188ce): the helper thread was started before the try that handles KeyboardInterrupt. A Ctrl-C while Thread.start() waits for the thread to initialize escapes without recovery, while the helper can still send StartCalculationExecution, leaving the calculation running with no ID on the cursor.

Verified and repaired in 6d0e5e9: Thread.start() moved inside the try; the helper claims the future with set_running_or_notify_cancel() before sending, and the recovery first calls future.cancel() — if that succeeds the request was never sent and never will be, so the interrupt propagates at once; otherwise the recovery waits for the request as before. New test test_calculate_interrupted_before_request_is_sent (helper not started / started after the interrupt) fails 2 of 4 cases on c0188ce and passes now.
The asyncio path has no equivalent window: _calculate has no await before asyncio.shield(start), and the start task's first step, which dispatches the request to the worker thread, is scheduled before the outer task can resume with CancelledError.

[laughingman7743]

Repair self-review of 6d0e5e9 (range-diff 0ba0875..c0188ce → 0ba0875..6d0e5e9: the first two commits are unchanged; one new commit touching pyathena/spark/common.py and tests/pyathena/spark/test_common.py).

Round one (behavior): CLEAN. Future.set_running_or_notify_cancel() and Future.cancel() serialize on the future's condition lock, so exactly one outcome holds. Either the helper claims the future first, in which case cancel() returns False and the recovery waits for the ID as before, or cancel() wins and the helper returns without sending. An interrupt before _start_new_thread leaves a pending future that cancel() settles, so the recovery never waits on a thread that does not exist. The normal path and kill_on_interrupt=False are unchanged; __wait_for_start is never called on a cancelled future. Offline Spark tests: 83 passed. A real-SIGINT probe on 3.13.1 still stops the returned ID.

Round two (claims): CLEAN. The commit message matches the control flow. The _calculate docstring now states that an interrupt before the request is sent propagates without sending it. The docs/spark.md sentence ("first waits for the StartCalculationExecution request to finish") remains accurate because nothing is sent in the new case.

[laughingman7743]

Independent follow-up (relayed): Codex CLI 0.157.0, read-only, model_reasoning_effort=high, run on a detached snapshot of 6d0e5e9. It reviewed the repair diff c0188ce..6d0e5e9 and the range-diff, and did not see this thread. It was a static review; the reviewer ran no tests.

Verdict: FINDINGS (2 × P2). Both findings were checked against the code, and neither leaves a calculation running without a stop request:

1. pyathena/spark/common.py:405 (at 6d0e5e9). An interrupt between the helper's set_running_or_notify_cancel() and the HTTP call makes cancel() return False, so the request is still sent and then stopped.
   Disposition: behavior kept, docstring corrected in a2008d2. The claim is the commit point: once the helper has begun, the cursor waits for the calculation ID and stops that calculation. That is the agreed design, "finish the in-flight start, then stop it." The docstring had said "before the helper sends the request", which over-promised. It now says "before the helper begins the request", and the two code comments were reworded to match.
2. pyathena/aio/spark/cursor.py:154. A cancellation requested after ensure_future() but before the start task's first step still dispatches the request, which is then stopped.
   Disposition: rejected, no change. The outcome is the documented one: the cursor waits for the start, stops the calculation, and re-raises. The calculation is never left running unseen, which is what Spark: an interrupt during StartCalculationExecution leaves the calculation running #841 is about. Detecting "cancel requested but not yet delivered" before the first step would need Task.cancelling(), which is 3.11+ (the floor is 3.10), or private task state. AioSparkCursor._calculate documents that a cancellation waits for the request to finish, and that claim holds.

Repair self-review of a2008d2 (docstring and comments only):

• Round one: CLEAN. No code change; lint passed.
• Round two: CLEAN. The new wording matches the control flow in both outcomes of the set_running_or_notify_cancel() / cancel() race.

[laughingman7743]

Independent follow-ups on the docstring repair (relayed). Codex CLI 0.157.0, read-only, model_reasoning_effort=high, detached snapshots, run without this thread. Static reviews only.

• a2008d2, reviewing 6d0e5e9..a2008d2: FINDINGS, P2 at pyathena/spark/common.py:370. "An interrupt before the helper begins the request … does not send it" is false in one race: the interrupt can be raised while the future is pending, and the helper can then claim it before the handler calls cancel(). The outcome is decided by which call wins, not by when the signal arrives. Verified; reworded in 9a9965a.
• 9a9965a, reviewing a2008d2..9a9965a: FINDINGS, P2 at the same line. "When handled, a request the helper has not yet begun is abandoned" still did not make the handler's cancel() attempt the condition. Verified; reworded in eeafbc1: "the cursor first tries to abandon the request. This succeeds only if the helper has not begun the request by then …"
• eeafbc1, reviewing 9a9965a..eeafbc1: CLEAN.

Repair self-review of 9a9965a and eeafbc1: these are docstring-only changes.

• Round one is CLEAN: there is no code change, and just lint passed.
• Round two is CLEAN: the text now matches both outcomes of the race, and the code comments ("Begin the request only if no interrupt has given up on it yet", "The helper has not begun the request and never will") describe the successful-cancel() branch they sit in.

[laughingman7743]

Rebase onto ceb9ad7 (#844 merged): repair review. The series moved from 0ba0875..eeafbc1 to ceb9ad7..bdfa630, and the PR went back to Draft for the review.

git range-diff:

• Commits 2–6 are identical.
• Commit 1 differs only in tests/pyathena/aio/spark/test_cursor.py. That was an add/add conflict next to Terminate only the Spark sessions a cursor started on close #844's test_session_ownership, resolved by keeping both sides: the imports are merged (AthenaSessionStatus from Terminate only the Spark sessions a cursor started on close #844, plus DatabaseError and RetryConfig), and the Terminate only the Spark sessions a cursor started on close #844 test is kept ahead of the new start-cancellation tests.
• pyathena/spark/common.py, pyathena/aio/spark/cursor.py and docs/spark.md applied cleanly.

Self-review round one (behavior): CLEAN.

• Terminate only the Spark sessions a cursor started on close #844 changes __init__ and close() only. _calculate and __cancel_and_wait stop by calculation ID and never terminate a session, so a borrowed session or a per-call execute(session_id=...) override is unaffected.
• just lint passed. pytest --noconftest over tests/pyathena/spark and tests/pyathena/aio/spark: 113 passed and 2 skipped; the 14 errors are live-fixture setups that --noconftest skips by design.

Self-review round two (claims): CLEAN.

• The docs paragraphs sit after Terminate only the Spark sessions a cursor started on close #844's rewritten lifecycle section and do not contradict it. just docs lint reports 0 errors.

Independent follow-up (relayed): CLEAN.

• Reviewer: Codex CLI 0.157.0, read-only, model_reasoning_effort=high, run on a detached snapshot of bdfa630 with the range-diff and the full ceb9ad7..bdfa630 delta, without this thread. Static review; it ran no tests.
• Codex reported that all six commits carried across, and that the conflict resolution kept both the ownership test and the start-cancellation tests.
• It found that recovery cancels by calculation ID even under execute(session_id=...), that closing a borrowed cursor leaves its session running, and no actionable interaction.

[laughingman7743]

Repair: rename _START_WAIT_INTERVAL to _INTERRUPT_CHECK_INTERVAL and rebase onto 9a1a72d (maintainer feedback that the old name was confusing). The PR was moved back to Draft for this repair.

The series moved from ceb9ad7..bdfa630 to 9a1a72d..1466e7b. git range-diff shows commits 1–6 identical and one new commit, 1466e7b. Upstream changes between the two bases touch only .github/workflows/* and docs/testing.md; per #863, PRs now run the AWS tests on the newest Python only. No Spark source or test changed upstream, and the rebase was clean.

What changed:

• Constant renamed; the value stays 0.1. Its single use is __wait_for_start. No other reference exists in the code, tests, docs, or PR text.
• The comment now states its purpose first: "How often a wait for the start request wakes up to check for Ctrl-C, so that a KeyboardInterrupt is raised promptly where an untimed lock wait cannot be interrupted by signals (Windows before Python 3.14)."
• The platform claim was checked against the Python docs for threading.Lock.acquire: "Changed in version 3.2: … can now be interrupted by signals on POSIX" and "Changed in version 3.14: … on Windows". The previous comment's "on every platform" was therefore narrowed to the case the timeout actually covers within the supported 3.10+ range.

Self-review round one (behavior): CLEAN. The rename is behavior-neutral. After the rebase, just lint passes; offline Spark tests give 113 passed, 2 skipped (plus the 14 live-fixture setup errors expected with --noconftest). A real-SIGINT probe on 3.13.1 still cancels the returned ID.

Self-review round two (claims): CLEAN. The comment's version and platform claim matches the Python docs, and the commit message says the same. just docs lint reports 0 errors.

Independent follow-up (relayed): CLEAN.

• Reviewer: Codex CLI 0.157.0, read-only, model_reasoning_effort=high, run on a detached snapshot of 1466e7b with the range-diff, the upstream delta, and the new commit, without the PR text or prior findings. Static review; it ran no tests.
• It confirmed the six commits are unchanged, found no interaction with the upstream CI changes, and found no remaining references to the old name.
• It traced wait() → Event.wait() → Condition.wait() and judged the Windows before 3.14 qualification correct.