PyAthena 3.28.0__tar.gz → 3.29.1__tar.gz

{pyathena-3.28.0 → pyathena-3.29.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: PyAthena
-Version: 3.28.0
+Version: 3.29.1
 Summary: Python DB API 2.0 (PEP 249) client for Amazon Athena
 Project-URL: homepage, https://github.com/pyathena-dev/PyAthena/
 Project-URL: repository, https://github.com/pyathena-dev/PyAthena/
@@ -43,6 +43,8 @@ Requires-Dist: botocore>=1.29.4
 Requires-Dist: fsspec
 Requires-Dist: python-dateutil
 Requires-Dist: tenacity>=4.1.0
+Provides-Extra: aiosqlalchemy
+Requires-Dist: sqlalchemy[asyncio]>=2.0.0; extra == 'aiosqlalchemy'
 Provides-Extra: arrow
 Requires-Dist: pyarrow>=10.0.0; (python_version < '3.14') and extra == 'arrow'
 Requires-Dist: pyarrow>=22.0.0; (python_version >= '3.14') and extra == 'arrow'
@@ -109,6 +111,22 @@ print(cursor.description)
 print(cursor.fetchall())
 ```
 
+Native asyncio is also supported:
+
+```python
+import asyncio
+from pyathena import aconnect
+
+async def main():
+    async with await aconnect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
+                              region_name="us-west-2") as conn:
+        cursor = conn.cursor()
+        await cursor.execute("SELECT 1")
+        print(await cursor.fetchone())
+
+asyncio.run(main())
+```
+
 ## License
 
 [MIT license](LICENSE)

{pyathena-3.28.0 → pyathena-3.29.1}/README.md

@@ -52,6 +52,22 @@ print(cursor.description)
 print(cursor.fetchall())
 ```
 
+Native asyncio is also supported:
+
+```python
+import asyncio
+from pyathena import aconnect
+
+async def main():
+    async with await aconnect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
+                              region_name="us-west-2") as conn:
+        cursor = conn.cursor()
+        await cursor.execute("SELECT 1")
+        print(await cursor.fetchone())
+
+asyncio.run(main())
+```
+
 ## License
 
 [MIT license](LICENSE)

{pyathena-3.28.0 → pyathena-3.29.1}/pyathena/__init__.py

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, FrozenSet, Type, overload
 from pyathena.error import *  # noqa
 
 if TYPE_CHECKING:
+    from pyathena.aio.connection import AioConnection
     from pyathena.connection import Connection, ConnectionCursor
     from pyathena.cursor import Cursor
 
@@ -128,3 +129,32 @@ def connect(*args, **kwargs) -> "Connection[Any]":
     from pyathena.connection import Connection
 
     return Connection(*args, **kwargs)
+
+
+async def aconnect(*args, **kwargs) -> "AioConnection":
+    """Create a new async database connection to Amazon Athena.
+
+    This is the async counterpart of :func:`connect`. It returns an
+    ``AioConnection`` whose cursors use native ``asyncio`` for polling
+    and API calls, keeping the event loop free.
+
+    Args:
+        **kwargs: Arguments forwarded to ``AioConnection.create()``.
+            See :func:`connect` for the full list of supported arguments.
+
+    Returns:
+        An ``AioConnection`` that produces ``AioCursor`` instances by default.
+
+    Example:
+        >>> import pyathena
+        >>> conn = await pyathena.aconnect(
+        ...     s3_staging_dir='s3://my-bucket/staging/',
+        ...     region_name='us-east-1',
+        ... )
+        >>> async with conn.cursor() as cursor:
+        ...     await cursor.execute("SELECT 1")
+        ...     print(await cursor.fetchone())
+    """
+    from pyathena.aio.connection import AioConnection
+
+    return await AioConnection.create(*args, **kwargs)

{pyathena-3.28.0 → pyathena-3.29.1}/pyathena/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '3.28.0'
-__version_tuple__ = version_tuple = (3, 28, 0)
+__version__ = version = '3.29.1'
+__version_tuple__ = version_tuple = (3, 29, 1)
 
 __commit_id__ = commit_id = None

pyathena-3.29.1/pyathena/aio/arrow/cursor.py

