Cancel a Spark calculation interrupted while it is being started

docs/spark.md

@@ -285,6 +285,18 @@ requests cancellation, waits until the calculation reaches a terminal state, and
 The `state` property returns that terminal state.
 If the cancellation request fails, the `KeyboardInterrupt` propagates with the error as its cause.
 
+A `KeyboardInterrupt` while `execute()` is still starting the calculation first waits for the
+[StartCalculationExecution](https://docs.aws.amazon.com/athena/latest/APIReference/API_StartCalculationExecution.html)
+request to finish, and then cancels the calculation it started in the same way.
+The `calculation_id` property returns that calculation's ID.
+A second `KeyboardInterrupt` during this wait propagates at once without cancelling the calculation.
+A cancellation request sent right after a calculation starts can occasionally have no effect, so the calculation can still end in the `COMPLETED` state.
+
+Unless `client_request_token` is passed to `execute()`, the cursor sends a generated `ClientRequestToken` with each calculation.
+A retried start request then returns the calculation that an earlier attempt started instead of starting another one.
+A token passed to `execute()` must be unique for each calculation:
+Athena returns the earlier calculation for a reused token, even when the code differs.
+
 (async-spark-cursor)=
 
 ## AsyncSparkCursor
@@ -498,3 +510,6 @@ async with await aio_connect(work_group="YOUR_SPARK_WORKGROUP",
 
 With `kill_on_interrupt` enabled, which is the default, cancelling the task while `execute()` waits for the calculation
 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.
+Cancelling the task again during this wait raises `asyncio.CancelledError` at once without cancelling the calculation.

pyathena/aio/spark/cursor.py

@@ -9,6 +9,7 @@
 
 import asyncio
 import logging
+import uuid
 from typing import Any, cast
 
 from pyathena.aio.util import async_retry_api_call
@@ -133,24 +134,78 @@ async def _calculate(  # type: ignore[override]
         description: str | None = None,
         client_request_token: str | None = None,
     ) -> str:
+        """Start a calculation execution with ``StartCalculationExecution``.
+
+        Without ``client_request_token``, a generated token is sent, so that a
+        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
+        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
+        propagates at once.
+
+        Args:
+            session_id: The session ID.
+            code_block: The code to run.
+            description: The calculation description.
+            client_request_token: The idempotency token of the request.
+
+        Returns:
+            The calculation execution ID.
+
+        Raises:
+            asyncio.CancelledError: If the task is cancelled while starting the
+                calculation. A failure to start, cancel, or wait for the
+                calculation becomes its ``__cause__``.
+            DatabaseError: If the request fails.
+        """
         request = self._build_start_calculation_execution_request(
             session_id=session_id,
             code_block=code_block,
             description=description,
-            client_request_token=client_request_token,
+            client_request_token=client_request_token or str(uuid.uuid4()),
         )
+        if not self._kill_on_interrupt:
+            return await self.__start_calculation(request)
+
+        start = asyncio.ensure_future(self.__start_calculation(request))
+        try:
+            return 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)
+            except Exception as e:
+                raise cancellation from e
+            raise
+
+    async def __start_calculation(self, request: dict[str, Any]) -> str:
+        """Send a ``StartCalculationExecution`` request.
+
+        Args:
+            request: The request parameters.
+
+        Returns:
+            The calculation execution ID.
+
+        Raises:
+            DatabaseError: If the request fails.
+        """
         try:
             response = await async_retry_api_call(
                 self._connection.client.start_calculation_execution,
                 config=self._retry_config,
                 logger=_logger,
                 **request,
             )
-            calculation_id = response.get("CalculationExecutionId")
         except Exception as e:
             _logger.exception("Failed to execute calculation.")
             raise DatabaseError(*e.args) from e
-        return cast(str, calculation_id)
+        return cast(str, response.get("CalculationExecutionId"))
 
     async def __poll(self, query_id: str) -> AthenaQueryExecution | AthenaCalculationExecution:
         while True:
@@ -195,14 +250,25 @@ async def _poll(  # type: ignore[override]
                 raise
             _logger.warning("Query canceled by user.")
             try:
-                await self._cancel(query_id)
-                self._calculation_execution = cast(
-                    AthenaCalculationExecution, await self.__poll(query_id)
-                )
+                await self.__cancel_and_wait(query_id)
             except Exception as e:
                 raise cancellation from e
             raise
 
+    async def __cancel_and_wait(self, calculation_id: str) -> None:
+        """Request cancellation and store the calculation's terminal state.
+
+        Args:
+            calculation_id: The calculation execution ID.
+
+        Raises:
+            OperationalError: If the cancellation or a status request fails.
+        """
+        await self._cancel(calculation_id)
+        self._calculation_execution = cast(
+            AthenaCalculationExecution, await self.__poll(calculation_id)
+        )
+
     async def _cancel(self, query_id: str) -> None:  # type: ignore[override]
         request: dict[str, Any] = {"CalculationExecutionId": query_id}
         try:

pyathena/spark/common.py

@@ -9,14 +9,17 @@
 
 import contextlib
 import logging
+import threading
 import time
+import uuid
 from abc import ABCMeta, abstractmethod
+from concurrent.futures import Future, wait
 from datetime import datetime
 from typing import Any, cast
 
 import botocore
 
-from pyathena import NotSupportedError, OperationalError
+from pyathena import DatabaseError, NotSupportedError, OperationalError
 from pyathena.common import BaseCursor
 from pyathena.model import (
     AthenaCalculationExecution,
@@ -28,6 +31,11 @@
 
 _logger = logging.getLogger(__name__)
 
+# 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).
+_INTERRUPT_CHECK_INTERVAL = 0.1
+
 
 class SparkBaseCursor(BaseCursor, metaclass=ABCMeta):
     """Abstract base class for Spark-enabled cursor implementations.
@@ -340,14 +348,136 @@ def _poll(self, query_id: str) -> AthenaQueryExecution | AthenaCalculationExecut
                 raise
             _logger.warning("Query canceled by user.")
             try:
-                self._cancel(query_id)
-                self._calculation_execution = cast(
-                    AthenaCalculationExecution, self.__poll(query_id)
-                )
+                self.__cancel_and_wait(query_id)
             except Exception as e:
                 raise interrupt from e
             raise
 
+    def __cancel_and_wait(self, calculation_id: str) -> None:
+        """Request cancellation and store the calculation's terminal state.
+
+        Args:
+            calculation_id: The calculation execution ID.
+
+        Raises:
+            OperationalError: If the cancellation or a status request fails.
+        """
+        self._cancel(calculation_id)
+        self._calculation_execution = cast(AthenaCalculationExecution, self.__poll(calculation_id))
+
+    def _calculate(
+        self,
+        session_id: str,
+        code_block: str,
+        description: str | None = None,
+        client_request_token: str | None = None,
+    ) -> str:
+        """Start a calculation execution with ``StartCalculationExecution``.
+
+        Without ``client_request_token``, a generated token is sent, so that a
+        retried request returns the calculation an earlier attempt started instead
+        of starting another one.
+
+        With ``kill_on_interrupt`` enabled, the request runs on a helper thread.
+        On ``KeyboardInterrupt``, the cursor first tries to abandon the request.
+        This succeeds only if the helper has not begun the request by then; the
+        helper then never sends it, and the interrupt 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 the interrupt. Another
+        ``KeyboardInterrupt`` during that wait propagates at once.
+
+        Args:
+            session_id: The session ID.
+            code_block: The code to run.
+            description: The calculation description.
+            client_request_token: The idempotency token of the request.
+
+        Returns:
+            The calculation execution ID.
+
+        Raises:
+            KeyboardInterrupt: If interrupted while starting the calculation. A
+                failure to start, cancel, or wait for the calculation becomes its
+                ``__cause__``.
+            DatabaseError: If the request fails.
+        """
+        request = self._build_start_calculation_execution_request(
+            session_id=session_id,
+            code_block=code_block,
+            description=description,
+            client_request_token=client_request_token or str(uuid.uuid4()),
+        )
+        if not self._kill_on_interrupt:
+            return self.__start_calculation(request)
+
+        future: Future[str] = Future()
+
+        def start() -> None:
+            # Begin the request only if no interrupt has given up on it yet.
+            if not future.set_running_or_notify_cancel():
+                return
+            try:
+                future.set_result(self.__start_calculation(request))
+            except BaseException as e:
+                future.set_exception(e)
+
+        try:
+            threading.Thread(target=start, name="pyathena-spark-start", daemon=True).start()
+            return self.__wait_for_start(future)
+        except KeyboardInterrupt as interrupt:
+            if future.cancel():
+                # The helper has not begun the request and never will.
+                raise
+            _logger.warning("Query canceled by user.")
+            try:
+                self._calculation_id = self.__wait_for_start(future)
+                self.__cancel_and_wait(self._calculation_id)
+            except Exception as e:
+                raise interrupt from e
+            raise
+
+    def __start_calculation(self, request: dict[str, Any]) -> str:
+        """Send a ``StartCalculationExecution`` request.
+
+        Args:
+            request: The request parameters.
+
+        Returns:
+            The calculation execution ID.
+
+        Raises:
+            DatabaseError: If the request fails.
+        """
+        try:
+            response = retry_api_call(
+                self._connection.client.start_calculation_execution,
+                config=self._retry_config,
+                logger=_logger,
+                **request,
+            )
+        except Exception as e:
+            _logger.exception("Failed to execute calculation.")
+            raise DatabaseError(*e.args) from e
+        return cast(str, response.get("CalculationExecutionId"))
+
+    @staticmethod
+    def __wait_for_start(future: Future[str]) -> str:
+        """Wait for the start request on a helper thread to finish.
+
+        Args:
+            future: The future of the start request.
+
+        Returns:
+            The calculation execution ID.
+
+        Raises:
+            DatabaseError: If the request failed.
+        """
+        while not future.done():
+            wait((future,), timeout=_INTERRUPT_CHECK_INTERVAL)
+        return future.result()
+
     def _cancel(self, query_id: str) -> None:
         """Stop a calculation execution with ``StopCalculationExecution``.