PyAthena 3.31.0__tar.gz → 3.32.0__tar.gz

{pyathena-3.31.0 → pyathena-3.32.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: PyAthena
-Version: 3.31.0
+Version: 3.32.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/

{pyathena-3.31.0 → pyathena-3.32.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.0'
-__version_tuple__ = version_tuple = (3, 31, 0)
+__version__ = version = '3.32.0'
+__version_tuple__ = version_tuple = (3, 32, 0)
 
 __commit_id__ = commit_id = None

{pyathena-3.31.0 → pyathena-3.32.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.0 → pyathena-3.32.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.0 → pyathena-3.32.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.0 → pyathena-3.32.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.0 → pyathena-3.32.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.0 → pyathena-3.32.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.0 → pyathena-3.32.0}/pyathena/filesystem/s3.py

@@ -1409,11 +1409,18 @@ class S3File(AbstractBufferedFile):
                 )
 
     def _upload_chunk(self, final: bool = False) -> bool:
+        # The return value controls whether fsspec's flush() resets self.buffer
+        # afterwards: it does so only when this returns a value other than False.
+        # Returning ``not final`` keeps the buffer intact on the final flush so a
+        # deferred commit() (autocommit=False, i.e. inside an fsspec transaction)
+        # can still read the bytes; resetting it there would upload an empty
+        # object for small files. Mid-stream chunks (final=False) return True so
+        # fsspec clears the already-uploaded buffer between parts.
         if self.tell() < self.blocksize:
             # Files smaller than block size in size cannot be multipart uploaded.
             if self.autocommit and final:
                 self.commit()
-            return True
+            return not final
 
         if not self.multipart_upload:
             raise RuntimeError("Multipart upload is not initialized.")
@@ -1456,7 +1463,7 @@ class S3File(AbstractBufferedFile):
 
         if self.autocommit and final:
             self.commit()
-        return True
+        return not final
 
     def commit(self) -> None:
         if self.tell() == 0:

{pyathena-3.31.0 → pyathena-3.32.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.0 → pyathena-3.32.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.0 → pyathena-3.32.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.0 → pyathena-3.32.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,