@@ -0,0 +1,230 @@
+# -*- coding: utf-8 -*-
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union, cast
+
+from pyathena.aio.common import WithAsyncFetch
+from pyathena.arrow.converter import (
+    DefaultArrowTypeConverter,
+    DefaultArrowUnloadTypeConverter,
+)
+from pyathena.arrow.result_set import AthenaArrowResultSet
+from pyathena.common import CursorIterator
+from pyathena.error import OperationalError, ProgrammingError
+from pyathena.model import AthenaQueryExecution
+
+if TYPE_CHECKING:
+    import polars as pl
+    from pyarrow import Table
+
+_logger = logging.getLogger(__name__)  # type: ignore
+
+
+class AioArrowCursor(WithAsyncFetch):
+    """Native asyncio cursor that returns results as Apache Arrow Tables.
+
+    Uses ``asyncio.to_thread()`` for both result set creation and fetch
+    operations, keeping the event loop free.
+
+    Example:
+        >>> async with await pyathena.aconnect(...) as conn:
+        ...     cursor = conn.cursor(AioArrowCursor)
+        ...     await cursor.execute("SELECT * FROM my_table")
+        ...     table = cursor.as_arrow()
+    """
+
+    def __init__(
+        self,
+        s3_staging_dir: Optional[str] = None,
+        schema_name: Optional[str] = None,
+        catalog_name: Optional[str] = None,
+        work_group: Optional[str] = None,
+        poll_interval: float = 1,
+        encryption_option: Optional[str] = None,
+        kms_key: Optional[str] = None,
+        kill_on_interrupt: bool = True,
+        unload: bool = False,
+        result_reuse_enable: bool = False,
+        result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES,
+        connect_timeout: Optional[float] = None,
+        request_timeout: Optional[float] = 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._unload = unload
+        self._connect_timeout = connect_timeout
+        self._request_timeout = request_timeout
+        self._result_set: Optional[AthenaArrowResultSet] = None
+
+    @staticmethod
+    def get_default_converter(
+        unload: bool = False,
+    ) -> Union[DefaultArrowTypeConverter, DefaultArrowUnloadTypeConverter, Any]:
+        if unload:
+            return DefaultArrowUnloadTypeConverter()
+        return DefaultArrowTypeConverter()
+
+    async def execute(  # type: ignore[override]
+        self,
+        operation: str,
+        parameters: Optional[Union[Dict[str, Any], List[str]]] = None,
+        work_group: Optional[str] = None,
+        s3_staging_dir: Optional[str] = None,
+        cache_size: Optional[int] = 0,
+        cache_expiration_time: Optional[int] = 0,
+        result_reuse_enable: Optional[bool] = None,
+        result_reuse_minutes: Optional[int] = None,
+        paramstyle: Optional[str] = None,
+        **kwargs,
+    ) -> "AioArrowCursor":
+        """Execute a SQL query asynchronously and return results as Arrow Tables.
+
+        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').
+            **kwargs: Additional execution parameters.
+
+        Returns:
+            Self reference for method chaining.
+        """
+        self._reset_state()
+        operation, unload_location = self._prepare_unload(operation, s3_staging_dir)
+        self.query_id = await self._execute(
+            operation,
+            parameters=parameters,
+            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,
+        )
+
+        query_execution = await self._poll(self.query_id)
+        if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED:
+            self.result_set = await asyncio.to_thread(
+                AthenaArrowResultSet,
+                connection=self._connection,
+                converter=self._converter,
+                query_execution=query_execution,
+                arraysize=self.arraysize,
+                retry_config=self._retry_config,
+                unload=self._unload,
+                unload_location=unload_location,
+                connect_timeout=self._connect_timeout,
+                request_timeout=self._request_timeout,
+                **kwargs,
+            )
+        else:
+            raise OperationalError(query_execution.state_change_reason)
+        return self
+
+    async def fetchone(  # type: ignore[override]
+        self,
+    ) -> Optional[Union[Tuple[Optional[Any], ...], Dict[Any, Optional[Any]]]]:
+        """Fetch the next row of the result set.
+
+        Wraps the synchronous fetch in ``asyncio.to_thread`` to avoid
+        blocking the event loop.
+
+        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 await asyncio.to_thread(result_set.fetchone)
+
+    async def fetchmany(  # type: ignore[override]
+        self, size: Optional[int] = None
+    ) -> List[Union[Tuple[Optional[Any], ...], Dict[Any, Optional[Any]]]]:
+        """Fetch multiple rows from the result set.
+
+        Wraps the synchronous fetch in ``asyncio.to_thread`` to avoid
+        blocking the event loop.
+
+        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 await asyncio.to_thread(result_set.fetchmany, size)
+
+    async def fetchall(  # type: ignore[override]
+        self,
+    ) -> List[Union[Tuple[Optional[Any], ...], Dict[Any, Optional[Any]]]]:
+        """Fetch all remaining rows from the result set.
+
+        Wraps the synchronous fetch in ``asyncio.to_thread`` to avoid
+        blocking the event loop.
+
+        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 await asyncio.to_thread(result_set.fetchall)
+
+    async def __anext__(self):
+        row = await self.fetchone()
+        if row is None:
+            raise StopAsyncIteration
+        return row
+
+    def as_arrow(self) -> "Table":
+        """Return query results as an Apache Arrow Table.
+
+        Returns:
+            Apache Arrow Table containing all query results.
+        """
+        if not self.has_result_set:
+            raise ProgrammingError("No result set.")
+        result_set = cast(AthenaArrowResultSet, self.result_set)
+        return result_set.as_arrow()
+
+    def as_polars(self) -> "pl.DataFrame":
+        """Return query results as a Polars DataFrame.
+
+        Returns:
+            Polars DataFrame containing all query results.
+        """
+        if not self.has_result_set:
+            raise ProgrammingError("No result set.")
+        result_set = cast(AthenaArrowResultSet, self.result_set)
+        return result_set.as_polars()