Break result sets after failed or cancelled fetches and drop the cross-thread guards

docs/aio.md

@@ -171,7 +171,10 @@ No further attempt is made, and the wait before the next attempt ends.
 A request that has already been sent still runs to its end in its thread.
 The requests that stop a query or a calculation and the request that terminates a Spark session keep their retries.
 The start request of `execute()` with `kill_on_interrupt` enabled keeps its retries after a cancellation, and stops them when the task is cancelled again.
-The fetch methods of `AioPandasCursor`, `AioArrowCursor`, `AioPolarsCursor`, and `AioS3FSCursor` run to their end in their thread.
+After a fetch is cancelled, later fetches of its result raise `ProgrammingError`.
+The fetch methods of `AioPandasCursor`, `AioArrowCursor`, `AioPolarsCursor`, and `AioS3FSCursor` run in a thread.
+A cancellation skips a fetch that has not started, and a started fetch keeps running in its thread without the retries of its requests.
+The cursor drops the result, which is closed once its fetch is no longer running.
 
 (aio-dict-cursor)=
 
@@ -227,7 +230,7 @@ All aio cursors use `await` for fetch operations, so fetching does not block the
   Parquet) inside `execute()`, wrapped in `asyncio.to_thread()`. When CSV results are read in
   chunks (`chunksize` for pandas and Polars, or a chunk size chosen by `auto_optimize_chunksize`
   for pandas), or with `chunksize` on Polars UNLOAD results, fetch calls read S3 lazily instead.
-  The fetch methods are also wrapped in `asyncio.to_thread()`.
+  The fetch methods also run in a thread.
 - `AioS3FSCursor` streams rows from the result file in S3 as they are fetched.
 
 With managed query result storage, where the query has no S3 output location, the pandas, Arrow,

docs/arrow.md

