from __future__ import annotations
import asyncio
import logging
from collections.abc import Callable
from typing import Any, cast
from pyathena.aio.common import WithAsyncFetch
from pyathena.common import CursorIterator
from pyathena.error import OperationalError, ProgrammingError
from pyathena.filesystem.s3_async import AioS3FileSystem
from pyathena.model import AthenaQueryExecution
from pyathena.options import ExecuteOptions
from pyathena.s3fs.converter import DefaultS3FSTypeConverter
from pyathena.s3fs.result_set import AthenaS3FSResultSet, CSVReaderType
_logger = logging.getLogger(__name__)
[docs]
class AioS3FSCursor(WithAsyncFetch):
"""Native asyncio cursor that reads CSV results via AioS3FileSystem.
Uses ``AioS3FileSystem`` for S3 operations, which replaces
``ThreadPoolExecutor`` parallelism with ``asyncio.gather`` +
``asyncio.to_thread``. Fetch operations are wrapped in
``asyncio.to_thread()`` because CSV reading is blocking I/O.
Example:
>>> async with await pyathena.aio_connect(...) as conn:
... cursor = conn.cursor(AioS3FSCursor)
... await cursor.execute("SELECT * FROM my_table")
... row = await cursor.fetchone()
"""
[docs]
def __init__(
self,
s3_staging_dir: str | None = None,
schema_name: str | None = None,
catalog_name: str | None = None,
work_group: str | None = None,
poll_interval: float = 1,
encryption_option: str | None = None,
kms_key: str | None = None,
kill_on_interrupt: bool = True,
result_reuse_enable: bool = False,
result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES,
csv_reader: CSVReaderType | None = None,
**kwargs,
) -> None:
super().__init__(
s3_staging_dir=s3_staging_dir,
schema_name=schema_name,
catalog_name=catalog_name,
work_group=work_group,
poll_interval=poll_interval,
encryption_option=encryption_option,
kms_key=kms_key,
kill_on_interrupt=kill_on_interrupt,
result_reuse_enable=result_reuse_enable,
result_reuse_minutes=result_reuse_minutes,
**kwargs,
)
self._csv_reader = csv_reader
self._result_set: AthenaS3FSResultSet | None = None
[docs]
@staticmethod
def get_default_converter(
unload: bool = False,
) -> DefaultS3FSTypeConverter:
"""Get the default type converter for S3FS cursor.
Args:
unload: Unused. S3FS cursor does not support UNLOAD operations.
Returns:
DefaultS3FSTypeConverter instance.
"""
return DefaultS3FSTypeConverter()
[docs]
async def execute( # type: ignore[override]
self,
operation: str,
parameters: dict[str, Any] | list[str] | None = None,
work_group: str | None = None,
s3_staging_dir: str | None = None,
cache_size: int | None = None,
cache_expiration_time: int | None = None,
result_reuse_enable: bool | None = None,
result_reuse_minutes: int | None = None,
paramstyle: str | None = None,
on_start_query_execution: Callable[[str], None] | None = None,
result_set_type_hints: dict[str | int, str] | None = None,
*,
options: ExecuteOptions | None = None,
**kwargs,
) -> AioS3FSCursor:
"""Execute a SQL query asynchronously via S3FileSystem CSV reader.
Args:
operation: SQL query string to execute.
parameters: Query parameters for parameterized queries.
work_group: Athena workgroup to use for this query.
s3_staging_dir: S3 location for query results.
cache_size: Number of queries to check for result caching.
cache_expiration_time: Cache expiration time in seconds.
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.
result_set_type_hints: Optional dictionary mapping column names to
Athena DDL type signatures for precise type conversion within
complex types.
options: Shared execution options as an
:class:`~pyathena.options.ExecuteOptions` instance. Individual
keyword arguments take precedence over ``options`` fields.
**kwargs: Additional execution parameters.
Returns:
Self reference for method chaining.
"""
self._reset_state()
options = ExecuteOptions.resolve(
options,
work_group=work_group,
s3_staging_dir=s3_staging_dir,
cache_size=cache_size,
cache_expiration_time=cache_expiration_time,
result_reuse_enable=result_reuse_enable,
result_reuse_minutes=result_reuse_minutes,
paramstyle=paramstyle,
on_start_query_execution=on_start_query_execution,
result_set_type_hints=result_set_type_hints,
)
self.query_id = await self._execute(
operation,
parameters=parameters,
options=options,
)
# Call user callbacks immediately after start_query_execution
self._call_on_start_query_execution(self.query_id, options)
query_execution = await self._poll(self.query_id)
if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED:
self.result_set = await asyncio.to_thread(
AthenaS3FSResultSet,
connection=self._connection,
converter=self._converter,
query_execution=query_execution,
arraysize=self.arraysize,
retry_config=self._retry_config,
csv_reader=self._csv_reader,
filesystem_class=AioS3FileSystem,
result_set_type_hints=options.result_set_type_hints,
**kwargs,
)
else:
raise OperationalError(query_execution.state_change_reason)
return self
[docs]
async def fetchone( # type: ignore[override]
self,
) -> tuple[Any | None, ...] | dict[Any, Any | None] | None:
"""Fetch the next row of the result set.
Wraps the synchronous fetch in ``asyncio.to_thread`` because
``AthenaS3FSResultSet`` reads rows lazily from S3.
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(AthenaS3FSResultSet, self.result_set)
return await asyncio.to_thread(result_set.fetchone)
[docs]
async def fetchmany( # type: ignore[override]
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`` because
``AthenaS3FSResultSet`` reads rows lazily from S3.
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(AthenaS3FSResultSet, self.result_set)
return await asyncio.to_thread(result_set.fetchmany, size)
[docs]
async def fetchall( # type: ignore[override]
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`` because
``AthenaS3FSResultSet`` reads rows lazily from S3.
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(AthenaS3FSResultSet, self.result_set)
return await asyncio.to_thread(result_set.fetchall)
async def __anext__(self):
row = await self.fetchone()
if row is None:
raise StopAsyncIteration
return row