Merge pull request #668 from pyathena-dev/feature/native-asyncio-cursor-phase3

Author: laughingman7743

pyathena/aio/arrow/cursor.py

@@ -3,9 +3,9 @@
 
 import asyncio
 import logging
-from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Union, cast
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union, cast
 
-from pyathena.aio.common import AioBaseCursor
+from pyathena.aio.common import WithAsyncFetch
 from pyathena.arrow.converter import (
     DefaultArrowTypeConverter,
     DefaultArrowUnloadTypeConverter,
@@ -14,7 +14,6 @@
 from pyathena.common import CursorIterator
 from pyathena.error import OperationalError, ProgrammingError
 from pyathena.model import AthenaCompression, AthenaFileFormat, AthenaQueryExecution
-from pyathena.result_set import WithResultSet
 
 if TYPE_CHECKING:
     import polars as pl
@@ -23,7 +22,7 @@
 _logger = logging.getLogger(__name__)  # type: ignore
 
 
-class AioArrowCursor(AioBaseCursor, CursorIterator, WithResultSet):
+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.
@@ -50,7 +49,6 @@ def __init__(
         unload: bool = False,
         result_reuse_enable: bool = False,
         result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES,
-        on_start_query_execution: Optional[Callable[[str], None]] = None,
         connect_timeout: Optional[float] = None,
         request_timeout: Optional[float] = None,
         **kwargs,
@@ -69,10 +67,8 @@ def __init__(
             **kwargs,
         )
         self._unload = unload
-        self._on_start_query_execution = on_start_query_execution
         self._connect_timeout = connect_timeout
         self._request_timeout = request_timeout
-        self._query_id: Optional[str] = None
         self._result_set: Optional[AthenaArrowResultSet] = None
 
     @staticmethod
@@ -83,45 +79,6 @@ def get_default_converter(
             return DefaultArrowUnloadTypeConverter()
         return DefaultArrowTypeConverter()
 
-    @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[AthenaArrowResultSet]:
-        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 execute(  # type: ignore[override]
         self,
         operation: str,
@@ -133,7 +90,6 @@ async def execute(  # type: ignore[override]
         result_reuse_enable: Optional[bool] = None,
         result_reuse_minutes: Optional[int] = None,
         paramstyle: Optional[str] = None,
-        on_start_query_execution: Optional[Callable[[str], None]] = None,
         **kwargs,
     ) -> "AioArrowCursor":
         """Execute a SQL query asynchronously and return results as Arrow Tables.
@@ -148,7 +104,6 @@ async def execute(  # type: ignore[override]
             result_reuse_enable: Enable Athena result reuse for this query.
             result_reuse_minutes: Minutes to reuse cached results.
             paramstyle: Parameter style ('qmark' or 'pyformat').
-            on_start_query_execution: Callback called when query starts.
             **kwargs: Additional execution parameters.
 
         Returns:
@@ -179,10 +134,6 @@ async def execute(  # type: ignore[override]
             paramstyle=paramstyle,
         )
 
-        if self._on_start_query_execution:
-            self._on_start_query_execution(self.query_id)
-        if on_start_query_execution:
-            on_start_query_execution(self.query_id)
         query_execution = await self._poll(self.query_id)
         if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED:
             self.result_set = await asyncio.to_thread(
@@ -202,84 +153,6 @@ async def execute(  # type: ignore[override]
             raise OperationalError(query_execution.state_change_reason)
         return self
 
-    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)
-        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(AthenaArrowResultSet, 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(AthenaArrowResultSet, 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(AthenaArrowResultSet, self.result_set)
-        return result_set.fetchall()
-
     def as_arrow(self) -> "Table":
         """Return query results as an Apache Arrow Table.
 
@@ -301,18 +174,3 @@ def as_polars(self) -> "pl.DataFrame":
             raise ProgrammingError("No result set.")
         result_set = cast(AthenaArrowResultSet, self.result_set)
         return result_set.as_polars()
-
-    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

@@ -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,154 @@ 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
+
+    @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()