PyAthena 3.29.0__tar.gz → 3.29.2__tar.gz

{pyathena-3.29.0 → pyathena-3.29.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: PyAthena
-Version: 3.29.0
+Version: 3.29.2
 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'
@@ -113,10 +115,10 @@ Native asyncio is also supported:
 
 ```python
 import asyncio
-from pyathena import aconnect
+from pyathena import aio_connect
 
 async def main():
-    async with await aconnect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
+    async with await aio_connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
                               region_name="us-west-2") as conn:
         cursor = conn.cursor()
         await cursor.execute("SELECT 1")

{pyathena-3.29.0 → pyathena-3.29.2}/README.md

@@ -56,10 +56,10 @@ Native asyncio is also supported:
 
 ```python
 import asyncio
-from pyathena import aconnect
+from pyathena import aio_connect
 
 async def main():
-    async with await aconnect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
+    async with await aio_connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
                               region_name="us-west-2") as conn:
         cursor = conn.cursor()
         await cursor.execute("SELECT 1")

{pyathena-3.29.0 → pyathena-3.29.2}/pyathena/__init__.py

@@ -131,7 +131,7 @@ def connect(*args, **kwargs) -> "Connection[Any]":
     return Connection(*args, **kwargs)
 
 
-async def aconnect(*args, **kwargs) -> "AioConnection":
+async def aio_connect(*args, **kwargs) -> "AioConnection":
     """Create a new async database connection to Amazon Athena.
 
     This is the async counterpart of :func:`connect`. It returns an
@@ -147,7 +147,7 @@ async def aconnect(*args, **kwargs) -> "AioConnection":
 
     Example:
         >>> import pyathena
-        >>> conn = await pyathena.aconnect(
+        >>> conn = await pyathena.aio_connect(
         ...     s3_staging_dir='s3://my-bucket/staging/',
         ...     region_name='us-east-1',
         ... )

{pyathena-3.29.0 → pyathena-3.29.2}/pyathena/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '3.29.0'
-__version_tuple__ = version_tuple = (3, 29, 0)
+__version__ = version = '3.29.2'
+__version_tuple__ = version_tuple = (3, 29, 2)
 
 __commit_id__ = commit_id = None

{pyathena-3.29.0 → pyathena-3.29.2}/pyathena/aio/arrow/cursor.py

@@ -29,7 +29,7 @@ class AioArrowCursor(WithAsyncFetch):
     operations, keeping the event loop free.
 
     Example:
-        >>> async with await pyathena.aconnect(...) as conn:
+        >>> async with await pyathena.aio_connect(...) as conn:
         ...     cursor = conn.cursor(AioArrowCursor)
         ...     await cursor.execute("SELECT * FROM my_table")
         ...     table = cursor.as_arrow()

{pyathena-3.29.0 → pyathena-3.29.2}/pyathena/aio/pandas/cursor.py

@@ -40,7 +40,7 @@ class AioPandasCursor(WithAsyncFetch):
     when ``chunksize`` is set, as fetch calls trigger lazy S3 reads.
 
     Example:
-        >>> async with await pyathena.aconnect(...) as conn:
+        >>> async with await pyathena.aio_connect(...) as conn:
         ...     cursor = conn.cursor(AioPandasCursor)
         ...     await cursor.execute("SELECT * FROM my_table")
         ...     df = cursor.as_pandas()

{pyathena-3.29.0 → pyathena-3.29.2}/pyathena/aio/polars/cursor.py

@@ -31,7 +31,7 @@ class AioPolarsCursor(WithAsyncFetch):
     when ``chunksize`` is set, as fetch calls trigger lazy S3 reads.
 
     Example:
-        >>> async with await pyathena.aconnect(...) as conn:
+        >>> async with await pyathena.aio_connect(...) as conn:
         ...     cursor = conn.cursor(AioPolarsCursor)
         ...     await cursor.execute("SELECT * FROM my_table")
         ...     df = cursor.as_polars()

{pyathena-3.29.0 → pyathena-3.29.2}/pyathena/aio/s3fs/cursor.py

@@ -23,7 +23,7 @@ class AioS3FSCursor(WithAsyncFetch):
     reader, making fetch calls blocking I/O.
 
     Example:
-        >>> async with await pyathena.aconnect(...) as conn:
+        >>> async with await pyathena.aio_connect(...) as conn:
         ...     cursor = conn.cursor(AioS3FSCursor)
         ...     await cursor.execute("SELECT * FROM my_table")
         ...     row = await cursor.fetchone()

{pyathena-3.29.0 → pyathena-3.29.2}/pyathena/aio/spark/cursor.py

@@ -33,7 +33,7 @@ class AioSparkCursor(SparkBaseCursor, WithCalculationExecution):
 
     Example:
         >>> import asyncio
-        >>> async with await pyathena.aconnect(
+        >>> async with await pyathena.aio_connect(
         ...     work_group="spark-workgroup",
         ...     cursor_class=AioSparkCursor,
         ... ) as conn:

pyathena-3.29.2/pyathena/aio/sqlalchemy/arrow.py

@@ -0,0 +1,55 @@
+# -*- coding: utf-8 -*-
+from typing import TYPE_CHECKING
+
+from pyathena.aio.sqlalchemy.base import AthenaAioDialect
+from pyathena.util import strtobool
+
+if TYPE_CHECKING:
+    from types import ModuleType
+
+
+class AthenaAioArrowDialect(AthenaAioDialect):
+    """Async SQLAlchemy dialect for Amazon Athena with Apache Arrow result format.
+
+    This dialect uses ``AioArrowCursor`` for native asyncio query execution
+    with Apache Arrow Table results.
+
+    Connection URL Format:
+        ``awsathena+aioarrow://{access_key}:{secret_key}@athena.{region}.amazonaws.com/{schema}``
+
+    Query Parameters:
+        In addition to the base dialect parameters:
+        - unload: If "true", use UNLOAD for Parquet output
+
+    Example:
+        >>> from sqlalchemy.ext.asyncio import create_async_engine
+        >>> engine = create_async_engine(
+        ...     "awsathena+aioarrow://:@athena.us-west-2.amazonaws.com/default"
+        ...     "?s3_staging_dir=s3://my-bucket/athena-results/"
+        ...     "&unload=true"
+        ... )
+
+    See Also:
+        :class:`~pyathena.aio.arrow.cursor.AioArrowCursor`: The underlying async cursor.
+        :class:`~pyathena.aio.sqlalchemy.base.AthenaAioDialect`: Base async dialect.
+    """
+
+    driver = "aioarrow"
+    supports_statement_cache = True
+
+    def create_connect_args(self, url):
+        from pyathena.aio.arrow.cursor import AioArrowCursor
+
+        opts = super()._create_connect_args(url)
+        opts.update({"cursor_class": AioArrowCursor})
+        cursor_kwargs = {}
+        if "unload" in opts:
+            cursor_kwargs.update({"unload": bool(strtobool(opts.pop("unload")))})
+        if cursor_kwargs:
+            opts.update({"cursor_kwargs": cursor_kwargs})
+        self._connect_options = opts
+        return [[], opts]
+
+    @classmethod
+    def import_dbapi(cls) -> "ModuleType":
+        return super().import_dbapi()

pyathena-3.29.2/pyathena/aio/sqlalchemy/base.py

@@ -0,0 +1,221 @@
+# -*- coding: utf-8 -*-
+from __future__ import annotations
+
+from collections import deque
+from typing import TYPE_CHECKING, Any, Dict, List, MutableMapping, Optional, Tuple, Union, cast
+
+from sqlalchemy import pool
+from sqlalchemy.engine import AdaptedConnection
+from sqlalchemy.util.concurrency import await_only
+
+import pyathena
+from pyathena.aio.connection import AioConnection
+from pyathena.error import (
+    DatabaseError,
+    DataError,
+    Error,
+    IntegrityError,
+    InterfaceError,
+    InternalError,
+    NotSupportedError,
+    OperationalError,
+    ProgrammingError,
+)
+from pyathena.sqlalchemy.base import AthenaDialect
+
+if TYPE_CHECKING:
+    from types import ModuleType
+
+    from sqlalchemy import URL
+
+
+class AsyncAdapt_pyathena_cursor:  # noqa: N801 - follows SQLAlchemy's internal async adapter naming convention (e.g. AsyncAdapt_asyncpg_dbapi)
+    """Wraps any async PyAthena cursor with a sync DBAPI interface.
+
+    SQLAlchemy's async engine uses greenlet-based ``await_only()`` to call
+    async methods from synchronous code running inside the greenlet context.
+    This adapter wraps an ``AioCursor`` (or variant) so that the dialect can
+    use a normal synchronous DBAPI interface while the underlying I/O is async.
+    """
+
+    server_side = False
+    __slots__ = ("_cursor", "_rows")
+
+    def __init__(self, cursor: Any) -> None:
+        self._cursor = cursor
+        self._rows: deque[Any] = deque()
+
+    @property
+    def description(self) -> Any:
+        return self._cursor.description
+
+    @property
+    def rowcount(self) -> int:
+        return self._cursor.rowcount  # type: ignore[no-any-return]
+
+    def close(self) -> None:
+        self._cursor.close()
+        self._rows.clear()
+
+    def execute(self, operation: str, parameters: Any = None, **kwargs: Any) -> Any:
+        result = await_only(self._cursor.execute(operation, parameters, **kwargs))
+        if self._cursor.description:
+            self._rows = deque(await_only(self._cursor.fetchall()))
+        else:
+            self._rows.clear()
+        return result
+
+    def executemany(
+        self,
+        operation: str,
+        seq_of_parameters: List[Optional[Union[Dict[str, Any], List[str]]]],
+        **kwargs: Any,
+    ) -> None:
+        for parameters in seq_of_parameters:
+            await_only(self._cursor.execute(operation, parameters, **kwargs))
+        self._rows.clear()
+
+    def fetchone(self) -> Any:
+        if self._rows:
+            return self._rows.popleft()
+        return None
+
+    def fetchmany(self, size: Optional[int] = None) -> Any:
+        if size is None:
+            size = self._cursor.arraysize if hasattr(self._cursor, "arraysize") else 1
+        return [self._rows.popleft() for _ in range(min(size, len(self._rows)))]
+
+    def fetchall(self) -> Any:
+        items = list(self._rows)
+        self._rows.clear()
+        return items
+
+    def setinputsizes(self, sizes: Any) -> None:
+        self._cursor.setinputsizes(sizes)
+
+    async def _async_soft_close(self) -> None:
+        return
+
+    # PyAthena-specific methods used by AthenaDialect reflection
+    def list_databases(self, *args: Any, **kwargs: Any) -> Any:
+        return await_only(self._cursor.list_databases(*args, **kwargs))
+
+    def get_table_metadata(self, *args: Any, **kwargs: Any) -> Any:
+        return await_only(self._cursor.get_table_metadata(*args, **kwargs))
+
+    def list_table_metadata(self, *args: Any, **kwargs: Any) -> Any:
+        return await_only(self._cursor.list_table_metadata(*args, **kwargs))
+
+    def __enter__(self) -> "AsyncAdapt_pyathena_cursor":
+        return self
+
+    def __exit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
+        self.close()
+
+
+class AsyncAdapt_pyathena_connection(AdaptedConnection):  # noqa: N801 - follows SQLAlchemy's internal async adapter naming convention (e.g. AsyncAdapt_asyncpg_dbapi)
+    """Wraps ``AioConnection`` with a sync DBAPI interface.
+
+    This adapted connection delegates ``cursor()`` to the underlying
+    ``AioConnection`` and wraps each returned async cursor with
+    ``AsyncAdapt_pyathena_cursor``.
+    """
+
+    __slots__ = ("dbapi", "_connection")
+
+    def __init__(self, dbapi: "AsyncAdapt_pyathena_dbapi", connection: AioConnection) -> None:
+        self.dbapi = dbapi
+        self._connection = connection  # type: ignore[assignment]
+
+    @property
+    def driver_connection(self) -> AioConnection:
+        return self._connection  # type: ignore[return-value]
+
+    @property
+    def catalog_name(self) -> Optional[str]:
+        return self._connection.catalog_name  # type: ignore[no-any-return]
+
+    @property
+    def schema_name(self) -> Optional[str]:
+        return self._connection.schema_name  # type: ignore[no-any-return]
+
+    def cursor(self) -> AsyncAdapt_pyathena_cursor:
+        raw_cursor = self._connection.cursor()
+        return AsyncAdapt_pyathena_cursor(raw_cursor)
+
+    def close(self) -> None:
+        self._connection.close()
+
+    def commit(self) -> None:
+        self._connection.commit()  # type: ignore[unused-coroutine]
+
+    def rollback(self) -> None:
+        pass
+
+
+class AsyncAdapt_pyathena_dbapi:  # noqa: N801 - follows SQLAlchemy's internal async adapter naming convention (e.g. AsyncAdapt_asyncpg_dbapi)
+    """Fake DBAPI module for the async SQLAlchemy engine.
+
+    SQLAlchemy expects ``import_dbapi()`` to return a module-like object
+    with ``connect()``, ``paramstyle``, and the standard DBAPI exception
+    hierarchy.  This class fulfils that contract while routing connections
+    through ``AioConnection``.
+    """
+
+    paramstyle = "pyformat"
+
+    # DBAPI exception hierarchy
+    Error = Error
+    Warning = pyathena.Warning
+    InterfaceError = InterfaceError
+    DatabaseError = DatabaseError
+    InternalError = InternalError
+    OperationalError = OperationalError
+    ProgrammingError = ProgrammingError
+    IntegrityError = IntegrityError
+    DataError = DataError
+    NotSupportedError = NotSupportedError
+
+    def connect(self, **kwargs: Any) -> AsyncAdapt_pyathena_connection:
+        connection = await_only(AioConnection.create(**kwargs))
+        return AsyncAdapt_pyathena_connection(self, connection)
+
+
+class AthenaAioDialect(AthenaDialect):
+    """Base async SQLAlchemy dialect for Amazon Athena.
+
+    Extends the synchronous ``AthenaDialect`` with async capability
+    by setting ``is_async = True`` and providing an adapted DBAPI module
+    that wraps ``AioConnection`` and async cursors via greenlet-based
+    ``await_only()``.
+
+    Subclasses (e.g. ``AthenaAioRestDialect``, ``AthenaAioPandasDialect``)
+    register concrete ``awsathena+aio*`` drivers.
+
+    See Also:
+        :class:`~pyathena.sqlalchemy.base.AthenaDialect`: Synchronous base dialect.
+        :class:`~pyathena.aio.connection.AioConnection`: Native async connection.
+    """
+
+    is_async = True
+    supports_statement_cache = True
+
+    @classmethod
+    def get_pool_class(cls, url: "URL") -> type:
+        return pool.AsyncAdaptedQueuePool
+
+    @classmethod
+    def import_dbapi(cls) -> "ModuleType":
+        return AsyncAdapt_pyathena_dbapi()  # type: ignore[return-value]
+
+    @classmethod
+    def dbapi(cls) -> "ModuleType":  # type: ignore[override]
+        return AsyncAdapt_pyathena_dbapi()  # type: ignore[return-value]
+
+    def create_connect_args(self, url: "URL") -> Tuple[Tuple[str], MutableMapping[str, Any]]:
+        opts = self._create_connect_args(url)
+        self._connect_options = opts
+        return cast(Tuple[str], ()), opts
+
+    def get_driver_connection(self, connection: Any) -> Any:
+        return connection

pyathena-3.29.2/pyathena/aio/sqlalchemy/pandas.py

@@ -0,0 +1,61 @@
+# -*- coding: utf-8 -*-
+from typing import TYPE_CHECKING
+
+from pyathena.aio.sqlalchemy.base import AthenaAioDialect
+from pyathena.util import strtobool
+
+if TYPE_CHECKING:
+    from types import ModuleType
+
+
+class AthenaAioPandasDialect(AthenaAioDialect):
+    """Async SQLAlchemy dialect for Amazon Athena with pandas DataFrame result format.
+
+    This dialect uses ``AioPandasCursor`` for native asyncio query execution
+    with pandas DataFrame results.
+
+    Connection URL Format:
+        ``awsathena+aiopandas://{access_key}:{secret_key}@athena.{region}.amazonaws.com/{schema}``
+
+    Query Parameters:
+        In addition to the base dialect parameters:
+        - unload: If "true", use UNLOAD for Parquet output
+        - engine: CSV parsing engine ("c", "python", or "pyarrow")
+        - chunksize: Number of rows per chunk for memory-efficient processing
+
+    Example:
+        >>> from sqlalchemy.ext.asyncio import create_async_engine
+        >>> engine = create_async_engine(
+        ...     "awsathena+aiopandas://:@athena.us-west-2.amazonaws.com/default"
+        ...     "?s3_staging_dir=s3://my-bucket/athena-results/"
+        ...     "&unload=true&chunksize=10000"
+        ... )
+
+    See Also:
+        :class:`~pyathena.aio.pandas.cursor.AioPandasCursor`: The underlying async cursor.
+        :class:`~pyathena.aio.sqlalchemy.base.AthenaAioDialect`: Base async dialect.
+    """
+
+    driver = "aiopandas"
+    supports_statement_cache = True
+
+    def create_connect_args(self, url):
+        from pyathena.aio.pandas.cursor import AioPandasCursor
+
+        opts = super()._create_connect_args(url)
+        opts.update({"cursor_class": AioPandasCursor})
+        cursor_kwargs = {}
+        if "unload" in opts:
+            cursor_kwargs.update({"unload": bool(strtobool(opts.pop("unload")))})
+        if "engine" in opts:
+            cursor_kwargs.update({"engine": opts.pop("engine")})
+        if "chunksize" in opts:
+            cursor_kwargs.update({"chunksize": int(opts.pop("chunksize"))})  # type: ignore
+        if cursor_kwargs:
+            opts.update({"cursor_kwargs": cursor_kwargs})
+        self._connect_options = opts
+        return [[], opts]
+
+    @classmethod
+    def import_dbapi(cls) -> "ModuleType":
+        return super().import_dbapi()

pyathena-3.29.2/pyathena/aio/sqlalchemy/polars.py

@@ -0,0 +1,55 @@
+# -*- coding: utf-8 -*-
+from typing import TYPE_CHECKING
+
+from pyathena.aio.sqlalchemy.base import AthenaAioDialect
+from pyathena.util import strtobool
+
+if TYPE_CHECKING:
+    from types import ModuleType
+
+
+class AthenaAioPolarsDialect(AthenaAioDialect):
+    """Async SQLAlchemy dialect for Amazon Athena with Polars DataFrame result format.
+
+    This dialect uses ``AioPolarsCursor`` for native asyncio query execution
+    with Polars DataFrame results.
+
+    Connection URL Format:
+        ``awsathena+aiopolars://{access_key}:{secret_key}@athena.{region}.amazonaws.com/{schema}``
+
+    Query Parameters:
+        In addition to the base dialect parameters:
+        - unload: If "true", use UNLOAD for Parquet output
+
+    Example:
+        >>> from sqlalchemy.ext.asyncio import create_async_engine
+        >>> engine = create_async_engine(
+        ...     "awsathena+aiopolars://:@athena.us-west-2.amazonaws.com/default"
+        ...     "?s3_staging_dir=s3://my-bucket/athena-results/"
+        ...     "&unload=true"
+        ... )
+
+    See Also:
+        :class:`~pyathena.aio.polars.cursor.AioPolarsCursor`: The underlying async cursor.
+        :class:`~pyathena.aio.sqlalchemy.base.AthenaAioDialect`: Base async dialect.
+    """
+
+    driver = "aiopolars"
+    supports_statement_cache = True
+
+    def create_connect_args(self, url):
+        from pyathena.aio.polars.cursor import AioPolarsCursor
+
+        opts = super()._create_connect_args(url)
+        opts.update({"cursor_class": AioPolarsCursor})
+        cursor_kwargs = {}
+        if "unload" in opts:
+            cursor_kwargs.update({"unload": bool(strtobool(opts.pop("unload")))})
+        if cursor_kwargs:
+            opts.update({"cursor_kwargs": cursor_kwargs})
+        self._connect_options = opts
+        return [[], opts]
+
+    @classmethod
+    def import_dbapi(cls) -> "ModuleType":
+        return super().import_dbapi()

pyathena-3.29.2/pyathena/aio/sqlalchemy/rest.py

@@ -0,0 +1,37 @@
+# -*- coding: utf-8 -*-
+from typing import TYPE_CHECKING
+
+from pyathena.aio.sqlalchemy.base import AthenaAioDialect
+
+if TYPE_CHECKING:
+    from types import ModuleType
+
+
+class AthenaAioRestDialect(AthenaAioDialect):
+    """Async SQLAlchemy dialect for Amazon Athena using the standard REST API cursor.
+
+    This dialect uses ``AioCursor`` for native asyncio query execution.
+    Results are returned as Python tuples with type conversion handled by
+    the default converter.
+
+    Connection URL Format:
+        ``awsathena+aiorest://{access_key}:{secret_key}@athena.{region}.amazonaws.com/{schema}``
+
+    Example:
+        >>> from sqlalchemy.ext.asyncio import create_async_engine
+        >>> engine = create_async_engine(
+        ...     "awsathena+aiorest://:@athena.us-west-2.amazonaws.com/default"
+        ...     "?s3_staging_dir=s3://my-bucket/athena-results/"
+        ... )
+
+    See Also:
+        :class:`~pyathena.aio.cursor.AioCursor`: The underlying async cursor.
+        :class:`~pyathena.aio.sqlalchemy.base.AthenaAioDialect`: Base async dialect.
+    """
+
+    driver = "aiorest"
+    supports_statement_cache = True
+
+    @classmethod
+    def import_dbapi(cls) -> "ModuleType":
+        return super().import_dbapi()

pyathena-3.29.2/pyathena/aio/sqlalchemy/s3fs.py

@@ -0,0 +1,44 @@
+# -*- coding: utf-8 -*-
+from typing import TYPE_CHECKING
+
+from pyathena.aio.sqlalchemy.base import AthenaAioDialect
+
+if TYPE_CHECKING:
+    from types import ModuleType
+
+
+class AthenaAioS3FSDialect(AthenaAioDialect):
+    """Async SQLAlchemy dialect for PyAthena with S3FS cursor.
+
+    This dialect uses ``AioS3FSCursor`` for native asyncio query execution
+    with S3 filesystem-based CSV result reading.
+
+    Connection URL Format:
+        ``awsathena+aios3fs://{access_key}:{secret_key}@athena.{region}.amazonaws.com/{schema}``
+
+    Example:
+        >>> from sqlalchemy.ext.asyncio import create_async_engine
+        >>> engine = create_async_engine(
+        ...     "awsathena+aios3fs://:@athena.us-east-1.amazonaws.com/database"
+        ...     "?s3_staging_dir=s3://bucket/path"
+        ... )
+
+    See Also:
+        :class:`~pyathena.aio.s3fs.cursor.AioS3FSCursor`: The underlying async cursor.
+        :class:`~pyathena.aio.sqlalchemy.base.AthenaAioDialect`: Base async dialect.
+    """
+
+    driver = "aios3fs"
+    supports_statement_cache = True
+
+    def create_connect_args(self, url):
+        from pyathena.aio.s3fs.cursor import AioS3FSCursor
+
+        opts = super()._create_connect_args(url)
+        opts.update({"cursor_class": AioS3FSCursor})
+        self._connect_options = opts
+        return [[], opts]
+
+    @classmethod
+    def import_dbapi(cls) -> "ModuleType":
+        return super().import_dbapi()

pyathena-3.29.2/pyathena/sqlalchemy/__init__.py

@@ -0,0 +1 @@
+# -*- coding: utf-8 -*-

{pyathena-3.29.0 → pyathena-3.29.2}/pyathena/sqlalchemy/compiler.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Type, Union, cast
+from typing import TYPE_CHECKING, Any, Dict, List, Mapping, Optional, Tuple, Union, cast
 
 from sqlalchemy import exc, types, util
 from sqlalchemy.sql.compiler import (
@@ -31,12 +31,12 @@ if TYPE_CHECKING:
         UniqueConstraint,
     )
     from sqlalchemy.sql.ddl import CreateTable
-    from sqlalchemy.sql.elements import FunctionElement
+    from sqlalchemy.sql.functions import Function
     from sqlalchemy.sql.selectable import GenerativeSelect
 
     from pyathena.sqlalchemy.base import AthenaDialect
 
-    _DialectArgDict = Dict[str, Any]
+    _DialectArgDict = Mapping[str, Any]
     CreateColumn = Any
 
 
@@ -61,10 +61,10 @@ class AthenaTypeCompiler(GenericTypeCompiler):
         https://docs.aws.amazon.com/athena/latest/ug/data-types.html
     """
 