@@ -485,8 +485,8 @@ cursor = connect(
 ## AioArrowCursor
 
 AioArrowCursor is a native asyncio cursor that returns results as Apache Arrow Tables.
-Unlike AsyncArrowCursor which uses `concurrent.futures`, this cursor uses
-`asyncio.to_thread()` for both result set creation and fetch operations,
+Unlike AsyncArrowCursor which uses `concurrent.futures`, this cursor runs
+both result set creation and fetch operations in threads,
 keeping the event loop free.
 
 ```python

docs/pandas.md

@@ -825,8 +825,8 @@ cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
 ## AioPandasCursor
 
 AioPandasCursor is a native asyncio cursor that returns results as pandas DataFrames.
-Unlike AsyncPandasCursor which uses `concurrent.futures`, this cursor uses
-`asyncio.to_thread()` for both result set creation and fetch operations,
+Unlike AsyncPandasCursor which uses `concurrent.futures`, this cursor runs
+both result set creation and fetch operations in threads,
 keeping the event loop free.
 
 ```python

docs/polars.md

@@ -618,8 +618,8 @@ for chunk in result_set.iter_chunks():
 ## AioPolarsCursor
 
 AioPolarsCursor is a native asyncio cursor that returns results as Polars DataFrames.
-Unlike AsyncPolarsCursor which uses `concurrent.futures`, this cursor uses
-`asyncio.to_thread()` for both result set creation and fetch operations,
+Unlike AsyncPolarsCursor which uses `concurrent.futures`, this cursor runs
+both result set creation and fetch operations in threads,
 keeping the event loop free.
 
 ```python

docs/s3fs.md

@@ -410,8 +410,8 @@ cursor.cancel(query_id)
 ## AioS3FSCursor
 
 AioS3FSCursor is a native asyncio cursor that uses the same lightweight CSV parsing as S3FSCursor.
-Unlike AsyncS3FSCursor which uses `concurrent.futures`, this cursor uses
-`asyncio.to_thread()` for both result set creation and fetch operations,
+Unlike AsyncS3FSCursor which uses `concurrent.futures`, this cursor runs
+both result set creation and fetch operations in threads,
 keeping the event loop free.
 
 Since `AthenaS3FSResultSet` lazily streams rows from S3 via a CSV reader,

docs/usage.md

@@ -130,6 +130,13 @@ for row in cursor:
     print(row)
 ```
 
+After a fetch raises an exception, including `KeyboardInterrupt`, later fetches of the same result raise `ProgrammingError`.
+An invalid argument, such as a `fetchmany()` size that is not an integer, raises its error without that effect.
+Execute the query again to read its result.
+
+As `pyathena.threadsafety = 2` declares, threads can share connections, but not cursors.
+Calling the methods of one cursor or its result set from several threads, including `close()` and `cancel()`, is not supported.
+
 (usage-query-with-parameters)=
 
 ## Query with parameters
@@ -399,7 +406,7 @@ Result columns from a passthrough query come back with the source system's types
 
 PyAthena provides a callback mechanism that allows you to get immediate access to the query ID
 as soon as the `start_query_execution` API call is made, before waiting for query completion.
-This is useful for monitoring, logging, or cancelling long-running queries from another thread.
+This is useful for monitoring, logging, or stopping long-running queries by their ID, for example with `AsyncCursor.cancel(query_id)`.
 When `cache_size` finds a reusable query, no new query starts, and the callback receives the reused query's ID.
 
 The `on_start_query_execution` callback can be configured at both the connection level and
@@ -446,91 +453,6 @@ cursor.execute(
 )
 ```
 
-### Query cancellation example
-
-A common use case is to cancel long-running analytical queries after a timeout:
-
-```python
-import threading
-from pyathena import connect
-
-def cancel_long_running_query():
-    """Example: Cancel a complex analytical query after 10 minutes."""
-    timeout_minutes = 10
-
-    def track_query_start(query_id):
-        print(f"Long-running analysis started: {query_id}")
-
-    def cancel_on_timeout(cursor):
-        """Cancel the query that is still running when the timer fires."""
-        try:
-            cursor.cancel()
-            print(f"Query cancelled after {timeout_minutes} minutes timeout")
-        except Exception as e:
-            print(f"Cancellation failed: {e}")
-
-    cursor = connect(
-        s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
-        region_name="us-west-2",
-        on_start_query_execution=track_query_start
-    ).cursor()
-
-    # Complex analytical query that might run for a long time
-    long_query = """
-    WITH daily_metrics AS (
-        SELECT
-            date_trunc('day', timestamp_col) as day,
-            user_id,
-            COUNT(*) as events,
-            AVG(duration) as avg_duration
-        FROM large_events_table
-        WHERE timestamp_col >= current_date - interval '1' year
-        GROUP BY 1, 2
-    ),
-    user_segments AS (
-        SELECT
-            user_id,
-            CASE
-                WHEN AVG(events) > 100 THEN 'high_activity'
-                WHEN AVG(events) > 10 THEN 'medium_activity'
-                ELSE 'low_activity'
-            END as segment
-        FROM daily_metrics
-        GROUP BY user_id
-    )
-    SELECT
-        segment,
-        COUNT(DISTINCT dm.user_id) as users,
-        AVG(events) as avg_daily_events
-    FROM daily_metrics dm
-    JOIN user_segments us ON dm.user_id = us.user_id
-    GROUP BY segment
-    ORDER BY avg_daily_events DESC
-    """
-
-    # Cancel the query if it is still running after the timeout
-    timer = threading.Timer(timeout_minutes * 60, cancel_on_timeout, args=(cursor,))
-    timer.start()
-    try:
-        print("Starting complex analytical query (10-minute timeout)...")
-        cursor.execute(long_query)
-
-        # Process results
-        results = cursor.fetchall()
-        print(f"Analysis completed successfully: {len(results)} segments found")
-        for row in results:
-            print(f"  {row[0]}: {row[1]} users, {row[2]:.1f} avg events")
-
-    except Exception as e:
-        print(f"Query failed or was cancelled: {e}")
-    finally:
-        # Stop the timer once the query has finished
-        timer.cancel()
-
-# Run the example
-cancel_long_running_query()
-```
-
 ### Multiple callbacks
 
 When both connection-level and execute-level callbacks are specified,

pyathena/aio/common.py

@@ -4,12 +4,13 @@
 
 import asyncio
 import logging
+import operator
 from collections.abc import Awaitable, Callable, Coroutine
 from typing import Any, NoReturn, TypeVar, cast
 
 from botocore.exceptions import BotoCoreError, ClientError
 
-from pyathena.aio.util import async_retry_api_call, close_in_thread
+from pyathena.aio.util import async_retry_api_call, close_in_thread, run_in_thread_or_close
 from pyathena.common import BaseCursor, CursorIterator
 from pyathena.error import DatabaseError, OperationalError, ProgrammingError
 from pyathena.glue import GlueMetadataClient
@@ -709,11 +710,11 @@ class WithAsyncFetch(WithResultSet, AioBaseCursor, CursorIterator):
     Combines ``WithResultSet`` with ``AioBaseCursor`` and ``CursorIterator``,
     and provides async fetch, ``executemany``, and ``cancel``, async
     iteration, and the async context manager protocol. The fetch methods run
-    the result set's synchronous fetch with ``asyncio.to_thread``; a subclass
-    whose result set fetches asynchronously overrides them. Closing the
+    the result set's synchronous fetch in a thread (see ``_fetch_in_thread()``);
+    a subclass whose result set fetches asynchronously overrides them. Closing the
     previous result set when a query is executed, and closing the cursor at
     the end of an ``async with`` block, also run in a thread, because closing
-    a result stream can wait for a read in progress. Synchronous iteration
+    a result stream can wait for its read-ahead requests in progress. Synchronous iteration
     raises ``TypeError``.
 
     Subclasses override ``execute()`` and optionally ``__init__`` and
@@ -793,28 +794,25 @@ async def fetchone(
     ) -> tuple[Any | None, ...] | dict[Any, Any | None] | None:
         """Fetch the next row of the result set.
 
-        Wraps the synchronous fetch in ``asyncio.to_thread`` to avoid
-        blocking the event loop.
+        Runs the synchronous fetch in a thread to avoid blocking the event
+        loop (see ``_fetch_in_thread()``).
 
         Returns:
             A tuple representing the next row, or None if no more rows.
 
         Raises:
             ProgrammingError: If no result set is available.
         """
-        if not self.has_result_set:
-            raise ProgrammingError("No result set.")
-        result_set = cast(AthenaResultSet, self.result_set)
-        return await asyncio.to_thread(result_set.fetchone)
+        return await self._fetch_in_thread(operator.methodcaller("fetchone"))
 
     @override
     async def fetchmany(
         self, size: int | None = None
     ) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]:
         """Fetch multiple rows from the result set.
 
-        Wraps the synchronous fetch in ``asyncio.to_thread`` to avoid
-        blocking the event loop.
+        Runs the synchronous fetch in a thread to avoid blocking the event
+        loop (see ``_fetch_in_thread()``).
 
         Args:
             size: Maximum number of rows to fetch. If None or not positive,
@@ -826,30 +824,51 @@ async def fetchmany(
         Raises:
             ProgrammingError: If no result set is available.
         """
-        if not self.has_result_set:
-            raise ProgrammingError("No result set.")
-        result_set = cast(AthenaResultSet, self.result_set)
-        return await asyncio.to_thread(result_set.fetchmany, size)
+        return await self._fetch_in_thread(operator.methodcaller("fetchmany", size))
 
     @override
     async def fetchall(
         self,
     ) -> list[tuple[Any | None, ...] | dict[Any, Any | None]]:
         """Fetch all remaining rows from the result set.
 
-        Wraps the synchronous fetch in ``asyncio.to_thread`` to avoid
-        blocking the event loop.
+        Runs the synchronous fetch in a thread to avoid blocking the event
+        loop (see ``_fetch_in_thread()``).
 
         Returns:
             List of tuples representing all remaining rows.
 
+        Raises:
+            ProgrammingError: If no result set is available.
+        """
+        return await self._fetch_in_thread(operator.methodcaller("fetchall"))
+
+    async def _fetch_in_thread(self, fetch: Callable[[AthenaResultSet], _T]) -> _T:
+        """Run a synchronous fetch of the result set in a thread.
+
+        A cancellation stops the retries of the AWS API calls that the fetch makes
+        and detaches the result set from the cursor, so that later fetches raise
+        ``ProgrammingError``. The result set is closed once the fetch is no longer
+        running (see :func:`~pyathena.aio.util.run_in_thread_or_close`).
+
+        Args:
+            fetch: The function that fetches from the result set.
+
+        Returns:
+            The return value of ``fetch``.
+
         Raises:
             ProgrammingError: If no result set is available.
         """
         if not self.has_result_set:
             raise ProgrammingError("No result set.")
         result_set = cast(AthenaResultSet, self.result_set)
-        return await asyncio.to_thread(result_set.fetchall)
+        try:
+            return await run_in_thread_or_close(result_set, fetch, result_set)
+        except asyncio.CancelledError:
+            if self.result_set is result_set:
+                self.result_set = None
+            raise
 
     @override
     def __iter__(self) -> NoReturn: