Terminate only the Spark sessions a cursor started on close

[laughingman7743]

WHAT

• Add terminate_session_on_close: bool | None = None to SparkCursor, AsyncSparkCursor, and AioSparkCursor (forwarded to SparkBaseCursor).
  • None (default): close() terminates only a session the cursor started. A session supplied with session_id is left running.
  • True: close() also terminates a supplied session (the previous behavior).
  • False: close() never terminates the session, so a started session can be reused later with session_id.
• After a successful termination, close() does not send TerminateSession again. After a failed termination, a later close() retries it. AsyncSparkCursor.close() still shuts down its executor in every case.
• The cleanup of a newly started session when cursor initialization fails (Spark: clean up newly created sessions when cursor initialization fails #794) is unchanged and does not depend on the new argument; a supplied session is never terminated by it.
• docs/spark.md: rewrite "Session lifecycle" with the ownership table and a reuse example, state that closing a Connection or AioConnection (including their context managers) does not terminate Spark sessions, and close the cursors in the Spark examples.
• Connection, AioConnection, and the shared cursor code are unchanged.

Behavior change for release notes: closing a Spark cursor created with session_id no longer terminates that session. Pass terminate_session_on_close=True to keep the previous behavior.

WHY

Closes #793
Closes #796

Part of #791. Passing session_id attaches a cursor to a session that is managed elsewhere, so closing that cursor must not end the session for its owner. For #796, session cleanup stays the responsibility of the cursor: the connection does not track cursors, and the documentation now says so.

TEST

Tested commit: 60669e4 (rebased onto master 0ba0875; the live AWS run below was on 22bc5a8, before a behavior-preserving refactor of how close() stores the resolved setting).

• just lint: passed.
• uv run --env-file .env pytest -n 1 tests/pyathena/spark/test_common.py plus the existing offline AsyncSparkCursor close tests: 64 passed.
  • New tests cover the three cursor variants × {started, supplied} × {None, True, False}, a repeated close after success, a retry after a failed termination, and executor shutdown of AsyncSparkCursor when the session is not terminated.
  • Against the source of origin/master (c01c56f), 16 of the 25 new tests fail; the other 9 are the started-session and True cases that already terminated before this change (the base cursor accepts unknown keyword arguments, so the new argument is ignored there).
• Live AWS, run once locally with -n 1: test_session_ownership in TestSparkCursor, TestAsyncSparkCursor, and TestAioSparkCursor passed (3 passed in 147 s). Each test attaches a second cursor to the fixture cursor's session, runs a calculation, closes it, runs another calculation with the fixture cursor, then closes the fixture cursor and waits for the session to become TERMINATED. ListSessions afterwards showed the three sessions TERMINATED and no active sessions in the Spark work group.
• The documentation statement that Athena terminates a session left running after its idle timeout matches an observed session from earlier Spark runs on this account: it ended TERMINATED with the reason "Session has timed out after idle" about 21 minutes after it started with a 20-minute idle timeout.
• just docs lint: passed. A single-version sphinx-build of the working tree produced no warnings for the Spark page.

🤖 Generated with Claude Code

[laughingman7743]

Self-review round one (behavior and implementation) — CLEAN

Base c01c56f73c7dbf973fe73b52b6093f0a32952321 (merge-base with master), head 22bc5a8c27a31aa1f0abba6fc57e3c340465263e.

Covered: SparkBaseCursor.__init__ ownership assignment and close()/_should_terminate_session(), SparkCursor/AsyncSparkCursor/AioSparkCursor constructor forwarding and positional order (new parameter follows session_idle_timeout_minutes, and follows max_workers in AsyncSparkCursor), AioSparkCursor.close(), AsyncSparkCursor.close() executor shutdown, the #794 init cleanup path, fixtures that close the cursor after a test already closed it, the offline and live tests, and the docs examples.

Checked scenarios:

• Started session with None/True terminates once; supplied session with None/False never calls TerminateSession; True restores the previous behavior.
• Repeated close after success is a no-op (_session_terminated), so the fixture teardown after the live tests' explicit close sends nothing; a failed termination leaves the flag unset and a later close retries.
• AsyncSparkCursor.close() keeps the try/finally from Shut down the AsyncSparkCursor executor when session termination fails #830, so the executor is shut down even when no termination is attempted.
• _start_session() cleanup (Spark: clean up newly created sessions when cursor initialization fails #794) runs before ownership is assigned and ignores terminate_session_on_close, so False cannot leak a session from a failed init; a supplied session is never terminated there.
• execute(session_id=...) overrides are not tracked and are never terminated, as before.

Test quality: the new offline tests fail on the master source for the 16 combinations whose behavior changed (the other 9 already terminated before); the live test proves a borrower close leaves the session usable and the owner close reaches TERMINATED.

Findings: none.

[laughingman7743]

Repair — self-review round one (behavior), affected scope — CLEAN

Old: base c01c56f73c7dbf973fe73b52b6093f0a32952321, head 22bc5a8c27a31aa1f0abba6fc57e3c340465263e. New: base 0ba0875 (rebased onto current master), head 60669e494b9a24ccd8d52e8444757591d45196df. Range-diff: 63e6770 differs from 22bc5a8 only in the tests/pyathena/util.py context (the add/add merge with decorated() from master; both helpers kept); 60669e4 is the repair.

Repair: terminate_session_on_close=None is resolved once in SparkBaseCursor.__init__ to not session_id, and close() (sync and AioSparkCursor) clears the flag after a successful termination. _owns_session, _session_terminated, and _should_terminate_session() are removed.

Equivalence checked:

• None resolves with the same truthiness test as the start/attach branch (if session_id:), so a cursor terminates exactly the sessions it starts, including the empty-string case.
• Repeated close after success is a no-op; a failed termination raises before the flag is cleared, so a later close retries; explicit True/False are unchanged.
• The Spark: clean up newly created sessions when cursor initialization fails #794 cleanup in _start_session() does not read the flag; an init failure returns no cursor, so where the flag is resolved does not affect it.
• The flag has no other readers (grep).

Tests: the offline ownership tests assert TerminateSession calls rather than internals and pass unchanged (64 passed incl. the #830 close tests, whose __new__ helper now sets the single flag). just lint passed. Upstream changes brought in by the rebase (#836, #843, #858) touch only the SQLAlchemy dialect and do not affect these contracts. Live tests were not rerun: the observable behavior they cover is unchanged and the AWS jobs run them after Ready.

[laughingman7743]

Self-review round two (claims, compatibility, operations) — CLEAN, PR body corrected

Base c01c56f73c7dbf973fe73b52b6093f0a32952321, head 22bc5a8c27a31aa1f0abba6fc57e3c340465263e.

Claims checked:

• "Closing the connection does not close its cursors or terminate their sessions ... Connection and AioConnection, including their context managers": Connection.close() is a no-op, Connection.cursor() keeps no reference, and AioConnection.__aexit__ only calls close() (pyathena/connection.py, pyathena/aio/connection.py). Holds.
• "A session that is left running is terminated by Athena after its idle timeout": an unclosed session from earlier Spark runs on this account (started 12:25:14 JST, 20-minute idle timeout) ended TERMINATED at 12:46:47 with reason "Session has timed out after idle". Holds.
• "These arguments apply only to a session started by the cursor": engine_configuration, notebook_version, description, and session_idle_timeout_minutes are read only by _start_session(). Holds.
• The reuse example (False then session_id=..., True) matches _should_terminate_session().
• PR body: "16 of the new tests fail on master" was incomplete; corrected to 16 of 25, explaining that the 9 passing cases already terminated and that BaseCursor ignores the unknown keyword on master.

Existing callers: the new parameter is appended after the existing positional parameters (after max_workers for AsyncSparkCursor), so positional calls are unchanged. The only behavior change is for callers passing session_id who relied on close() terminating it; the PR states it for release notes with terminate_session_on_close=True as the migration. A public code search found no callers passing session_id beyond vendored copies of PyAthena.

AWS operator: the change only removes TerminateSession calls (borrowed sessions, repeated close); a borrowed session stays billed until its owner or the idle timeout ends it, which the docs state.

Adversarial: two cursors on one session with True both try to terminate it; the second close may fail against an already terminated session and can be retried or ignored by the caller; not claimed otherwise in the docs.

Findings: none in code; PR body corrected.

[laughingman7743]

Repair — self-review round two (claims), affected scope — CLEAN, PR body updated

New head 60669e494b9a24ccd8d52e8444757591d45196df (base 0ba0875).

Claims rechecked against the repair:

• SparkBaseCursor.close() / AioSparkCursor.close() docstrings ("After a successful termination, further calls do not terminate the session again; after a failed one, calling this method again retries it") hold for the single flag; they name no removed helper.
• Constructor docstrings ("If None, only a session started by this cursor is terminated; a session supplied with session_id is left running") hold for not session_id.
• Commit message "the behavior is unchanged": supported by the unchanged offline tests covering 3 variants × {started, supplied} × {None, True, False}, repeated close, and retry.
• docs/spark.md is unchanged by the repair; its table and reuse example still match.
• PR body TEST section now names the tested commit and that the local live run was on 22bc5a8, before this behavior-preserving refactor.

Existing callers: public signatures and defaults are unchanged.

[laughingman7743]

Independent review (relayed) — FINDINGS: 1, pre-existing, deferred

Reviewer: Codex CLI 0.157.0, model gpt-6-sol, reasoning effort high, sandbox read-only, session 01a0dcb5-9730-7551-92d0-8bf74fd1522d. Static review of a detached snapshot at head 22bc5a8c27a31aa1f0abba6fc57e3c340465263e against merge-base c01c56f73c7dbf973fe73b52b6093f0a32952321, given the literal diff and the intended contract, without the PR number, description, or prior findings. The snapshot was unchanged afterwards.

Covered (as reported): the diff; all three Spark cursor variants; constructor cleanup, close retries, and executor shutdown; connection and context-manager behavior; the tests, fixtures, helper, and Spark documentation.

Finding (P2): pyathena/connection.py:555 — Connection.cursor() applies kwargs.update(self.cursor_kwargs), so connection-level cursor_kwargs override arguments passed to cursor(). With cursor_kwargs={"terminate_session_on_close": False}, conn.cursor(session_id=..., terminate_session_on_close=True) receives False.

Disposition: verified as described, and pre-existing — the same precedence already applies to every cursor argument (session_id, engine_configuration, max_workers, and all non-Spark cursors). Changing it means changing shared Connection.cursor() behavior for all cursors, which is outside this Spark-only change; no code change here. Recorded for a separate issue.

[laughingman7743]

Repair — independent follow-up (relayed) — CLEAN

Reviewer: Codex CLI 0.157.0, model gpt-6-sol, reasoning effort high, sandbox read-only, session 01a0e0b0-6ef0-7510-b6b0-ae94b06e6002. Static review of a detached snapshot at head 60669e494b9a24ccd8d52e8444757591d45196df, given the range-diff c01c56f..22bc5a8 → 0ba0875..60669e4 and the repair diff 63e6770..60669e4 with the intended contract; no PR number, description, or prior findings. The snapshot was unchanged afterwards.

Covered (as reported): the range diff and repair diff; constructor and close behavior in all three Spark cursor variants; falsy session_id, termination failure and repeated calls, construction cleanup, and executor shutdown; the affected callers, tests, test utility, and documentation.

Result: CLEAN. The previous P2 (Connection.cursor() cursor_kwargs precedence, pre-existing) remains deferred as recorded above.