-    def visit_FLOAT(self, type_: Type[Any], **kw) -> str:  # noqa: N802
-        return self.visit_REAL(type_, **kw)
+    def visit_FLOAT(self, type_: types.Float[Any], **kw: Any) -> str:  # noqa: N802
+        return self.visit_REAL(type_, **kw)  # type: ignore[arg-type]
 
-    def visit_REAL(self, type_: Type[Any], **kw) -> str:  # noqa: N802
+    def visit_REAL(self, type_: types.REAL[Any], **kw: Any) -> str:  # noqa: N802
         return "FLOAT"
 
     def visit_DOUBLE(self, type_, **kw) -> str:  # noqa: N802
@@ -73,78 +73,78 @@ class AthenaTypeCompiler(GenericTypeCompiler):
     def visit_DOUBLE_PRECISION(self, type_, **kw) -> str:  # noqa: N802
         return "DOUBLE"
 
-    def visit_NUMERIC(self, type_: Type[Any], **kw) -> str:  # noqa: N802
-        return self.visit_DECIMAL(type_, **kw)
+    def visit_NUMERIC(self, type_: types.Numeric[Any], **kw: Any) -> str:  # noqa: N802
+        return self.visit_DECIMAL(type_, **kw)  # type: ignore[arg-type]
 
-    def visit_DECIMAL(self, type_: Type[Any], **kw) -> str:  # noqa: N802
+    def visit_DECIMAL(self, type_: types.DECIMAL[Any], **kw: Any) -> str:  # noqa: N802
         if type_.precision is None:
             return "DECIMAL"
         if type_.scale is None:
             return f"DECIMAL({type_.precision})"
         return f"DECIMAL({type_.precision}, {type_.scale})"
 
-    def visit_TINYINT(self, type_: Type[Any], **kw) -> str:  # noqa: N802
+    def visit_TINYINT(self, type_: types.Integer, **kw: Any) -> str:  # noqa: N802
         return "TINYINT"
 
-    def visit_INTEGER(self, type_: Type[Any], **kw) -> str:  # noqa: N802
+    def visit_INTEGER(self, type_: types.Integer, **kw: Any) -> str:  # noqa: N802
         return "INTEGER"
 
-    def visit_SMALLINT(self, type_: Type[Any], **kw) -> str:  # noqa: N802
+    def visit_SMALLINT(self, type_: types.SmallInteger, **kw: Any) -> str:  # noqa: N802
         return "SMALLINT"
 
-    def visit_BIGINT(self, type_: Type[Any], **kw) -> str:  # noqa: N802
+    def visit_BIGINT(self, type_: types.BigInteger, **kw: Any) -> str:  # noqa: N802
         return "BIGINT"
 
-    def visit_TIMESTAMP(self, type_: Type[Any], **kw) -> str:  # noqa: N802
+    def visit_TIMESTAMP(self, type_: types.TIMESTAMP, **kw: Any) -> str:  # noqa: N802
         return "TIMESTAMP"
 
-    def visit_DATETIME(self, type_: Type[Any], **kw) -> str:  # noqa: N802
-        return self.visit_TIMESTAMP(type_, **kw)
+    def visit_DATETIME(self, type_: types.DateTime, **kw: Any) -> str:  # noqa: N802
+        return self.visit_TIMESTAMP(type_, **kw)  # type: ignore[arg-type]
 
-    def visit_DATE(self, type_: Type[Any], **kw) -> str:  # noqa: N802
+    def visit_DATE(self, type_: types.Date, **kw: Any) -> str:  # noqa: N802
         return "DATE"
 
-    def visit_TIME(self, type_: Type[Any], **kw) -> str:  # noqa: N802
+    def visit_TIME(self, type_: types.Time, **kw: Any) -> str:  # noqa: N802
         raise exc.CompileError(f"Data type `{type_}` is not supported")
 
-    def visit_CLOB(self, type_: Type[Any], **kw) -> str:  # noqa: N802
-        return self.visit_BINARY(type_, **kw)
+    def visit_CLOB(self, type_: types.CLOB, **kw: Any) -> str:  # noqa: N802
+        return self.visit_BINARY(type_, **kw)  # type: ignore[arg-type]
 
-    def visit_NCLOB(self, type_: Type[Any], **kw) -> str:  # noqa: N802
-        return self.visit_BINARY(type_, **kw)
+    def visit_NCLOB(self, type_: types.Text, **kw: Any) -> str:  # noqa: N802
+        return self.visit_BINARY(type_, **kw)  # type: ignore[arg-type]
 
-    def visit_CHAR(self, type_: Type[Any], **kw) -> str:  # noqa: N802
+    def visit_CHAR(self, type_: types.CHAR, **kw: Any) -> str:  # noqa: N802
         if type_.length:
-            return cast(str, self._render_string_type(type_, "CHAR"))
+            return self._render_string_type("CHAR", type_.length, type_.collation)
         return "STRING"
 
-    def visit_NCHAR(self, type_: Type[Any], **kw) -> str:  # noqa: N802
-        return self.visit_CHAR(type_, **kw)
+    def visit_NCHAR(self, type_: types.NCHAR, **kw: Any) -> str:  # noqa: N802
+        return self.visit_CHAR(type_, **kw)  # type: ignore[arg-type]
 
-    def visit_VARCHAR(self, type_: Type[Any], **kw) -> str:  # noqa: N802
+    def visit_VARCHAR(self, type_: types.String, **kw: Any) -> str:  # noqa: N802
         if type_.length:
-            return cast(str, self._render_string_type(type_, "VARCHAR"))
+            return self._render_string_type("VARCHAR", type_.length, type_.collation)
         return "STRING"
 
-    def visit_NVARCHAR(self, type_: Type[Any], **kw) -> str:  # noqa: N802
-        return self.visit_VARCHAR(type_, **kw)
+    def visit_NVARCHAR(self, type_: types.NVARCHAR, **kw: Any) -> str:  # noqa: N802
+        return self.visit_VARCHAR(type_, **kw)  # type: ignore[arg-type]
 
-    def visit_TEXT(self, type_: Type[Any], **kw) -> str:  # noqa: N802
+    def visit_TEXT(self, type_: types.Text, **kw: Any) -> str:  # noqa: N802
         return "STRING"
 
-    def visit_BLOB(self, type_: Type[Any], **kw) -> str:  # noqa: N802
-        return self.visit_BINARY(type_, **kw)
+    def visit_BLOB(self, type_: types.LargeBinary, **kw: Any) -> str:  # noqa: N802
+        return self.visit_BINARY(type_, **kw)  # type: ignore[arg-type]
 
-    def visit_BINARY(self, type_: Type[Any], **kw) -> str:  # noqa: N802
+    def visit_BINARY(self, type_: types.BINARY, **kw: Any) -> str:  # noqa: N802
         return "BINARY"
 
-    def visit_VARBINARY(self, type_: Type[Any], **kw) -> str:  # noqa: N802
-        return self.visit_BINARY(type_, **kw)
+    def visit_VARBINARY(self, type_: types.VARBINARY, **kw: Any) -> str:  # noqa: N802
+        return self.visit_BINARY(type_, **kw)  # type: ignore[arg-type]
 
-    def visit_BOOLEAN(self, type_: Type[Any], **kw) -> str:  # noqa: N802
+    def visit_BOOLEAN(self, type_: types.Boolean, **kw: Any) -> str:  # noqa: N802
         return "BOOLEAN"
 
-    def visit_JSON(self, type_: Type[Any], **kw) -> str:  # noqa: N802
+    def visit_JSON(self, type_: types.JSON, **kw: Any) -> str:  # noqa: N802
         return "JSON"
 
     def visit_string(self, type_, **kw):  # noqa: N802
@@ -219,10 +219,10 @@ class AthenaStatementCompiler(SQLCompiler):
         https://docs.aws.amazon.com/athena/latest/ug/ddl-sql-reference.html
     """
 
-    def visit_char_length_func(self, fn: "FunctionElement[Any]", **kw):
+    def visit_char_length_func(self, fn: "Function[Any]", **kw: Any) -> str:
         return f"length{self.function_argspec(fn, **kw)}"
 
-    def visit_filter_func(self, fn: "FunctionElement[Any]", **kw) -> str:
+    def visit_filter_func(self, fn: "Function[Any]", **kw: Any) -> str:
         """Compile Athena filter() function with lambda expressions.
 
         Supports syntax: filter(array_expr, lambda_expr)
@@ -370,7 +370,7 @@ class AthenaDDLCompiler(DDLCompiler):
         return f"COMMENT {self._escape_comment(comment)}"
 
     def _get_bucket_count(
-        self, dialect_opts: "_DialectArgDict", connect_opts: Dict[str, Any]
+        self, dialect_opts: "_DialectArgDict", connect_opts: Mapping[str, Any]
     ) -> Optional[str]:
         if dialect_opts["bucket_count"]:
             bucket_count = dialect_opts["bucket_count"]
@@ -381,7 +381,7 @@ class AthenaDDLCompiler(DDLCompiler):
         return cast(str, bucket_count) if bucket_count is not None else None
 
     def _get_file_format(
-        self, dialect_opts: "_DialectArgDict", connect_opts: Dict[str, Any]
+        self, dialect_opts: "_DialectArgDict", connect_opts: Mapping[str, Any]
     ) -> Optional[str]:
         if dialect_opts["file_format"]:
             file_format = dialect_opts["file_format"]
@@ -392,7 +392,7 @@ class AthenaDDLCompiler(DDLCompiler):
         return cast(Optional[str], file_format)
 
     def _get_file_format_specification(
-        self, dialect_opts: "_DialectArgDict", connect_opts: Dict[str, Any]
+        self, dialect_opts: "_DialectArgDict", connect_opts: Mapping[str, Any]
     ) -> str:
         file_format = self._get_file_format(dialect_opts, connect_opts)
         text = []
@@ -401,7 +401,7 @@ class AthenaDDLCompiler(DDLCompiler):
         return "\n".join(text)
 
     def _get_row_format(
-        self, dialect_opts: "_DialectArgDict", connect_opts: Dict[str, Any]
+        self, dialect_opts: "_DialectArgDict", connect_opts: Mapping[str, Any]
     ) -> Optional[str]:
         if dialect_opts["row_format"]:
             row_format = dialect_opts["row_format"]
@@ -412,7 +412,7 @@ class AthenaDDLCompiler(DDLCompiler):
         return cast(Optional[str], row_format)
 
     def _get_row_format_specification(
-        self, dialect_opts: "_DialectArgDict", connect_opts: Dict[str, Any]
+        self, dialect_opts: "_DialectArgDict", connect_opts: Mapping[str, Any]
     ) -> str:
         row_format = self._get_row_format(dialect_opts, connect_opts)
         text = []
@@ -421,7 +421,7 @@ class AthenaDDLCompiler(DDLCompiler):
         return "\n".join(text)
 
     def _get_serde_properties(
-        self, dialect_opts: "_DialectArgDict", connect_opts: Dict[str, Any]
+        self, dialect_opts: "_DialectArgDict", connect_opts: Mapping[str, Any]
     ) -> Optional[Union[str, Dict[str, Any]]]:
         if dialect_opts["serdeproperties"]:
             serde_properties = dialect_opts["serdeproperties"]
@@ -432,7 +432,7 @@ class AthenaDDLCompiler(DDLCompiler):
         return cast(Optional[str], serde_properties)
 
     def _get_serde_properties_specification(
-        self, dialect_opts: "_DialectArgDict", connect_opts: Dict[str, Any]
+        self, dialect_opts: "_DialectArgDict", connect_opts: Mapping[str, Any]
     ) -> str:
         serde_properties = self._get_serde_properties(dialect_opts, connect_opts)
         text = []
@@ -446,7 +446,7 @@ class AthenaDDLCompiler(DDLCompiler):
         return "\n".join(text)
 
     def _get_table_location(
-        self, table: "Table", dialect_opts: "_DialectArgDict", connect_opts: Dict[str, Any]
+        self, table: "Table", dialect_opts: "_DialectArgDict", connect_opts: Mapping[str, Any]
     ) -> Optional[str]:
         if dialect_opts["location"]:
             location = cast(str, dialect_opts["location"])
@@ -464,7 +464,7 @@ class AthenaDDLCompiler(DDLCompiler):
         return location
 
     def _get_table_location_specification(
-        self, table: "Table", dialect_opts: "_DialectArgDict", connect_opts: Dict[str, Any]
+        self, table: "Table", dialect_opts: "_DialectArgDict", connect_opts: Mapping[str, Any]
     ) -> str:
         location = self._get_table_location(table, dialect_opts, connect_opts)
         text = []
@@ -482,7 +482,7 @@ class AthenaDDLCompiler(DDLCompiler):
         return "\n".join(text)
 
     def _get_table_properties(
-        self, dialect_opts: "_DialectArgDict", connect_opts: Dict[str, Any]
+        self, dialect_opts: "_DialectArgDict", connect_opts: Mapping[str, Any]
     ) -> Optional[Union[Dict[str, str], str]]:
         if dialect_opts["tblproperties"]:
             table_properties = cast(str, dialect_opts["tblproperties"])
@@ -493,7 +493,7 @@ class AthenaDDLCompiler(DDLCompiler):
         return table_properties
 
     def _get_compression(
-        self, dialect_opts: "_DialectArgDict", connect_opts: Dict[str, Any]
+        self, dialect_opts: "_DialectArgDict", connect_opts: Mapping[str, Any]
     ) -> Optional[str]:
         if dialect_opts["compression"]:
             compression = cast(str, dialect_opts["compression"])
@@ -504,7 +504,7 @@ class AthenaDDLCompiler(DDLCompiler):
         return compression
 
     def _get_table_properties_specification(
-        self, dialect_opts: "_DialectArgDict", connect_opts: Dict[str, Any]
+        self, dialect_opts: "_DialectArgDict", connect_opts: Mapping[str, Any]
     ) -> str:
         properties = self._get_table_properties(dialect_opts, connect_opts)
         if properties:
@@ -554,26 +554,22 @@ class AthenaDDLCompiler(DDLCompiler):
             text.append(f"{self._get_comment_specification(column.comment)}")
         return " ".join(text)
 
-    def visit_check_constraint(self, constraint: "CheckConstraint", **kw) -> Optional[str]:
-        return None
+    def visit_check_constraint(self, constraint: "CheckConstraint", **kw: Any) -> str:
+        return ""
 
-    def visit_column_check_constraint(self, constraint: "CheckConstraint", **kw) -> Optional[str]:
-        return None
+    def visit_column_check_constraint(self, constraint: "CheckConstraint", **kw: Any) -> str:
+        return ""
 
-    def visit_foreign_key_constraint(
-        self, constraint: "ForeignKeyConstraint", **kw
-    ) -> Optional[str]:
-        return None
+    def visit_foreign_key_constraint(self, constraint: "ForeignKeyConstraint", **kw: Any) -> str:
+        return ""
 
-    def visit_primary_key_constraint(
-        self, constraint: "PrimaryKeyConstraint", **kw
-    ) -> Optional[str]:
-        return None
+    def visit_primary_key_constraint(self, constraint: "PrimaryKeyConstraint", **kw: Any) -> str:
+        return ""
 
-    def visit_unique_constraint(self, constraint: "UniqueConstraint", **kw) -> Optional[str]:
-        return None
+    def visit_unique_constraint(self, constraint: "UniqueConstraint", **kw: Any) -> str:
+        return ""
 
-    def _get_connect_option_partitions(self, connect_opts: Dict[str, Any]) -> List[str]:
+    def _get_connect_option_partitions(self, connect_opts: Mapping[str, Any]) -> List[str]:
         if connect_opts:
             partition = cast(str, connect_opts.get("partition"))
             partitions = partition.split(",") if partition else []
@@ -581,7 +577,7 @@ class AthenaDDLCompiler(DDLCompiler):
             partitions = []
         return partitions
 
-    def _get_connect_option_buckets(self, connect_opts: Dict[str, Any]) -> List[str]:
+    def _get_connect_option_buckets(self, connect_opts: Mapping[str, Any]) -> List[str]:
         if connect_opts:
             bucket = cast(str, connect_opts.get("cluster"))
             buckets = bucket.split(",") if bucket else []
@@ -624,7 +620,7 @@ class AthenaDDLCompiler(DDLCompiler):
         table: "Table",
         is_iceberg: bool,
         create_columns: List["CreateColumn"],
-        connect_opts: Dict[str, Any],
+        connect_opts: Mapping[str, Any],
     ) -> Tuple[List[str], List[str], List[str]]:
         columns, partitions, buckets = [], [], []
         conn_partitions = self._get_connect_option_partitions(connect_opts)

{pyathena-3.29.0 → pyathena-3.29.2}/pyproject.toml

@@ -42,9 +42,17 @@ awsathena = "pyathena.sqlalchemy.base:AthenaDialect"
 "awsathena.arrow" = "pyathena.sqlalchemy.arrow:AthenaArrowDialect"
 "awsathena.polars" = "pyathena.sqlalchemy.polars:AthenaPolarsDialect"
 "awsathena.s3fs" = "pyathena.sqlalchemy.s3fs:AthenaS3FSDialect"
+"awsathena.aiorest" = "pyathena.aio.sqlalchemy.rest:AthenaAioRestDialect"
+"awsathena.aiopandas" = "pyathena.aio.sqlalchemy.pandas:AthenaAioPandasDialect"
+"awsathena.aioarrow" = "pyathena.aio.sqlalchemy.arrow:AthenaAioArrowDialect"
+"awsathena.aiopolars" = "pyathena.aio.sqlalchemy.polars:AthenaAioPolarsDialect"
+"awsathena.aios3fs" = "pyathena.aio.sqlalchemy.s3fs:AthenaAioS3FSDialect"
 
 [project.optional-dependencies]
 sqlalchemy = ["sqlalchemy>=1.0.0"]
+aiosqlalchemy = [
+    "sqlalchemy[asyncio]>=2.0.0",
+]
 pandas = [
     "pandas>=1.3.0; python_version<'3.13'",
     "pandas>=2.3.0; python_version>='3.13'",
@@ -59,7 +67,7 @@ polars = [
 
 [dependency-groups]
 dev = [
-    "sqlalchemy>=1.0.0",
+    "sqlalchemy[asyncio]>=1.0.0",
     "pandas>=1.3.0; python_version<'3.13'",
     "pandas>=2.3.0; python_version>='3.13'",
     "numpy>=1.26.0; python_version<'3.13'",
@@ -170,7 +178,7 @@ exclude = [
 legacy_tox_ini = """
 [tox]
 isolated_build = true
-envlist = py{310,311,312,313,314}
+envlist = py{310,311,312,313,314}-{pyathena,sqla,sqla_async}
 
 [gh-actions]
 python =
@@ -180,6 +188,12 @@ python =
     3.13: py313
     3.14: py314
 
+[gh-actions:env]
+TEST_TYPE =
+    pyathena: pyathena
+    sqla: sqla
+    sqla_async: sqla_async
+
 [testenv]
 allowlist_externals =
     uv
@@ -187,8 +201,9 @@ allowlist_externals =
     make
 commands =
     uv sync --group dev
-    make test
-    make test-sqla
+    pyathena: make test
+    sqla: make test-sqla
+    sqla_async: make test-sqla-async
 passenv =
     TOXENV
     AWS_*