Rename AioCursorBase to WithAsyncFetch and consolidate into common.py

Move the shared SQL cursor mixin from aio/base.py into aio/common.py
and rename from AioCursorBase to WithAsyncFetch to follow the existing
WithXXX naming convention (WithResultSet, WithCalculationExecution)
where XXX describes the functionality provided.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

Author: laughingman7743

pyathena/aio/arrow/cursor.py

@@ -5,7 +5,7 @@
 import logging
 from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union, cast
 
-from pyathena.aio.base import AioCursorBase
+from pyathena.aio.common import WithAsyncFetch
 from pyathena.arrow.converter import (
     DefaultArrowTypeConverter,
     DefaultArrowUnloadTypeConverter,
@@ -22,7 +22,7 @@
 _logger = logging.getLogger(__name__)  # type: ignore
 
 
-class AioArrowCursor(AioCursorBase):
+class AioArrowCursor(WithAsyncFetch):
     """Native asyncio cursor that returns results as Apache Arrow Tables.
 
     Uses ``asyncio.to_thread()`` to create the result set off the event loop.

pyathena/aio/base.py

@@ -5,12 +5,13 @@
 import logging
 import sys
 from datetime import datetime, timedelta, timezone
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union, cast
 
 from pyathena.aio.util import async_retry_api_call
-from pyathena.common import BaseCursor
-from pyathena.error import DatabaseError, OperationalError
+from pyathena.common import BaseCursor, CursorIterator
+from pyathena.error import DatabaseError, OperationalError, ProgrammingError
 from pyathena.model import AthenaDatabase, AthenaQueryExecution, AthenaTableMetadata
+from pyathena.result_set import AthenaResultSet, WithResultSet
 
 _logger = logging.getLogger(__name__)  # type: ignore
 
@@ -346,3 +347,155 @@ async def list_table_metadata(  # type: ignore[override]
             if not next_token:
                 break
         return metadata
+
+
+class WithAsyncFetch(AioBaseCursor, CursorIterator, WithResultSet):
+    """Mixin providing shared fetch, lifecycle, and async protocol for SQL cursors.
+
+    Provides properties (``arraysize``, ``result_set``, ``query_id``,
+    ``rownumber``, ``rowcount``), lifecycle methods (``close``, ``executemany``,
+    ``cancel``), default sync fetch (for cursors whose result sets load all
+    data eagerly in ``__init__``), and the async iteration protocol.
+
+    Subclasses override ``execute()`` and optionally ``__init__`` and
+    format-specific helpers.
+    """
+
+    def __init__(self, **kwargs) -> None:
+        super().__init__(**kwargs)
+        self._query_id: Optional[str] = None
+        self._result_set: Optional[AthenaResultSet] = None
+        self._on_start_query_execution = kwargs.get("on_start_query_execution")
+
+    @property
+    def arraysize(self) -> int:
+        return self._arraysize
+
+    @arraysize.setter
+    def arraysize(self, value: int) -> None:
+        if value <= 0:
+            raise ProgrammingError("arraysize must be a positive integer value.")
+        self._arraysize = value
+
+    @property  # type: ignore
+    def result_set(self) -> Optional[AthenaResultSet]:
+        return self._result_set
+
+    @result_set.setter
+    def result_set(self, val) -> None:
+        self._result_set = val
+
+    @property
+    def query_id(self) -> Optional[str]:
+        return self._query_id
+
+    @query_id.setter
+    def query_id(self, val) -> None:
+        self._query_id = val
+
+    @property
+    def rownumber(self) -> Optional[int]:
+        return self.result_set.rownumber if self.result_set else None
+
+    @property
+    def rowcount(self) -> int:
+        return self.result_set.rowcount if self.result_set else -1
+
+    def close(self) -> None:
+        """Close the cursor and release associated resources."""
+        if self.result_set and not self.result_set.is_closed:
+            self.result_set.close()
+
+    async def executemany(  # type: ignore[override]
+        self,
+        operation: str,
+        seq_of_parameters: List[Optional[Union[Dict[str, Any], List[str]]]],
+        **kwargs,
+    ) -> None:
+        """Execute a SQL query multiple times with different parameters.
+
+        Args:
+            operation: SQL query string to execute.
+            seq_of_parameters: Sequence of parameter sets, one per execution.
+            **kwargs: Additional keyword arguments passed to each ``execute()``.
+        """
+        for parameters in seq_of_parameters:
+            await self.execute(operation, parameters, **kwargs)
+        # Operations that have result sets are not allowed with executemany.
+        self._reset_state()
+
+    async def cancel(self) -> None:
+        """Cancel the currently executing query.
+
+        Raises:
+            ProgrammingError: If no query is currently executing.
+        """
+        if not self.query_id:
+            raise ProgrammingError("QueryExecutionId is none or empty.")
+        await self._cancel(self.query_id)
+
+    def fetchone(
+        self,
+    ) -> Optional[Union[Tuple[Optional[Any], ...], Dict[Any, Optional[Any]]]]:
+        """Fetch the next row of the result set.
+
+        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 result_set.fetchone()
+
+    def fetchmany(
+        self, size: Optional[int] = None
+    ) -> List[Union[Tuple[Optional[Any], ...], Dict[Any, Optional[Any]]]]:
+        """Fetch multiple rows from the result set.
+
+        Args:
+            size: Maximum number of rows to fetch. Defaults to arraysize.
+
+        Returns:
+            List of tuples representing the fetched 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 result_set.fetchmany(size)
+
+    def fetchall(
+        self,
+    ) -> List[Union[Tuple[Optional[Any], ...], Dict[Any, Optional[Any]]]]:
+        """Fetch all remaining rows from the result set.
+
+        Returns:
+            List of tuples representing all remaining 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 result_set.fetchall()
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        row = self.fetchone()
+        if row is None:
+            raise StopAsyncIteration
+        return row
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        self.close()

pyathena/aio/common.py

@@ -4,7 +4,7 @@
 import logging
 from typing import Any, Callable, Dict, List, Optional, Union, cast
 
-from pyathena.aio.base import AioCursorBase
+from pyathena.aio.common import WithAsyncFetch
 from pyathena.aio.result_set import AthenaAioDictResultSet, AthenaAioResultSet
 from pyathena.common import CursorIterator
 from pyathena.error import OperationalError, ProgrammingError
@@ -13,7 +13,7 @@
 _logger = logging.getLogger(__name__)  # type: ignore
 
 
-class AioCursor(AioCursorBase):
+class AioCursor(WithAsyncFetch):
     """Native asyncio cursor for Amazon Athena.
 
     Unlike ``AsyncCursor`` (which uses ``ThreadPoolExecutor``), this cursor

pyathena/aio/cursor.py

@@ -16,7 +16,7 @@
     cast,
 )
 
-from pyathena.aio.base import AioCursorBase
+from pyathena.aio.common import WithAsyncFetch
 from pyathena.common import CursorIterator
 from pyathena.error import OperationalError, ProgrammingError
 from pyathena.model import AthenaCompression, AthenaFileFormat, AthenaQueryExecution
@@ -32,7 +32,7 @@
 _logger = logging.getLogger(__name__)  # type: ignore
 
 
-class AioPandasCursor(AioCursorBase):
+class AioPandasCursor(WithAsyncFetch):
     """Native asyncio cursor that returns results as pandas DataFrames.
 
     Uses ``asyncio.to_thread()`` to create the result set off the event loop.