PyAthena 3.31.1__tar.gz → 3.33.0__tar.gz

{pyathena-3.31.1 → pyathena-3.33.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: PyAthena
-Version: 3.31.1
+Version: 3.33.0
 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/
@@ -38,7 +38,7 @@ Classifier: Programming Language :: Python :: 3.13
 Classifier: Programming Language :: Python :: 3.14
 Classifier: Topic :: Database :: Front-Ends
 Requires-Python: >=3.10
-Requires-Dist: boto3>=1.38.2
+Requires-Dist: boto3>=1.41.2
 Requires-Dist: botocore>=1.41.2
 Requires-Dist: fsspec
 Requires-Dist: python-dateutil

{pyathena-3.31.1 → pyathena-3.33.0}/pyathena/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '3.31.1'
-__version_tuple__ = version_tuple = (3, 31, 1)
+__version__ = version = '3.33.0'
+__version_tuple__ = version_tuple = (3, 33, 0)
 
 __commit_id__ = commit_id = None

{pyathena-3.31.1 → pyathena-3.33.0}/pyathena/aio/common.py

@@ -83,6 +83,8 @@ class AioBaseCursor(BaseCursor):
     async def __poll(self, query_id: str) -> AthenaQueryExecution:
         while True:
             query_execution = await self._get_query_execution(query_id)
+            if self._on_poll:
+                self._on_poll(query_execution)
             if query_execution.state in [
                 AthenaQueryExecution.STATE_SUCCEEDED,
                 AthenaQueryExecution.STATE_FAILED,

{pyathena-3.31.1 → pyathena-3.33.0}/pyathena/aio/spark/cursor.py

@@ -128,6 +128,8 @@ class AioSparkCursor(SparkBaseCursor, WithCalculationExecution):
     async def __poll(self, query_id: str) -> AthenaQueryExecution | AthenaCalculationExecution:
         while True:
             calculation_status = await self._get_calculation_execution_status(query_id)
+            if self._on_poll:
+                self._on_poll(calculation_status)
             if calculation_status.state in [
                 AthenaCalculationExecutionStatus.STATE_COMPLETED,
                 AthenaCalculationExecutionStatus.STATE_FAILED,

{pyathena-3.31.1 → pyathena-3.33.0}/pyathena/arrow/cursor.py

@@ -63,7 +63,6 @@ class ArrowCursor(WithFetch):
         unload: bool = False,
         result_reuse_enable: bool = False,
         result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES,
-        on_start_query_execution: Callable[[str], None] | None = None,
         connect_timeout: float | None = None,
         request_timeout: float | None = None,
         **kwargs,
@@ -82,7 +81,6 @@ class ArrowCursor(WithFetch):
             unload: Enable UNLOAD for high-performance Parquet output.
             result_reuse_enable: Enable Athena query result reuse.
             result_reuse_minutes: Minutes to reuse cached results.
-            on_start_query_execution: Callback invoked when query starts.
             connect_timeout: Socket connection timeout in seconds for S3 operations.
                 Defaults to AWS SDK default (typically 1 second) if not specified.
             request_timeout: Request timeout in seconds for S3 operations.
@@ -113,7 +111,6 @@ class ArrowCursor(WithFetch):
             **kwargs,
         )
         self._unload = unload
-        self._on_start_query_execution = on_start_query_execution
         self._connect_timeout = connect_timeout
         self._request_timeout = request_timeout
 

{pyathena-3.31.1 → pyathena-3.33.0}/pyathena/common.py

@@ -4,6 +4,7 @@ import logging
 import sys
 import time
 from abc import ABCMeta, abstractmethod
+from collections.abc import Callable
 from datetime import datetime, timedelta, timezone
 from typing import TYPE_CHECKING, Any, cast
 
@@ -27,6 +28,14 @@ if TYPE_CHECKING:
 
 _logger = logging.getLogger(__name__)
 
+OnPollCallback = Callable[[AthenaQueryExecution | AthenaCalculationExecutionStatus], None]
+"""Type of the optional ``on_poll`` callback.
+
+Invoked once per poll iteration with the current execution object: an
+:class:`~pyathena.model.AthenaQueryExecution` for SQL queries, or an
+:class:`~pyathena.model.AthenaCalculationExecutionStatus` for Spark calculations.
+"""
+
 
 class CursorIterator(metaclass=ABCMeta):
     """Abstract base class providing iteration and result fetching capabilities for cursors.
@@ -164,6 +173,8 @@ class BaseCursor(metaclass=ABCMeta):
         kill_on_interrupt: bool,
         result_reuse_enable: bool,
         result_reuse_minutes: int,
+        on_start_query_execution: Callable[[str], None] | None = None,
+        on_poll: OnPollCallback | None = None,
         **kwargs,
     ) -> None:
         super().__init__()
@@ -181,6 +192,11 @@ class BaseCursor(metaclass=ABCMeta):
         self._kill_on_interrupt = kill_on_interrupt
         self._result_reuse_enable = result_reuse_enable
         self._result_reuse_minutes = result_reuse_minutes
+        # ``on_start_query_execution`` is invoked only by cursors whose ``execute()``
+        # supports it (the synchronous cursors). Async/aio/Spark cursors return the
+        # query id immediately through their execution model and do not invoke it.
+        self._on_start_query_execution = on_start_query_execution
+        self._on_poll = on_poll
 
     @staticmethod
     def get_default_converter(unload: bool = False) -> DefaultTypeConverter | Any:
@@ -558,6 +574,8 @@ class BaseCursor(metaclass=ABCMeta):
     def __poll(self, query_id: str) -> AthenaQueryExecution | AthenaCalculationExecution:
         while True:
             query_execution = self._get_query_execution(query_id)
+            if self._on_poll:
+                self._on_poll(query_execution)
             if query_execution.state in [
                 AthenaQueryExecution.STATE_SUCCEEDED,
                 AthenaQueryExecution.STATE_FAILED,

{pyathena-3.31.1 → pyathena-3.33.0}/pyathena/connection.py

@@ -18,7 +18,7 @@ from boto3.session import Session
 from botocore.config import Config
 
 import pyathena
-from pyathena.common import BaseCursor, CursorIterator
+from pyathena.common import BaseCursor, CursorIterator, OnPollCallback
 from pyathena.converter import Converter
 from pyathena.cursor import Cursor
 from pyathena.error import NotSupportedError, ProgrammingError
@@ -127,6 +127,7 @@ class Connection(Generic[ConnectionCursor]):
         result_reuse_enable: bool = ...,
         result_reuse_minutes: int = ...,
         on_start_query_execution: Callable[[str], None] | None = ...,
+        on_poll: OnPollCallback | None = ...,
         **kwargs,
     ) -> None: ...
 
@@ -158,6 +159,7 @@ class Connection(Generic[ConnectionCursor]):
         result_reuse_enable: bool = ...,
         result_reuse_minutes: int = ...,
         on_start_query_execution: Callable[[str], None] | None = ...,
+        on_poll: OnPollCallback | None = ...,
         **kwargs,
     ) -> None: ...
 
@@ -188,6 +190,7 @@ class Connection(Generic[ConnectionCursor]):
         result_reuse_enable: bool = False,
         result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES,
         on_start_query_execution: Callable[[str], None] | None = None,
+        on_poll: OnPollCallback | None = None,
         **kwargs,
     ) -> None:
         """Initialize a new Athena database connection.
@@ -224,6 +227,10 @@ class Connection(Generic[ConnectionCursor]):
             result_reuse_enable: Enable Athena query result reuse. Defaults to False.
             result_reuse_minutes: Minutes to reuse cached results.
             on_start_query_execution: Callback function called when query starts.
+            on_poll: Callback invoked once per poll iteration with the current
+                execution object (``AthenaQueryExecution``, or
+                ``AthenaCalculationExecutionStatus`` for Spark). Useful for
+                monitoring live query progress. Defaults to None.
             **kwargs: Additional arguments passed to boto3 Session and client.
 
         Raises:
@@ -331,6 +338,7 @@ class Connection(Generic[ConnectionCursor]):
         self.result_reuse_enable = result_reuse_enable
         self.result_reuse_minutes = result_reuse_minutes
         self.on_start_query_execution = on_start_query_execution
+        self.on_poll = on_poll
 
     def _assume_role(
         self,
@@ -555,6 +563,7 @@ class Connection(Generic[ConnectionCursor]):
             on_start_query_execution=kwargs.pop(
                 "on_start_query_execution", self.on_start_query_execution
             ),
+            on_poll=kwargs.pop("on_poll", self.on_poll),
             **kwargs,
         )
 

{pyathena-3.31.1 → pyathena-3.33.0}/pyathena/cursor.py

@@ -52,7 +52,6 @@ class Cursor(WithFetch):
         kill_on_interrupt: bool = True,
         result_reuse_enable: bool = False,
         result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES,
-        on_start_query_execution: Callable[[str], None] | None = None,
         **kwargs,
     ) -> None:
         super().__init__(
@@ -69,7 +68,6 @@ class Cursor(WithFetch):
             **kwargs,
         )
         self._result_set_class = AthenaResultSet
-        self._on_start_query_execution = on_start_query_execution
 
     @property
     def arraysize(self) -> int:

{pyathena-3.31.1 → pyathena-3.33.0}/pyathena/pandas/cursor.py

@@ -79,7 +79,6 @@ class PandasCursor(WithFetch):
         result_reuse_enable: bool = False,
         result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES,
         auto_optimize_chunksize: bool = False,
-        on_start_query_execution: Callable[[str], None] | None = None,
         **kwargs,
     ) -> None:
         """Initialize PandasCursor with configuration options.
@@ -105,7 +104,6 @@ class PandasCursor(WithFetch):
             auto_optimize_chunksize: Enable automatic chunksize determination for
                                    large files. Only effective when chunksize is None.
                                    Default: False (no automatic chunking).
-            on_start_query_execution: Callback for query start events.
             **kwargs: Additional arguments passed to pandas.read_csv.
         """
         super().__init__(
@@ -128,7 +126,6 @@ class PandasCursor(WithFetch):
         self._cache_type = cache_type
         self._max_workers = max_workers
         self._auto_optimize_chunksize = auto_optimize_chunksize
-        self._on_start_query_execution = on_start_query_execution
 
     @staticmethod
     def get_default_converter(

{pyathena-3.31.1 → pyathena-3.33.0}/pyathena/polars/cursor.py

@@ -74,7 +74,6 @@ class PolarsCursor(WithFetch):
         unload: bool = False,
         result_reuse_enable: bool = False,
         result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES,
-        on_start_query_execution: Callable[[str], None] | None = None,
         block_size: int | None = None,
         cache_type: str | None = None,
         max_workers: int = (cpu_count() or 1) * 5,
@@ -95,7 +94,6 @@ class PolarsCursor(WithFetch):
             unload: Enable UNLOAD for high-performance Parquet output.
             result_reuse_enable: Enable Athena query result reuse.
             result_reuse_minutes: Minutes to reuse cached results.
-            on_start_query_execution: Callback invoked when query starts.
             block_size: S3 read block size.
             cache_type: S3 caching strategy.
             max_workers: Maximum worker threads for parallel S3 operations.
@@ -123,7 +121,6 @@ class PolarsCursor(WithFetch):
             **kwargs,
         )
         self._unload = unload
-        self._on_start_query_execution = on_start_query_execution
         self._block_size = block_size
         self._cache_type = cache_type
         self._max_workers = max_workers

{pyathena-3.31.1 → pyathena-3.33.0}/pyathena/s3fs/cursor.py

@@ -58,7 +58,6 @@ class S3FSCursor(WithFetch):
         kill_on_interrupt: bool = True,
         result_reuse_enable: bool = False,
         result_reuse_minutes: int = CursorIterator.DEFAULT_RESULT_REUSE_MINUTES,
-        on_start_query_execution: Callable[[str], None] | None = None,
         csv_reader: CSVReaderType | None = None,
         **kwargs,
     ) -> None:
@@ -75,7 +74,6 @@ class S3FSCursor(WithFetch):
             kill_on_interrupt: Cancel running query on keyboard interrupt.
             result_reuse_enable: Enable Athena query result reuse.
             result_reuse_minutes: Minutes to reuse cached results.
-            on_start_query_execution: Callback invoked when query starts.
             csv_reader: CSV reader class to use for parsing results.
                 Use AthenaCSVReader (default) to distinguish between NULL
                 (unquoted empty) and empty string (quoted empty "").
@@ -104,7 +102,6 @@ class S3FSCursor(WithFetch):
             result_reuse_minutes=result_reuse_minutes,
             **kwargs,
         )
-        self._on_start_query_execution = on_start_query_execution
         self._csv_reader = csv_reader
 
     @staticmethod

{pyathena-3.31.1 → pyathena-3.33.0}/pyathena/spark/common.py

@@ -201,6 +201,8 @@ class SparkBaseCursor(BaseCursor, metaclass=ABCMeta):
     def __poll(self, query_id: str) -> AthenaQueryExecution | AthenaCalculationExecution:
         while True:
             calculation_status = self._get_calculation_execution_status(query_id)
+            if self._on_poll:
+                self._on_poll(calculation_status)
             if calculation_status.state in [
                 AthenaCalculationExecutionStatus.STATE_COMPLETED,
                 AthenaCalculationExecutionStatus.STATE_FAILED,

{pyathena-3.31.1 → pyathena-3.33.0}/pyathena/sqlalchemy/base.py

@@ -207,8 +207,7 @@ class AthenaDialect(DefaultDialect):
         #   awsathena+rest://
         #   {aws_access_key_id}:{aws_secret_access_key}@athena.{region_name}.amazonaws.com:443/
         #   {schema_name}?s3_staging_dir={s3_staging_dir}&...
-        self._connect_options = self._create_connect_args(url)
-        return cast(tuple[str], ()), self._connect_options
+        return cast(tuple[str], ()), self._create_connect_args(url)
 
     def _create_connect_args(self, url: URL) -> dict[str, Any]:
         opts: dict[str, Any] = {
@@ -238,6 +237,12 @@ class AthenaDialect(DefaultDialect):
             opts.update({"result_reuse_enable": bool(strtobool(opts["result_reuse_enable"]))})
         if "result_reuse_minutes" in opts:
             opts.update({"result_reuse_minutes": int(opts["result_reuse_minutes"])})
+        # Store on the dialect so compilers can consult connection options
+        # (e.g. catalog_name for S3 Tables detection). Assigned here rather than
+        # in create_connect_args because subclass dialects call this method
+        # directly and mutate the returned dict afterwards; sharing the same
+        # object keeps _connect_options in sync with their updates.
+        self._connect_options = opts
         return opts
 
     @reflection.cache

{pyathena-3.31.1 → pyathena-3.33.0}/pyathena/sqlalchemy/compiler.py

@@ -39,6 +39,13 @@ if TYPE_CHECKING:
     _DialectArgDict = Mapping[str, Any]
     CreateColumn = Any
 
+# Prefix of the Athena data catalog name registered for an Amazon S3 Tables
+# table bucket (e.g. ``s3tablescatalog/my-bucket``). It is selected via the
+# connection ``catalog_name``. S3 Tables are Iceberg-backed and use managed
+# storage, so their CREATE TABLE statements must not include a LOCATION clause.
+# https://docs.aws.amazon.com/AmazonS3/latest/userguide/s3-tables-integrations-query-athena.html
+S3_TABLES_CATALOG_PREFIX = "s3tablescatalog/"
+
 
 class AthenaTypeCompiler(GenericTypeCompiler):
     """Type compiler for Amazon Athena SQL types.
@@ -305,6 +312,10 @@ class AthenaDDLCompiler(DDLCompiler):
 
     - External table creation (EXTERNAL keyword for Hive-style tables)
     - Iceberg table creation (managed tables with ACID support)
+    - Amazon S3 Tables (Iceberg-backed, managed storage): set the connection
+      ``catalog_name`` to ``s3tablescatalog/<table-bucket>`` and use the
+      namespace as the table ``schema``. The LOCATION clause is omitted since
+      storage is managed.
     - File formats: PARQUET, ORC, TEXTFILE, JSON, AVRO, etc.
     - Row formats with SerDe specifications
     - Compression settings for various file formats
@@ -445,6 +456,74 @@ class AthenaDDLCompiler(DDLCompiler):
             text.append(")")
         return "\n".join(text)
 
+    @staticmethod
+    def _is_s3_tables_catalog(connect_opts: Mapping[str, Any]) -> bool:
+        """Return whether the connection targets an Amazon S3 Tables catalog.
+
+        S3 Tables are queried by setting the connection ``catalog_name`` to
+        ``s3tablescatalog/<table-bucket>`` and using the namespace as the table
+        ``schema`` (a two-part ``namespace.table`` identifier). Athena rejects a
+        three-part ``catalog.namespace.table`` identifier in DDL, so the catalog
+        must be selected at the connection level. Such tables use managed
+        storage, so their CREATE TABLE statement must omit the LOCATION clause.
+
+        Args:
+            connect_opts: The dialect connection options.
+
+        Returns:
+            True if ``catalog_name`` names an S3 Tables catalog.
+        """
+        if not connect_opts:
+            return False
+        catalog = connect_opts.get("catalog_name") or ""
+        # Athena resolves catalog names case-insensitively.
+        return catalog.lower().startswith(S3_TABLES_CATALOG_PREFIX)
+
+    def _is_iceberg_table(
+        self, dialect_opts: _DialectArgDict, connect_opts: Mapping[str, Any]
+    ) -> bool:
+        """Return whether the table properties declare an Iceberg table.
+
+        Args:
+            dialect_opts: The table's ``awsathena_*`` dialect options.
+            connect_opts: The dialect connection options.
+
+        Returns:
+            True if the rendered TBLPROPERTIES set ``table_type`` to Iceberg.
+        """
+        table_properties = self._get_table_properties_specification(
+            dialect_opts, connect_opts
+        ).lower()
+        return ("table_type" in table_properties) and ("iceberg" in table_properties)
+
+    def _validate_s3_tables_create_table(
+        self, dialect_opts: _DialectArgDict, connect_opts: Mapping[str, Any]
+    ) -> None:
+        """Validate a CREATE TABLE compiled against an S3 Tables catalog.
+
+        S3 Tables support only Iceberg tables on managed storage, so the table
+        must declare ``table_type`` ICEBERG and must not specify a location.
+        Raising here surfaces a clear client-side error instead of emitting DDL
+        that Athena would reject.
+
+        Args:
+            dialect_opts: The table's ``awsathena_*`` dialect options.
+            connect_opts: The dialect connection options.
+
+        Raises:
+            exc.CompileError: If the table is not Iceberg or specifies a location.
+        """
+        if not self._is_iceberg_table(dialect_opts, connect_opts):
+            raise exc.CompileError(
+                "S3 Tables support only Iceberg tables; specify the dialect keyword "
+                "argument `awsathena_tblproperties={'table_type': 'ICEBERG'}`"
+            )
+        if dialect_opts["location"]:
+            raise exc.CompileError(
+                "S3 Tables use managed storage and do not accept a table location; "
+                "remove the dialect keyword argument `awsathena_location`"
+            )
+
     def _get_table_location(
         self, table: Table, dialect_opts: _DialectArgDict, connect_opts: Mapping[str, Any]
     ) -> str | None:
@@ -662,12 +741,7 @@ class AthenaDDLCompiler(DDLCompiler):
         dialect = cast("AthenaDialect", self.dialect)
         connect_opts = dialect._connect_options
 
-        table_properties = self._get_table_properties_specification(
-            dialect_opts, connect_opts
-        ).lower()
-        is_iceberg = False
-        if ("table_type" in table_properties) and ("iceberg" in table_properties):
-            is_iceberg = True
+        is_iceberg = self._is_iceberg_table(dialect_opts, connect_opts)
 
         # https://docs.aws.amazon.com/athena/latest/ug/querying-iceberg-creating-tables.html
         text = ["\nCREATE TABLE"] if is_iceberg else ["\nCREATE EXTERNAL TABLE"]
@@ -705,11 +779,19 @@ class AthenaDDLCompiler(DDLCompiler):
         dialect_opts: _DialectArgDict = table.dialect_options["awsathena"]
         dialect = cast("AthenaDialect", self.dialect)
         connect_opts = dialect._connect_options
-        text = [
-            self._get_row_format_specification(dialect_opts, connect_opts),
-            self._get_serde_properties_specification(dialect_opts, connect_opts),
-            self._get_file_format_specification(dialect_opts, connect_opts),
-            self._get_table_location_specification(table, dialect_opts, connect_opts),
-            self._get_table_properties_specification(dialect_opts, connect_opts),
-        ]
+        if self._is_s3_tables_catalog(connect_opts):
+            # S3 Tables are managed Iceberg tables: ROW FORMAT, SERDEPROPERTIES,
+            # STORED AS, and LOCATION are not accepted, so emit only TBLPROPERTIES.
+            self._validate_s3_tables_create_table(dialect_opts, connect_opts)
+            text = [
+                self._get_table_properties_specification(dialect_opts, connect_opts),
+            ]
+        else:
+            text = [
+                self._get_row_format_specification(dialect_opts, connect_opts),
+                self._get_serde_properties_specification(dialect_opts, connect_opts),
+                self._get_file_format_specification(dialect_opts, connect_opts),
+                self._get_table_location_specification(table, dialect_opts, connect_opts),
+                self._get_table_properties_specification(dialect_opts, connect_opts),
+            ]
         return "\n".join([t for t in text if t])

{pyathena-3.31.1 → pyathena-3.33.0}/pyproject.toml

@@ -5,7 +5,7 @@ authors = [
     {name = "laughingman7743", email = "laughingman7743@gmail.com"},
 ]
 dependencies = [
-    "boto3>=1.38.2",
+    "boto3>=1.41.2",
     "botocore>=1.41.2",
     "tenacity>=4.1.0",
     "fsspec",
@@ -223,12 +223,12 @@ TEST_TYPE =
 allowlist_externals =
     uv
     uvx
-    make
+    just
 commands =
     uv sync --group dev
-    pyathena: make test/pyathena
-    sqla: make test/sqla
-    sqla_async: make test/sqla-async
+    pyathena: just test pyathena
+    sqla: just test sqla
+    sqla_async: just test sqla-async
 passenv =
     TOXENV
     AWS_*