Re-raise and stop interrupted queries in the SQL cursors, sharing the Spark handling

docs/aio.md

@@ -130,6 +130,38 @@ async with await aio_connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
         await cursor.cancel()
 ```
 
+(aio-task-cancellation)=
+
+### Task cancellation
+
+With `kill_on_interrupt` enabled, which is the default, cancelling the task while `execute()` waits for the query
+requests cancellation of the query, waits until it reaches a terminal state, and then raises `asyncio.CancelledError`.
+Cancellation is a best-effort request, so the query can still end as `SUCCEEDED` or `FAILED`.
+The `query_id` property keeps the ID of the cancelled query.
+If the cancellation request fails, `asyncio.CancelledError` is raised with the error as its cause.
+Cancelling the task while `execute()` is still starting the query first waits for the start request to finish,
+and then cancels the query it started in the same way.
+If the task is cancelled before `execute()` begins the request, the request is never sent.
+Cancelling the task again during the cancellation request or these waits raises `asyncio.CancelledError` immediately, and the query can keep running.
+With `kill_on_interrupt=False`, `asyncio.CancelledError` is raised immediately, and a query that has already started keeps running.
+
+With `kill_on_interrupt` enabled, a timeout from `asyncio.wait_for()` while `execute()` starts or waits for the query therefore cancels it and raises `asyncio.TimeoutError`.
+If the timeout expires while `execute()` is still looking up a cached result, no query is started and `query_id` is `None`.
+
+```python
+import asyncio
+
+from pyathena import aio_connect
+
+async with await aio_connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
+                          region_name="us-west-2") as conn:
+    async with conn.cursor() as cursor:
+        try:
+            await asyncio.wait_for(cursor.execute("SELECT * FROM many_rows"), timeout=60)
+        except asyncio.TimeoutError:
+            print(f"Query timed out: {cursor.query_id}")
+```
+
 (aio-dict-cursor)=
 
 ## AioDictCursor

docs/spark.md

@@ -512,4 +512,5 @@ With `kill_on_interrupt` enabled, which is the default, cancelling the task whil
 requests cancellation of the calculation, waits until it reaches a terminal state, and then raises `asyncio.CancelledError`.
 Cancelling the task while `execute()` is still starting the calculation first waits for the start request to finish,
 and then cancels the calculation it started in the same way.
+If the task is cancelled before `execute()` begins the request, the request is never sent.
 Cancelling the task again during this wait raises `asyncio.CancelledError` at once without cancelling the calculation.

docs/usage.md

@@ -501,6 +501,40 @@ The `on_start_query_execution` callback is supported by the following cursor typ
 Note: `AsyncCursor` and its variants do not support this callback as they already
 return the query ID immediately through their different execution model.
 
+## Query cancellation on interrupt
+
+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`.
+The `query_id` property keeps the ID of the interrupted query.
+If the cancellation request fails, the `KeyboardInterrupt` propagates with the error as its cause.
+
+A `KeyboardInterrupt` while `execute()` is still starting the query first waits for the
+[StartQueryExecution](https://docs.aws.amazon.com/athena/latest/APIReference/API_StartQueryExecution.html)
+request to finish, and then cancels the query it started in the same way.
+The `query_id` property returns that query's ID.
+If `execute()` has not begun the request when the interrupt is handled, the request is never sent.
+`AsyncCursor` and its variants also stop a query whose start is interrupted in `execute()`,
+but they have no `query_id` property, so the ID of that query is not available.
+They wait for queries on worker threads, which do not receive `KeyboardInterrupt`.
+
+A second `KeyboardInterrupt` during the cancellation request or these waits propagates immediately, and the query can keep running.
+With `kill_on_interrupt=False`, the `KeyboardInterrupt` propagates immediately, and a query that has already started keeps running.
+
+```python
+from pyathena import connect
+
+cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
+                 region_name="us-west-2").cursor()
+try:
+    cursor.execute("SELECT * FROM many_rows")
+except KeyboardInterrupt:
+    print(f"Query {cursor.query_id} was interrupted")
+    raise
+```
+
+For the native asyncio cursors, see {ref}`aio-task-cancellation`.
+
 ## Query polling callback
 
 PyAthena provides an `on_poll` callback that is invoked once per poll iteration with the

pyathena/aio/common.py

@@ -2,7 +2,7 @@
 
 import asyncio
 import logging
-from collections.abc import Awaitable, Callable
+from collections.abc import Awaitable, Callable, Coroutine
 from typing import Any, NoReturn, TypeVar, cast
 
 from botocore.exceptions import BotoCoreError, ClientError
@@ -64,6 +64,8 @@ async def _execute(  # type: ignore[override]
             The query execution ID.
 
         Raises:
+            asyncio.CancelledError: If the task is cancelled while starting the
+                query; see ``_start_execution()``.
             ProgrammingError: If the formatter rejects the query or its parameters.
             DatabaseError: If the ``StartQueryExecution`` request fails.
         """
@@ -88,18 +90,89 @@ async def _execute(  # type: ignore[override]
             cache_expiration_time=options.cache_expiration_time,
         )
         if query_id is None:
+            query_id = await self._start_execution(lambda: self._start_query_execution(request))
+        return query_id
+
+    @override
+    async def _start_query_execution(self, request: dict[str, Any]) -> str:  # type: ignore[override]
+        """Send a ``StartQueryExecution`` request.
+
+        Args:
+            request: The request parameters.
+
+        Returns:
+            The query execution ID.
+
+        Raises:
+            DatabaseError: If the request fails.
+        """
+        try:
+            response = await async_retry_api_call(
+                self._connection.client.start_query_execution,
+                config=self._retry_config,
+                logger=_logger,
+                **request,
+            )
+        except Exception as e:
+            _logger.exception("Failed to execute query.")
+            raise DatabaseError(*e.args) from e
+        return cast(str, response.get("QueryExecutionId"))
+
+    @override
+    async def _start_execution(  # type: ignore[override]
+        self, start: Callable[[], Coroutine[Any, Any, str]]
+    ) -> str:
+        """Send a start request so that task cancellation stops the execution it starts.
+
+        With ``kill_on_interrupt`` enabled, the request runs in a task shielded
+        from task cancellation. On cancellation, the request is abandoned if that
+        task has not begun it by then; it is never sent, and the cancellation
+        propagates. Otherwise the cursor waits for the request to finish, records
+        the execution ID with ``_set_interrupted_execution_id()``, requests
+        cancellation with ``_cancel_and_wait()``, and re-raises
+        ``asyncio.CancelledError``. Another cancellation during that wait
+        propagates at once.
+
+        Args:
+            start: Returns a coroutine that sends the start request and returns the
+                execution ID.
+
+        Returns:
+            The execution ID.
+
+        Raises:
+            asyncio.CancelledError: If the task is cancelled while starting. A
+                failure to start, cancel, or wait for the execution becomes its
+                ``__cause__``.
+            DatabaseError: If the request fails.
+        """
+        if not self._kill_on_interrupt:
+            return await start()
+
+        caller = asyncio.current_task()
+        cancel_requests = caller.cancelling() if caller else 0
+
+        async def run() -> str | None:
+            # Begin the request only if the caller has not been cancelled since.
+            if caller and caller.cancelling() > cancel_requests:
+                return None
+            return await start()
+
+        task = asyncio.ensure_future(run())
+        try:
+            return cast(str, await asyncio.shield(task))
+        except asyncio.CancelledError as cancellation:
             try:
-                response = await async_retry_api_call(
-                    self._connection.client.start_query_execution,
-                    config=self._retry_config,
-                    logger=_logger,
-                    **request,
-                )
-                query_id = response.get("QueryExecutionId")
+                execution_id = await task
+                if execution_id is None:
+                    # The task did not begin the request, so it was never sent.
+                    raise cancellation
+                _logger.warning("Query canceled by user.")
+                self._set_interrupted_execution_id(execution_id)
+                await self._cancel_and_wait(execution_id)
             except Exception as e:
-                _logger.exception("Failed to execute query.")
-                raise DatabaseError(*e.args) from e
-        return query_id
+                raise cancellation from e
+            raise
 
     @override
     async def _get_query_execution(self, query_id: str) -> AthenaQueryExecution:  # type: ignore[override]
@@ -154,10 +227,12 @@ async def _poll_until_terminal(self, query_id: str) -> AthenaQueryExecution:  #
 
     @override
     async def _poll(self, query_id: str) -> AthenaQueryExecution:  # type: ignore[override]
-        """Wait for a query execution to finish.
+        """Wait for a query execution to reach a terminal state.
 
-        On ``asyncio.CancelledError`` with ``kill_on_interrupt`` enabled, stops the
-        query and returns its final execution instead of re-raising.
+        On task cancellation with ``kill_on_interrupt`` enabled, requests
+        cancellation with ``_cancel_and_wait()`` and re-raises
+        ``asyncio.CancelledError``. Cancellation is a best-effort request, so the
+        query can still end as ``SUCCEEDED`` or ``FAILED`` instead of ``CANCELLED``.
 
         Args:
             query_id: The query execution ID.
@@ -166,19 +241,34 @@ async def _poll(self, query_id: str) -> AthenaQueryExecution:  # type: ignore[ov
             The query execution in a terminal state.
 
         Raises:
-            asyncio.CancelledError: If cancelled and ``kill_on_interrupt`` is disabled.
-            OperationalError: If a status or stop request fails.
+            asyncio.CancelledError: If the task is cancelled while waiting. A failure
+                to cancel or wait for the query becomes its ``__cause__``.
+            OperationalError: If a status request fails.
         """
         try:
-            query_execution = await self._poll_until_terminal(query_id)
-        except asyncio.CancelledError:
-            if self._kill_on_interrupt:
-                _logger.warning("Query canceled by user.")
-                await self._cancel(query_id)
-                query_execution = await self._poll_until_terminal(query_id)
-            else:
+            return await self._poll_until_terminal(query_id)
+        except asyncio.CancelledError as cancellation:
+            if not self._kill_on_interrupt:
                 raise
-        return query_execution
+            _logger.warning("Query canceled by user.")
+            try:
+                await self._cancel_and_wait(query_id)
+            except Exception as e:
+                raise cancellation from e
+            raise
+
+    @override
+    async def _cancel_and_wait(self, query_id: str) -> None:  # type: ignore[override]
+        """Request cancellation of a query and wait for a terminal state.
+
+        Args:
+            query_id: The query execution ID.
+
+        Raises:
+            OperationalError: If the cancellation or a status request fails.
+        """
+        await self._cancel(query_id)
+        await self._poll_until_terminal(query_id)
 
     @override
     async def _cancel(self, query_id: str) -> None:  # type: ignore[override]

pyathena/aio/spark/cursor.py

@@ -106,8 +106,10 @@ async def _calculate(  # type: ignore[override]
         retried request returns the calculation an earlier attempt started instead
         of starting another one.
 
-        With ``kill_on_interrupt`` enabled, the request is shielded from task
-        cancellation. On cancellation, waits for the request to finish, requests
+        With ``kill_on_interrupt`` enabled, the request runs in a task shielded
+        from task cancellation. On cancellation, the request is abandoned if that
+        task has not begun it by then; it is never sent, and the cancellation
+        propagates. Otherwise the cursor waits for the request to finish, requests
         cancellation of the calculation it started, waits for a terminal state,
         stores the calculation ID and execution on the cursor, and re-raises
         ``asyncio.CancelledError``. Another cancellation during that wait
@@ -137,14 +139,27 @@ async def _calculate(  # type: ignore[override]
         if not self._kill_on_interrupt:
             return await self._start_calculation_execution(request)
 
-        start = asyncio.ensure_future(self._start_calculation_execution(request))
+        caller = asyncio.current_task()
+        cancel_requests = caller.cancelling() if caller else 0
+
+        async def run() -> str | None:
+            # Begin the request only if the caller has not been cancelled since.
+            if caller and caller.cancelling() > cancel_requests:
+                return None
+            return await self._start_calculation_execution(request)
+
+        start = asyncio.ensure_future(run())
         try:
-            return await asyncio.shield(start)
+            return cast(str, await asyncio.shield(start))
         except asyncio.CancelledError as cancellation:
-            _logger.warning("Query canceled by user.")
             try:
-                self._calculation_id = await start
-                await self._cancel_and_wait(self._calculation_id)
+                calculation_id = await start
+                if calculation_id is None:
+                    # The task did not begin the request, so it was never sent.
+                    raise cancellation
+                _logger.warning("Query canceled by user.")
+                self._calculation_id = calculation_id
+                await self._cancel_and_wait(calculation_id)
             except Exception as e:
                 raise cancellation from e
             raise
@@ -340,6 +355,9 @@ async def execute(
         Returns:
             Self reference for method chaining.
         """
+        # A failure below must not leave the previous calculation on the cursor.
+        self._calculation_id = None
+        self._calculation_execution = None
         self._calculation_id = await self._calculate(
             session_id=session_id if session_id else self._session_id,
             code_block=operation,