Merge pull request #733 from pyathena-dev/feature/execute-options

Centralize shared execute() kwargs into an ExecuteOptions dataclass

Author: laughingman7743

docs/api/connection.rst

@@ -15,6 +15,12 @@ Connection
    :members:
    :inherited-members:
 
+Execution Options
+-----------------
+
+.. autoclass:: pyathena.options.ExecuteOptions
+   :members:
+
 Standard Cursors
 ----------------
 

docs/usage.md

@@ -117,6 +117,53 @@ print(cursor.fetchall())
 
 You can find more information about the [considerations and limitations of parameterized queries](https://docs.aws.amazon.com/athena/latest/ug/querying-with-prepared-statements.html) in the official documentation.
 
+## Execution options
+
+The `execute()` method of every SQL cursor (`Cursor`, `AsyncCursor`, the aio cursors, and their
+pandas/arrow/polars/s3fs variants) accepts the same set of shared keyword arguments, such as
+`work_group`, `s3_staging_dir`, `cache_size`, `cache_expiration_time`, `result_reuse_enable`,
+`result_reuse_minutes`, `paramstyle`, `on_start_query_execution`, and `result_set_type_hints`.
+The Spark cursors execute calculations instead of SQL queries and do not accept these arguments.
+
+These arguments can also be passed together as an `ExecuteOptions` instance using the `options`
+keyword argument.
+
+```python
+from pyathena import ExecuteOptions, connect
+
+cursor = connect(s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/",
+                 region_name="us-west-2").cursor()
+options = ExecuteOptions(work_group="YOUR_WORK_GROUP",
+                         result_reuse_enable=True,
+                         result_reuse_minutes=60)
+cursor.execute("SELECT * FROM one_row", options=options)
+```
+
+When both `options` and individual keyword arguments are specified, the individual keyword
+arguments take precedence.
+
+```python
+# Executes with work_group="ANOTHER_WORK_GROUP"; the other fields of options still apply.
+cursor.execute("SELECT * FROM one_row",
+               options=options,
+               work_group="ANOTHER_WORK_GROUP")
+```
+
+Passing `None` for an individual keyword argument is treated as "not specified" and does not
+reset the corresponding `options` field. To run a single query without a field set on `options`,
+construct a new instance without that field.
+
+`ExecuteOptions` is immutable. To create a variation of an existing instance, use the `merge()`
+method, which returns a new instance with the specified fields applied.
+
+```python
+adhoc_options = options.merge(work_group="ANOTHER_WORK_GROUP")
+```
+
+The `on_start_query_execution` field is invoked by the synchronous and aio cursors;
+`AsyncCursor`-based cursors return the query ID directly through their execution model and do
+not invoke it.
+
 ## Quickly re-run queries
 
 ### Result reuse configuration
@@ -423,6 +470,7 @@ The `on_start_query_execution` callback is supported by the following cursor typ
 - `PandasCursor`
 - `PolarsCursor`
 - `S3FSCursor`
+- `AioCursor`, `AioDictCursor`, `AioArrowCursor`, `AioPandasCursor`, `AioPolarsCursor`, `AioS3FSCursor`
 
 Note: `AsyncCursor` and its variants do not support this callback as they already
 return the query ID immediately through their different execution model.

pyathena/__init__.py

@@ -4,6 +4,7 @@
 from typing import TYPE_CHECKING, Any, overload
 
 from pyathena.error import *  # noqa: F403
+from pyathena.options import ExecuteOptions as ExecuteOptions
 
 if TYPE_CHECKING:
     from pyathena.aio.connection import AioConnection

pyathena/aio/arrow/cursor.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import logging
+from collections.abc import Callable
 from typing import TYPE_CHECKING, Any, cast
 
 from pyathena.aio.common import WithAsyncFetch
@@ -13,6 +14,7 @@
 from pyathena.common import CursorIterator
 from pyathena.error import OperationalError, ProgrammingError
 from pyathena.model import AthenaQueryExecution
+from pyathena.options import ExecuteOptions
 
 if TYPE_CHECKING:
     import polars as pl
@@ -83,11 +85,15 @@ async def execute(  # type: ignore[override]
         parameters: dict[str, Any] | list[str] | None = None,
         work_group: str | None = None,
         s3_staging_dir: str | None = None,
-        cache_size: int | None = 0,
-        cache_expiration_time: int | None = 0,
+        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,
     ) -> AioArrowCursor:
         """Execute a SQL query asynchronously and return results as Arrow Tables.
@@ -102,25 +108,41 @@ async def execute(  # type: ignore[override]
             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()
-        operation, unload_location = self._prepare_unload(operation, s3_staging_dir)
-        self.query_id = await self._execute(
-            operation,
-            parameters=parameters,
+        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,
+        )
+        operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir)
+        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(
@@ -134,6 +156,7 @@ async def execute(  # type: ignore[override]
                 unload_location=unload_location,
                 connect_timeout=self._connect_timeout,
                 request_timeout=self._request_timeout,
+                result_set_type_hints=options.result_set_type_hints,
                 **kwargs,
             )
         else:

pyathena/aio/common.py

@@ -10,6 +10,7 @@
 from pyathena.common import BaseCursor, CursorIterator
 from pyathena.error import DatabaseError, OperationalError, ProgrammingError
 from pyathena.model import AthenaDatabase, AthenaQueryExecution, AthenaTableMetadata
+from pyathena.options import ExecuteOptions
 from pyathena.result_set import AthenaResultSet, WithResultSet
 
 _logger = logging.getLogger(__name__)
@@ -27,29 +28,24 @@ 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 = 0,
-        cache_expiration_time: int | None = 0,
-        result_reuse_enable: bool | None = None,
-        result_reuse_minutes: int | None = None,
-        paramstyle: str | None = None,
+        options: ExecuteOptions | None = None,
     ) -> str:
-        query, execution_parameters = self._prepare_query(operation, parameters, paramstyle)
+        options = ExecuteOptions.resolve(options)
+        query, execution_parameters = self._prepare_query(operation, parameters, options.paramstyle)
 
         request = self._build_start_query_execution_request(
             query=query,
-            work_group=work_group,
-            s3_staging_dir=s3_staging_dir,
-            result_reuse_enable=result_reuse_enable,
-            result_reuse_minutes=result_reuse_minutes,
+            work_group=options.work_group,
+            s3_staging_dir=options.s3_staging_dir,
+            result_reuse_enable=options.result_reuse_enable,
+            result_reuse_minutes=options.result_reuse_minutes,
             execution_parameters=execution_parameters,
         )
         query_id = await self._find_previous_query_id(
             query,
-            work_group,
-            cache_size=cache_size if cache_size else 0,
-            cache_expiration_time=cache_expiration_time if cache_expiration_time else 0,
+            options.work_group,
+            cache_size=options.cache_size,
+            cache_expiration_time=options.cache_expiration_time,
         )
         if query_id is None:
             try:

pyathena/aio/cursor.py

@@ -1,13 +1,15 @@
 from __future__ import annotations
 
 import logging
+from collections.abc import Callable
 from typing import Any, cast
 
 from pyathena.aio.common import WithAsyncFetch
 from pyathena.aio.result_set import AthenaAioDictResultSet, AthenaAioResultSet
 from pyathena.common import CursorIterator
 from pyathena.error import OperationalError, ProgrammingError
 from pyathena.model import AthenaQueryExecution
+from pyathena.options import ExecuteOptions
 
 _logger = logging.getLogger(__name__)
 
@@ -74,12 +76,15 @@ async def execute(  # type: ignore[override]
         parameters: dict[str, Any] | list[str] | None = None,
         work_group: str | None = None,
         s3_staging_dir: str | None = None,
-        cache_size: int = 0,
-        cache_expiration_time: int = 0,
+        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,
     ) -> AioCursor:
         """Execute a SQL query asynchronously.
@@ -94,27 +99,40 @@ async def execute(  # type: ignore[override]
             result_reuse_enable: Enable result reuse (optional).
             result_reuse_minutes: Result reuse duration in minutes (optional).
             paramstyle: Parameter style to use (optional).
+            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()
-        self.query_id = await self._execute(
-            operation,
-            parameters=parameters,
+        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 self._result_set_class.create(
@@ -123,7 +141,7 @@ async def execute(  # type: ignore[override]
                 query_execution,
                 self.arraysize,
                 self._retry_config,
-                result_set_type_hints=result_set_type_hints,
+                result_set_type_hints=options.result_set_type_hints,
             )
         else:
             raise OperationalError(query_execution.state_change_reason)

pyathena/aio/pandas/cursor.py

@@ -2,7 +2,7 @@
 
 import asyncio
 import logging
-from collections.abc import Iterable
+from collections.abc import Callable, Iterable
 from multiprocessing import cpu_count
 from typing import (
     TYPE_CHECKING,
@@ -14,6 +14,7 @@
 from pyathena.common import CursorIterator
 from pyathena.error import OperationalError, ProgrammingError
 from pyathena.model import AthenaQueryExecution
+from pyathena.options import ExecuteOptions
 from pyathena.pandas.converter import (
     DefaultPandasTypeConverter,
     DefaultPandasUnloadTypeConverter,
@@ -97,14 +98,18 @@ async def execute(  # type: ignore[override]
         parameters: dict[str, Any] | list[str] | None = None,
         work_group: str | None = None,
         s3_staging_dir: str | None = None,
-        cache_size: int | None = 0,
-        cache_expiration_time: int | None = 0,
+        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,
         keep_default_na: bool = False,
         na_values: Iterable[str] | None = ("",),
         quoting: int = 1,
+        on_start_query_execution: Callable[[str], None] | None = None,
+        result_set_type_hints: dict[str | int, str] | None = None,
+        *,
+        options: ExecuteOptions | None = None,
         **kwargs,
     ) -> AioPandasCursor:
         """Execute a SQL query asynchronously and return results as pandas DataFrames.
@@ -122,25 +127,41 @@ async def execute(  # type: ignore[override]
             keep_default_na: Whether to keep default pandas NA values.
             na_values: Additional values to treat as NA.
             quoting: CSV quoting behavior (pandas csv.QUOTE_* constants).
+            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 pandas read_csv/read_parquet 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,
+        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,
+        )
+        operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir)
+        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(
@@ -161,6 +182,7 @@ async def execute(  # type: ignore[override]
                 cache_type=kwargs.pop("cache_type", self._cache_type),
                 max_workers=kwargs.pop("max_workers", self._max_workers),
                 auto_optimize_chunksize=self._auto_optimize_chunksize,
+                result_set_type_hints=options.result_set_type_hints,
                 **kwargs,
             )
         else: