From a26c3fb0524e322264ac0a8c02dbd27c44998bc4 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 6 Jul 2026 23:14:46 +0900 Subject: [PATCH 1/4] Centralize shared execute() kwargs into an ExecuteOptions dataclass Extract the execution arguments shared by all cursor implementations (work_group, s3_staging_dir, cache_size, cache_expiration_time, result_reuse_enable, result_reuse_minutes, paramstyle, on_start_query_execution, result_set_type_hints) into a frozen ExecuteOptions dataclass in pyathena/options.py. - BaseCursor._execute() and AioBaseCursor._execute() now take a single ExecuteOptions instead of nine individual kwargs - Every cursor's execute() keeps its existing keyword arguments for backward compatibility and additionally accepts a keyword-only options parameter; individual kwargs take precedence over options - Adding a new shared execute() argument now only requires changes to the dataclass and the internals, not all cursor signatures - Add the explicit result_set_type_hints argument to the aio pandas/arrow/polars/s3fs cursors, which previously relied on kwargs forwarding, for consistency with the other cursors - Export ExecuteOptions from the pyathena package root and document it in the usage guide and API reference Closes #691 Co-Authored-By: Claude Fable 5 --- docs/api/connection.rst | 6 ++ docs/usage.md | 40 +++++++++++++ pyathena/__init__.py | 1 + pyathena/aio/arrow/cursor.py | 27 +++++++-- pyathena/aio/common.py | 28 +++++---- pyathena/aio/cursor.py | 22 +++++-- pyathena/aio/pandas/cursor.py | 27 +++++++-- pyathena/aio/polars/cursor.py | 27 +++++++-- pyathena/aio/s3fs/cursor.py | 25 ++++++-- pyathena/arrow/async_cursor.py | 21 ++++--- pyathena/arrow/cursor.py | 29 ++++++--- pyathena/async_cursor.py | 22 +++++-- pyathena/common.py | 28 +++++---- pyathena/cursor.py | 27 ++++++--- pyathena/options.py | 90 ++++++++++++++++++++++++++++ pyathena/pandas/async_cursor.py | 21 ++++--- pyathena/pandas/cursor.py | 29 ++++++--- pyathena/polars/async_cursor.py | 24 +++++--- pyathena/polars/cursor.py | 29 ++++++--- pyathena/s3fs/async_cursor.py | 22 +++++-- pyathena/s3fs/cursor.py | 27 ++++++--- tests/pyathena/test_cursor.py | 28 ++++++++- tests/pyathena/test_options.py | 100 ++++++++++++++++++++++++++++++++ 23 files changed, 564 insertions(+), 136 deletions(-) create mode 100644 pyathena/options.py create mode 100644 tests/pyathena/test_options.py diff --git a/docs/api/connection.rst b/docs/api/connection.rst index 8bdf0e37..bb186eb7 100644 --- a/docs/api/connection.rst +++ b/docs/api/connection.rst @@ -15,6 +15,12 @@ Connection :members: :inherited-members: +Execution Options +----------------- + +.. autoclass:: pyathena.options.ExecuteOptions + :members: + Standard Cursors ---------------- diff --git a/docs/usage.md b/docs/usage.md index 7b2a84d4..59e04a82 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -117,6 +117,46 @@ 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 cursor 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`. + +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") +``` + +`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 only invoked by synchronous cursors; asynchronous cursors +return the query ID directly through their execution model. + ## Quickly re-run queries ### Result reuse configuration diff --git a/pyathena/__init__.py b/pyathena/__init__.py index c3ad35c1..e462b28b 100644 --- a/pyathena/__init__.py +++ b/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 diff --git a/pyathena/aio/arrow/cursor.py b/pyathena/aio/arrow/cursor.py index 9baebc77..cc46a3eb 100644 --- a/pyathena/aio/arrow/cursor.py +++ b/pyathena/aio/arrow/cursor.py @@ -13,6 +13,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 +84,14 @@ 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, + 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,16 +106,19 @@ 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'). + 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 = (options if options is not None else ExecuteOptions()).merge( work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -119,6 +126,13 @@ async def execute( # type: ignore[override] result_reuse_enable=result_reuse_enable, result_reuse_minutes=result_reuse_minutes, paramstyle=paramstyle, + 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, ) query_execution = await self._poll(self.query_id) @@ -134,6 +148,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: diff --git a/pyathena/aio/common.py b/pyathena/aio/common.py index b24eb596..f971a641 100644 --- a/pyathena/aio/common.py +++ b/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,26 @@ 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 = options if options is not None else ExecuteOptions() + 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 if options.cache_size else 0, + cache_expiration_time=options.cache_expiration_time + if options.cache_expiration_time + else 0, ) if query_id is None: try: diff --git a/pyathena/aio/cursor.py b/pyathena/aio/cursor.py index 07bf507e..2e3fe4ca 100644 --- a/pyathena/aio/cursor.py +++ b/pyathena/aio/cursor.py @@ -8,6 +8,7 @@ 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 +75,14 @@ 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, result_set_type_hints: dict[str | int, str] | None = None, + *, + options: ExecuteOptions | None = None, **kwargs, ) -> AioCursor: """Execute a SQL query asynchronously. @@ -97,15 +100,16 @@ async def execute( # type: ignore[override] 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 = (options if options is not None else ExecuteOptions()).merge( work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -113,6 +117,12 @@ async def execute( # type: ignore[override] result_reuse_enable=result_reuse_enable, result_reuse_minutes=result_reuse_minutes, paramstyle=paramstyle, + result_set_type_hints=result_set_type_hints, + ) + self.query_id = await self._execute( + operation, + parameters=parameters, + options=options, ) query_execution = await self._poll(self.query_id) @@ -123,7 +133,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) diff --git a/pyathena/aio/pandas/cursor.py b/pyathena/aio/pandas/cursor.py index 88b0f1c6..f0eeb16d 100644 --- a/pyathena/aio/pandas/cursor.py +++ b/pyathena/aio/pandas/cursor.py @@ -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,17 @@ 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, + 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,16 +126,19 @@ 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). + 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 = (options if options is not None else ExecuteOptions()).merge( work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -139,6 +146,13 @@ async def execute( # type: ignore[override] result_reuse_enable=result_reuse_enable, result_reuse_minutes=result_reuse_minutes, paramstyle=paramstyle, + 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, ) query_execution = await self._poll(self.query_id) @@ -161,6 +175,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: diff --git a/pyathena/aio/polars/cursor.py b/pyathena/aio/polars/cursor.py index 484c413f..82227822 100644 --- a/pyathena/aio/polars/cursor.py +++ b/pyathena/aio/polars/cursor.py @@ -9,6 +9,7 @@ from pyathena.common import CursorIterator from pyathena.error import OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution +from pyathena.options import ExecuteOptions from pyathena.polars.converter import ( DefaultPolarsTypeConverter, DefaultPolarsUnloadTypeConverter, @@ -89,11 +90,14 @@ 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, + result_set_type_hints: dict[str | int, str] | None = None, + *, + options: ExecuteOptions | None = None, **kwargs, ) -> AioPolarsCursor: """Execute a SQL query asynchronously and return results as Polars DataFrames. @@ -108,16 +112,19 @@ 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'). + 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 passed to Polars read functions. 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 = (options if options is not None else ExecuteOptions()).merge( work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -125,6 +132,13 @@ async def execute( # type: ignore[override] result_reuse_enable=result_reuse_enable, result_reuse_minutes=result_reuse_minutes, paramstyle=paramstyle, + 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, ) query_execution = await self._poll(self.query_id) @@ -142,6 +156,7 @@ async def execute( # type: ignore[override] cache_type=self._cache_type, max_workers=self._max_workers, chunksize=self._chunksize, + result_set_type_hints=options.result_set_type_hints, **kwargs, ) else: diff --git a/pyathena/aio/s3fs/cursor.py b/pyathena/aio/s3fs/cursor.py index cf9c92a8..316a20a5 100644 --- a/pyathena/aio/s3fs/cursor.py +++ b/pyathena/aio/s3fs/cursor.py @@ -9,6 +9,7 @@ from pyathena.error import OperationalError, ProgrammingError from pyathena.filesystem.s3_async import AioS3FileSystem from pyathena.model import AthenaQueryExecution +from pyathena.options import ExecuteOptions from pyathena.s3fs.converter import DefaultS3FSTypeConverter from pyathena.s3fs.result_set import AthenaS3FSResultSet, CSVReaderType @@ -81,11 +82,14 @@ 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, + result_set_type_hints: dict[str | int, str] | None = None, + *, + options: ExecuteOptions | None = None, **kwargs, ) -> AioS3FSCursor: """Execute a SQL query asynchronously via S3FileSystem CSV reader. @@ -100,15 +104,19 @@ 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'). + 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 = (options if options is not None else ExecuteOptions()).merge( work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -116,6 +124,12 @@ async def execute( # type: ignore[override] result_reuse_enable=result_reuse_enable, result_reuse_minutes=result_reuse_minutes, paramstyle=paramstyle, + result_set_type_hints=result_set_type_hints, + ) + self.query_id = await self._execute( + operation, + parameters=parameters, + options=options, ) query_execution = await self._poll(self.query_id) @@ -129,6 +143,7 @@ async def execute( # type: ignore[override] retry_config=self._retry_config, csv_reader=self._csv_reader, filesystem_class=AioS3FileSystem, + result_set_type_hints=options.result_set_type_hints, **kwargs, ) else: diff --git a/pyathena/arrow/async_cursor.py b/pyathena/arrow/async_cursor.py index da84b5c9..386da908 100644 --- a/pyathena/arrow/async_cursor.py +++ b/pyathena/arrow/async_cursor.py @@ -14,6 +14,7 @@ from pyathena.async_cursor import AsyncCursor from pyathena.common import CursorIterator from pyathena.model import AthenaQueryExecution +from pyathena.options import ExecuteOptions _logger = logging.getLogger(__name__) @@ -176,18 +177,17 @@ def execute( 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, result_set_type_hints: dict[str | int, str] | None = None, + *, + options: ExecuteOptions | None = None, **kwargs, ) -> tuple[str, Future[AthenaArrowResultSet | Any]]: - operation, unload_location = self._prepare_unload(operation, s3_staging_dir) - query_id = self._execute( - operation, - parameters=parameters, + options = (options if options is not None else ExecuteOptions()).merge( work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -195,13 +195,20 @@ def execute( result_reuse_enable=result_reuse_enable, result_reuse_minutes=result_reuse_minutes, paramstyle=paramstyle, + result_set_type_hints=result_set_type_hints, + ) + operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) + query_id = self._execute( + operation, + parameters=parameters, + options=options, ) return ( query_id, self._executor.submit( self._collect_result_set, query_id, - result_set_type_hints, + options.result_set_type_hints, unload_location, kwargs, ), diff --git a/pyathena/arrow/cursor.py b/pyathena/arrow/cursor.py index b4775df6..9c566684 100644 --- a/pyathena/arrow/cursor.py +++ b/pyathena/arrow/cursor.py @@ -12,6 +12,7 @@ from pyathena.common import CursorIterator from pyathena.error import OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution +from pyathena.options import ExecuteOptions from pyathena.result_set import WithFetch if TYPE_CHECKING: @@ -128,13 +129,15 @@ def execute( 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, ) -> ArrowCursor: """Execute a SQL query and return results as Apache Arrow Tables. @@ -157,6 +160,9 @@ def execute( 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: @@ -167,10 +173,7 @@ def execute( >>> table = cursor.as_arrow() # Returns Apache Arrow Table """ self._reset_state() - operation, unload_location = self._prepare_unload(operation, s3_staging_dir) - self.query_id = self._execute( - operation, - parameters=parameters, + options = (options if options is not None else ExecuteOptions()).merge( work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -178,14 +181,22 @@ def execute( 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 = self._execute( + operation, + parameters=parameters, + options=options, ) # Call user callbacks immediately after start_query_execution # Both connection-level and execute-level callbacks are invoked if set if self._on_start_query_execution: self._on_start_query_execution(self.query_id) - if on_start_query_execution: - on_start_query_execution(self.query_id) + if options.on_start_query_execution: + options.on_start_query_execution(self.query_id) query_execution = cast(AthenaQueryExecution, self._poll(self.query_id)) if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED: self.result_set = AthenaArrowResultSet( @@ -198,7 +209,7 @@ def execute( unload_location=unload_location, connect_timeout=self._connect_timeout, request_timeout=self._request_timeout, - result_set_type_hints=result_set_type_hints, + result_set_type_hints=options.result_set_type_hints, **kwargs, ) else: diff --git a/pyathena/async_cursor.py b/pyathena/async_cursor.py index 8ba89073..7ede4da6 100644 --- a/pyathena/async_cursor.py +++ b/pyathena/async_cursor.py @@ -9,6 +9,7 @@ from pyathena.common import BaseCursor, CursorIterator from pyathena.error import NotSupportedError, ProgrammingError from pyathena.model import AthenaQueryExecution +from pyathena.options import ExecuteOptions from pyathena.result_set import AthenaDictResultSet, AthenaResultSet _logger = logging.getLogger(__name__) @@ -165,12 +166,14 @@ def execute( 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, result_set_type_hints: dict[str | int, str] | None = None, + *, + options: ExecuteOptions | None = None, **kwargs, ) -> tuple[str, Future[AthenaResultSet | Any]]: """Execute a SQL query asynchronously. @@ -192,6 +195,9 @@ def execute( 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: @@ -205,9 +211,7 @@ def execute( >>> # Do other work while query runs... >>> result_set = future.result() # Wait for completion """ - query_id = self._execute( - operation, - parameters=parameters, + options = (options if options is not None else ExecuteOptions()).merge( work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -215,9 +219,15 @@ def execute( result_reuse_enable=result_reuse_enable, result_reuse_minutes=result_reuse_minutes, paramstyle=paramstyle, + result_set_type_hints=result_set_type_hints, + ) + query_id = self._execute( + operation, + parameters=parameters, + options=options, ) return query_id, self._executor.submit( - self._collect_result_set, query_id, result_set_type_hints + self._collect_result_set, query_id, options.result_set_type_hints ) def executemany( diff --git a/pyathena/common.py b/pyathena/common.py index 04d919b4..50d78656 100644 --- a/pyathena/common.py +++ b/pyathena/common.py @@ -21,6 +21,7 @@ AthenaQueryExecution, AthenaTableMetadata, ) +from pyathena.options import ExecuteOptions from pyathena.util import RetryConfig, retry_api_call if TYPE_CHECKING: @@ -701,29 +702,26 @@ def _execute( 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 = options if options is not None else ExecuteOptions() + 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 = 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 if options.cache_size else 0, + cache_expiration_time=options.cache_expiration_time + if options.cache_expiration_time + else 0, ) if query_id is None: try: diff --git a/pyathena/cursor.py b/pyathena/cursor.py index a9438d0a..b279d83a 100644 --- a/pyathena/cursor.py +++ b/pyathena/cursor.py @@ -7,6 +7,7 @@ from pyathena.common import CursorIterator from pyathena.error import OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution +from pyathena.options import ExecuteOptions from pyathena.result_set import AthenaDictResultSet, AthenaResultSet, WithFetch _logger = logging.getLogger(__name__) @@ -87,13 +88,15 @@ def execute( 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, ) -> Cursor: """Execute a SQL query. @@ -110,6 +113,9 @@ def execute( Athena DDL type signatures for precise type conversion within complex types. For example: ``{"tags": "array(varchar)", "metadata": "map(varchar, integer)"}`` + options: Shared execution options as an + :class:`~pyathena.options.ExecuteOptions` instance. Individual + keyword arguments take precedence over ``options`` fields. **kwargs: Additional execution parameters. Returns: @@ -125,9 +131,7 @@ def execute( ... ) """ self._reset_state() - self.query_id = self._execute( - operation, - parameters=parameters, + options = (options if options is not None else ExecuteOptions()).merge( work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -135,14 +139,21 @@ def execute( 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 = self._execute( + operation, + parameters=parameters, + options=options, ) # Call user callbacks immediately after start_query_execution # Both connection-level and execute-level callbacks are invoked if set if self._on_start_query_execution: self._on_start_query_execution(self.query_id) - if on_start_query_execution: - on_start_query_execution(self.query_id) + if options.on_start_query_execution: + options.on_start_query_execution(self.query_id) query_execution = cast(AthenaQueryExecution, self._poll(self.query_id)) if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED: @@ -152,7 +163,7 @@ def execute( 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) diff --git a/pyathena/options.py b/pyathena/options.py new file mode 100644 index 00000000..066a2cd7 --- /dev/null +++ b/pyathena/options.py @@ -0,0 +1,90 @@ +from __future__ import annotations + +import logging +from collections.abc import Callable +from dataclasses import dataclass, fields, replace +from typing import Any + +_logger = logging.getLogger(__name__) + + +@dataclass(frozen=True) +class ExecuteOptions: + """Shared options for ``Cursor.execute()`` across all cursor implementations. + + This dataclass is the single source of truth for the query-execution + arguments shared by every cursor type (sync/async/aio and + pandas/arrow/polars/s3fs variants). It can be passed to ``execute()`` + via the ``options`` keyword argument as an alternative to individual + keyword arguments: + + >>> from pyathena.options import ExecuteOptions + >>> options = ExecuteOptions(work_group="primary", cache_size=100) + >>> cursor.execute("SELECT * FROM my_table", options=options) + + When both ``options`` and individual keyword arguments are provided, + the individual keyword arguments take precedence. This allows building + a base ``ExecuteOptions`` once and tweaking it per call: + + >>> cursor.execute("SELECT ...", options=options, work_group="adhoc") + + Attributes: + work_group: Athena workgroup to use for this query. Overrides the + connection-level workgroup. + s3_staging_dir: S3 location for query results. Overrides the + connection-level staging directory. + cache_size: Number of recent queries to scan for client-side result + caching. 0 (default) disables the cache lookup. + cache_expiration_time: Maximum age in seconds of a cached query + result to consider for reuse. 0 (default) means no age limit. + result_reuse_enable: Enable Athena server-side result reuse for this + query. None (default) falls back to the connection-level setting. + result_reuse_minutes: Maximum age in minutes of a previous query + result that Athena should consider for reuse. None (default) + falls back to the connection-level setting. + paramstyle: Parameter style for this query ('qmark' or 'pyformat'). + None (default) uses the module-level ``pyathena.paramstyle``. + on_start_query_execution: Callback invoked with the query ID + immediately after the StartQueryExecution API call. Only invoked + by synchronous cursors; asynchronous cursors return the query ID + directly through their execution model. + result_set_type_hints: Mapping of column names (or indices) to Athena + DDL type signatures for precise type conversion within complex + types. For example: + ``{"tags": "array(varchar)", "metadata": "map(varchar, integer)"}`` + """ + + 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 + on_start_query_execution: Callable[[str], None] | None = None + result_set_type_hints: dict[str | int, str] | None = None + + def merge(self, **overrides: Any) -> ExecuteOptions: + """Return a new instance with non-None ``overrides`` applied. + + Args: + **overrides: Field values to apply on top of this instance. + None values are ignored, so an omitted ``execute()`` keyword + argument never clobbers a value set on ``options``. + + Returns: + A new ``ExecuteOptions`` with the overrides applied. + + Raises: + TypeError: If an override name is not a field of this class. + """ + valid = {f.name for f in fields(self)} + unknown = set(overrides) - valid + if unknown: + raise TypeError( + f"Unknown {self.__class__.__name__} fields: {', '.join(sorted(unknown))}" + ) + applied = {k: v for k, v in overrides.items() if v is not None} + if not applied: + return self + return replace(self, **applied) diff --git a/pyathena/pandas/async_cursor.py b/pyathena/pandas/async_cursor.py index db9f3f9e..b7bcb535 100644 --- a/pyathena/pandas/async_cursor.py +++ b/pyathena/pandas/async_cursor.py @@ -10,6 +10,7 @@ from pyathena.async_cursor import AsyncCursor from pyathena.common import CursorIterator from pyathena.model import AthenaQueryExecution +from pyathena.options import ExecuteOptions from pyathena.pandas.converter import ( DefaultPandasTypeConverter, DefaultPandasUnloadTypeConverter, @@ -151,8 +152,8 @@ def execute( 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, @@ -160,12 +161,11 @@ def execute( keep_default_na: bool = False, na_values: Iterable[str] | None = ("",), quoting: int = 1, + *, + options: ExecuteOptions | None = None, **kwargs, ) -> tuple[str, Future[AthenaPandasResultSet | Any]]: - operation, unload_location = self._prepare_unload(operation, s3_staging_dir) - query_id = self._execute( - operation, - parameters=parameters, + options = (options if options is not None else ExecuteOptions()).merge( work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -173,13 +173,20 @@ def execute( result_reuse_enable=result_reuse_enable, result_reuse_minutes=result_reuse_minutes, paramstyle=paramstyle, + result_set_type_hints=result_set_type_hints, + ) + operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) + query_id = self._execute( + operation, + parameters=parameters, + options=options, ) return ( query_id, self._executor.submit( self._collect_result_set, query_id, - result_set_type_hints, + options.result_set_type_hints, keep_default_na, na_values, quoting, diff --git a/pyathena/pandas/cursor.py b/pyathena/pandas/cursor.py index 3b55e7b6..89a73c16 100644 --- a/pyathena/pandas/cursor.py +++ b/pyathena/pandas/cursor.py @@ -12,6 +12,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, @@ -141,8 +142,8 @@ def execute( 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, @@ -151,6 +152,8 @@ def execute( 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, ) -> PandasCursor: """Execute a SQL query and return results as pandas DataFrames. @@ -176,6 +179,9 @@ def execute( 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: @@ -187,10 +193,7 @@ def execute( >>> df = cursor.fetchall() # Returns pandas DataFrame """ self._reset_state() - operation, unload_location = self._prepare_unload(operation, s3_staging_dir) - self.query_id = self._execute( - operation, - parameters=parameters, + options = (options if options is not None else ExecuteOptions()).merge( work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -198,14 +201,22 @@ def execute( 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 = self._execute( + operation, + parameters=parameters, + options=options, ) # Call user callbacks immediately after start_query_execution # Both connection-level and execute-level callbacks are invoked if set if self._on_start_query_execution: self._on_start_query_execution(self.query_id) - if on_start_query_execution: - on_start_query_execution(self.query_id) + if options.on_start_query_execution: + options.on_start_query_execution(self.query_id) query_execution = cast(AthenaQueryExecution, self._poll(self.query_id)) if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED: self.result_set = AthenaPandasResultSet( @@ -225,7 +236,7 @@ def execute( 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=result_set_type_hints, + result_set_type_hints=options.result_set_type_hints, **kwargs, ) else: diff --git a/pyathena/polars/async_cursor.py b/pyathena/polars/async_cursor.py index 3774a47e..5c37857c 100644 --- a/pyathena/polars/async_cursor.py +++ b/pyathena/polars/async_cursor.py @@ -9,6 +9,7 @@ from pyathena.async_cursor import AsyncCursor from pyathena.common import CursorIterator from pyathena.model import AthenaQueryExecution +from pyathena.options import ExecuteOptions from pyathena.polars.converter import ( DefaultPolarsTypeConverter, DefaultPolarsUnloadTypeConverter, @@ -190,12 +191,14 @@ def execute( 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, result_set_type_hints: dict[str | int, str] | None = None, + *, + options: ExecuteOptions | None = None, **kwargs, ) -> tuple[str, Future[AthenaPolarsResultSet | Any]]: """Execute a SQL query asynchronously and return results as Polars DataFrames. @@ -216,6 +219,9 @@ def execute( 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 passed to Polars read functions. Returns: @@ -226,10 +232,7 @@ def execute( >>> result_set = future.result() >>> df = result_set.as_polars() # Returns Polars DataFrame """ - operation, unload_location = self._prepare_unload(operation, s3_staging_dir) - query_id = self._execute( - operation, - parameters=parameters, + options = (options if options is not None else ExecuteOptions()).merge( work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -237,13 +240,20 @@ def execute( result_reuse_enable=result_reuse_enable, result_reuse_minutes=result_reuse_minutes, paramstyle=paramstyle, + result_set_type_hints=result_set_type_hints, + ) + operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) + query_id = self._execute( + operation, + parameters=parameters, + options=options, ) return ( query_id, self._executor.submit( self._collect_result_set, query_id, - result_set_type_hints, + options.result_set_type_hints, unload_location, kwargs, ), diff --git a/pyathena/polars/cursor.py b/pyathena/polars/cursor.py index 12ab6b25..632fb44a 100644 --- a/pyathena/polars/cursor.py +++ b/pyathena/polars/cursor.py @@ -12,6 +12,7 @@ from pyathena.common import CursorIterator from pyathena.error import OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution +from pyathena.options import ExecuteOptions from pyathena.polars.converter import ( DefaultPolarsTypeConverter, DefaultPolarsUnloadTypeConverter, @@ -148,13 +149,15 @@ def execute( 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, ) -> PolarsCursor: """Execute a SQL query and return results as Polars DataFrames. @@ -176,6 +179,9 @@ def execute( 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 passed to Polars read functions. Returns: @@ -186,10 +192,7 @@ def execute( >>> df = cursor.as_polars() # Returns Polars DataFrame """ self._reset_state() - operation, unload_location = self._prepare_unload(operation, s3_staging_dir) - self.query_id = self._execute( - operation, - parameters=parameters, + options = (options if options is not None else ExecuteOptions()).merge( work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -197,14 +200,22 @@ def execute( 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 = self._execute( + operation, + parameters=parameters, + options=options, ) # Call user callbacks immediately after start_query_execution # Both connection-level and execute-level callbacks are invoked if set if self._on_start_query_execution: self._on_start_query_execution(self.query_id) - if on_start_query_execution: - on_start_query_execution(self.query_id) + if options.on_start_query_execution: + options.on_start_query_execution(self.query_id) query_execution = cast(AthenaQueryExecution, self._poll(self.query_id)) if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED: self.result_set = AthenaPolarsResultSet( @@ -219,7 +230,7 @@ def execute( cache_type=self._cache_type, max_workers=self._max_workers, chunksize=self._chunksize, - result_set_type_hints=result_set_type_hints, + result_set_type_hints=options.result_set_type_hints, **kwargs, ) else: diff --git a/pyathena/s3fs/async_cursor.py b/pyathena/s3fs/async_cursor.py index 73acede2..79f4defd 100644 --- a/pyathena/s3fs/async_cursor.py +++ b/pyathena/s3fs/async_cursor.py @@ -9,6 +9,7 @@ from pyathena.common import CursorIterator from pyathena.error import ProgrammingError from pyathena.model import AthenaQueryExecution +from pyathena.options import ExecuteOptions from pyathena.s3fs.converter import DefaultS3FSTypeConverter from pyathena.s3fs.result_set import AthenaS3FSResultSet, CSVReaderType @@ -177,12 +178,14 @@ def execute( 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, result_set_type_hints: dict[str | int, str] | None = None, + *, + options: ExecuteOptions | None = None, **kwargs, ) -> tuple[str, Future[AthenaS3FSResultSet | Any]]: """Execute a SQL query asynchronously. @@ -203,6 +206,9 @@ def execute( 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: @@ -213,9 +219,7 @@ def execute( >>> result_set = future.result() >>> rows = result_set.fetchall() """ - query_id = self._execute( - operation, - parameters=parameters, + options = (options if options is not None else ExecuteOptions()).merge( work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -223,13 +227,19 @@ def execute( result_reuse_enable=result_reuse_enable, result_reuse_minutes=result_reuse_minutes, paramstyle=paramstyle, + result_set_type_hints=result_set_type_hints, + ) + query_id = self._execute( + operation, + parameters=parameters, + options=options, ) return ( query_id, self._executor.submit( self._collect_result_set, query_id, - result_set_type_hints, + options.result_set_type_hints, kwargs, ), ) diff --git a/pyathena/s3fs/cursor.py b/pyathena/s3fs/cursor.py index f0521f7f..c5525baf 100644 --- a/pyathena/s3fs/cursor.py +++ b/pyathena/s3fs/cursor.py @@ -7,6 +7,7 @@ from pyathena.common import CursorIterator from pyathena.error import OperationalError from pyathena.model import AthenaQueryExecution +from pyathena.options import ExecuteOptions from pyathena.result_set import WithFetch from pyathena.s3fs.converter import DefaultS3FSTypeConverter from pyathena.s3fs.result_set import AthenaS3FSResultSet, CSVReaderType @@ -124,13 +125,15 @@ def execute( 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, ) -> S3FSCursor: """Execute a SQL query and return results. @@ -152,6 +155,9 @@ def execute( 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: @@ -162,9 +168,7 @@ def execute( >>> rows = cursor.fetchall() """ self._reset_state() - self.query_id = self._execute( - operation, - parameters=parameters, + options = (options if options is not None else ExecuteOptions()).merge( work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -172,13 +176,20 @@ def execute( 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 = self._execute( + operation, + parameters=parameters, + options=options, ) # Call user callbacks immediately after start_query_execution if self._on_start_query_execution: self._on_start_query_execution(self.query_id) - if on_start_query_execution: - on_start_query_execution(self.query_id) + if options.on_start_query_execution: + options.on_start_query_execution(self.query_id) query_execution = cast(AthenaQueryExecution, self._poll(self.query_id)) if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED: @@ -189,7 +200,7 @@ def execute( arraysize=self.arraysize, retry_config=self._retry_config, csv_reader=self._csv_reader, - result_set_type_hints=result_set_type_hints, + result_set_type_hints=options.result_set_type_hints, **kwargs, ) else: diff --git a/tests/pyathena/test_cursor.py b/tests/pyathena/test_cursor.py index 15155180..30c5bb7b 100644 --- a/tests/pyathena/test_cursor.py +++ b/tests/pyathena/test_cursor.py @@ -15,7 +15,7 @@ import pytest -from pyathena import BINARY, BOOLEAN, DATE, DATETIME, JSON, NUMBER, STRING, TIME +from pyathena import BINARY, BOOLEAN, DATE, DATETIME, JSON, NUMBER, STRING, TIME, ExecuteOptions from pyathena.converter import _to_array, _to_map, _to_struct from pyathena.cursor import Cursor from pyathena.error import DatabaseError, NotSupportedError, ProgrammingError @@ -853,6 +853,32 @@ def on_start(query_id): row = cursor.fetchone() assert row == (1,) + def test_execute_with_options(self, cursor): + """Test execute with an ExecuteOptions instance instead of individual kwargs.""" + callback_results = [] + options = ExecuteOptions( + on_start_query_execution=callback_results.append, + result_reuse_enable=False, + ) + result = cursor.execute("SELECT 1", options=options) + assert callback_results == [cursor.query_id] + assert result is cursor + assert cursor.fetchone() == (1,) + + def test_execute_kwargs_override_options(self, cursor): + """Test that individual execute() kwargs take precedence over options fields.""" + from_options = [] + from_kwargs = [] + options = ExecuteOptions(on_start_query_execution=from_options.append) + cursor.execute( + "SELECT 1", + options=options, + on_start_query_execution=from_kwargs.append, + ) + assert from_options == [] + assert from_kwargs == [cursor.query_id] + assert cursor.fetchone() == (1,) + def test_connection_level_callback(self): """Test connection-level default callback.""" callback_results = [] diff --git a/tests/pyathena/test_options.py b/tests/pyathena/test_options.py new file mode 100644 index 00000000..df5c096a --- /dev/null +++ b/tests/pyathena/test_options.py @@ -0,0 +1,100 @@ +import dataclasses + +import pytest + +from pyathena.options import ExecuteOptions + + +def test_execute_options_defaults(): + options = ExecuteOptions() + assert options.work_group is None + assert options.s3_staging_dir is None + assert options.cache_size == 0 + assert options.cache_expiration_time == 0 + assert options.result_reuse_enable is None + assert options.result_reuse_minutes is None + assert options.paramstyle is None + assert options.on_start_query_execution is None + assert options.result_set_type_hints is None + + +def test_execute_options_is_frozen(): + options = ExecuteOptions() + with pytest.raises(dataclasses.FrozenInstanceError): + options.work_group = "primary" # type: ignore[misc] + + +def test_merge_ignores_none_overrides(): + options = ExecuteOptions(work_group="primary", cache_size=100) + merged = options.merge(work_group=None, cache_size=None, s3_staging_dir=None) + assert merged.work_group == "primary" + assert merged.cache_size == 100 + assert merged.s3_staging_dir is None + + +def test_merge_overrides_take_precedence(): + options = ExecuteOptions( + work_group="primary", + cache_size=100, + result_reuse_enable=True, + paramstyle="pyformat", + ) + merged = options.merge( + work_group="adhoc", + cache_size=0, + result_reuse_enable=False, + paramstyle="qmark", + ) + assert merged.work_group == "adhoc" + assert merged.cache_size == 0 + assert merged.result_reuse_enable is False + assert merged.paramstyle == "qmark" + + +def test_merge_does_not_mutate_original(): + options = ExecuteOptions(work_group="primary") + merged = options.merge(work_group="adhoc") + assert options.work_group == "primary" + assert merged is not options + + +def test_merge_without_applied_overrides_returns_self(): + options = ExecuteOptions(work_group="primary") + assert options.merge() is options + assert options.merge(work_group=None) is options + + +def test_merge_raises_on_unknown_field(): + with pytest.raises(TypeError, match="unknown_field"): + ExecuteOptions().merge(unknown_field="value") + + +@pytest.mark.parametrize( + ("base", "overrides", "expected"), + [ + # Overrides fill in unset fields + ( + {}, + {"work_group": "wg", "result_reuse_minutes": 5}, + {"work_group": "wg", "result_reuse_minutes": 5}, + ), + # False is a real value and must override + ( + {"result_reuse_enable": True}, + {"result_reuse_enable": False}, + {"result_reuse_enable": False}, + ), + # 0 is a real value and must override + ({"cache_size": 100}, {"cache_size": 0}, {"cache_size": 0}), + # Callbacks and hints pass through + ( + {"result_set_type_hints": {"tags": "array(varchar)"}}, + {"result_set_type_hints": {"tags": "map(varchar, integer)"}}, + {"result_set_type_hints": {"tags": "map(varchar, integer)"}}, + ), + ], +) +def test_merge_precedence(base, overrides, expected): + merged = ExecuteOptions(**base).merge(**overrides) + for name, value in expected.items(): + assert getattr(merged, name) == value From df90095180d425af1a957457d930ce22880976aa Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 6 Jul 2026 23:37:02 +0900 Subject: [PATCH 2/4] Address self-review findings - Simplify ExecuteOptions.merge() to rely on dataclasses.replace() for unknown-field validation and remove the unused logging boilerplate - Type cache_size/cache_expiration_time as plain int with default 0 and drop the duplicated None coercion in both _execute implementations - Add ExecuteOptions.resolve() and use it in all execute() methods instead of repeating the default-construction-and-merge expression - Invoke on_start_query_execution callbacks in the aio cursors, which wait for query completion like the synchronous cursors, and add the explicit on_start_query_execution argument to their execute() methods - Add missing execute() docstrings to AsyncPandasCursor and AsyncArrowCursor - Clarify cache_size/cache_expiration_time interplay in the ExecuteOptions docstring, scope the usage guide section to SQL cursors, and document that None keyword arguments do not reset options fields Co-Authored-By: Claude Fable 5 --- docs/usage.md | 14 ++++++-- pyathena/aio/arrow/cursor.py | 14 +++++++- pyathena/aio/common.py | 6 ++-- pyathena/aio/cursor.py | 14 +++++++- pyathena/aio/pandas/cursor.py | 15 +++++++-- pyathena/aio/polars/cursor.py | 14 +++++++- pyathena/aio/s3fs/cursor.py | 14 +++++++- pyathena/arrow/async_cursor.py | 26 ++++++++++++++- pyathena/arrow/cursor.py | 3 +- pyathena/async_cursor.py | 3 +- pyathena/common.py | 12 +++---- pyathena/cursor.py | 3 +- pyathena/options.py | 54 +++++++++++++++++++------------ pyathena/pandas/async_cursor.py | 29 ++++++++++++++++- pyathena/pandas/cursor.py | 3 +- pyathena/polars/async_cursor.py | 3 +- pyathena/polars/cursor.py | 3 +- pyathena/s3fs/async_cursor.py | 3 +- pyathena/s3fs/cursor.py | 3 +- tests/pyathena/aio/test_cursor.py | 14 ++++++++ tests/pyathena/test_options.py | 17 ++++++++-- 21 files changed, 213 insertions(+), 54 deletions(-) diff --git a/docs/usage.md b/docs/usage.md index 59e04a82..d71e0004 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -119,9 +119,11 @@ You can find more information about the [considerations and limitations of param ## Execution options -The `execute()` method of every cursor accepts the same set of shared keyword arguments, such as +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. @@ -147,6 +149,10 @@ cursor.execute("SELECT * FROM one_row", 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. @@ -154,8 +160,9 @@ method, which returns a new instance with the specified fields applied. adhoc_options = options.merge(work_group="ANOTHER_WORK_GROUP") ``` -The `on_start_query_execution` field is only invoked by synchronous cursors; asynchronous cursors -return the query ID directly through their execution model. +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 @@ -463,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. diff --git a/pyathena/aio/arrow/cursor.py b/pyathena/aio/arrow/cursor.py index cc46a3eb..76276843 100644 --- a/pyathena/aio/arrow/cursor.py +++ b/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 @@ -89,6 +90,7 @@ async def execute( # type: ignore[override] 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, @@ -106,6 +108,7 @@ 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. @@ -118,7 +121,8 @@ async def execute( # type: ignore[override] Self reference for method chaining. """ self._reset_state() - options = (options if options is not None else ExecuteOptions()).merge( + options = ExecuteOptions.resolve( + options, work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -126,6 +130,7 @@ async def execute( # type: ignore[override] 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) @@ -135,6 +140,13 @@ async def execute( # type: ignore[override] options=options, ) + # Call user callbacks immediately after start_query_execution + # Both connection-level and execute-level callbacks are invoked if set + if self._on_start_query_execution: + self._on_start_query_execution(self.query_id) + if options.on_start_query_execution: + options.on_start_query_execution(self.query_id) + query_execution = await self._poll(self.query_id) if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED: self.result_set = await asyncio.to_thread( diff --git a/pyathena/aio/common.py b/pyathena/aio/common.py index f971a641..9c5c3e43 100644 --- a/pyathena/aio/common.py +++ b/pyathena/aio/common.py @@ -44,10 +44,8 @@ async def _execute( # type: ignore[override] query_id = await self._find_previous_query_id( query, options.work_group, - cache_size=options.cache_size if options.cache_size else 0, - cache_expiration_time=options.cache_expiration_time - if options.cache_expiration_time - else 0, + cache_size=options.cache_size, + cache_expiration_time=options.cache_expiration_time, ) if query_id is None: try: diff --git a/pyathena/aio/cursor.py b/pyathena/aio/cursor.py index 2e3fe4ca..27b3301d 100644 --- a/pyathena/aio/cursor.py +++ b/pyathena/aio/cursor.py @@ -1,6 +1,7 @@ from __future__ import annotations import logging +from collections.abc import Callable from typing import Any, cast from pyathena.aio.common import WithAsyncFetch @@ -80,6 +81,7 @@ async def execute( # type: ignore[override] 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, @@ -97,6 +99,7 @@ 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. @@ -109,7 +112,8 @@ async def execute( # type: ignore[override] Self reference for method chaining. """ self._reset_state() - options = (options if options is not None else ExecuteOptions()).merge( + options = ExecuteOptions.resolve( + options, work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -117,6 +121,7 @@ async def execute( # type: ignore[override] 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( @@ -125,6 +130,13 @@ async def execute( # type: ignore[override] options=options, ) + # Call user callbacks immediately after start_query_execution + # Both connection-level and execute-level callbacks are invoked if set + if self._on_start_query_execution: + self._on_start_query_execution(self.query_id) + if options.on_start_query_execution: + options.on_start_query_execution(self.query_id) + query_execution = await self._poll(self.query_id) if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED: self.result_set = await self._result_set_class.create( diff --git a/pyathena/aio/pandas/cursor.py b/pyathena/aio/pandas/cursor.py index f0eeb16d..b02f729e 100644 --- a/pyathena/aio/pandas/cursor.py +++ b/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, @@ -106,6 +106,7 @@ async def execute( # type: ignore[override] 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, @@ -126,6 +127,7 @@ 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. @@ -138,7 +140,8 @@ async def execute( # type: ignore[override] Self reference for method chaining. """ self._reset_state() - options = (options if options is not None else ExecuteOptions()).merge( + options = ExecuteOptions.resolve( + options, work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -146,6 +149,7 @@ async def execute( # type: ignore[override] 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) @@ -155,6 +159,13 @@ async def execute( # type: ignore[override] options=options, ) + # Call user callbacks immediately after start_query_execution + # Both connection-level and execute-level callbacks are invoked if set + if self._on_start_query_execution: + self._on_start_query_execution(self.query_id) + if options.on_start_query_execution: + options.on_start_query_execution(self.query_id) + query_execution = await self._poll(self.query_id) if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED: self.result_set = await asyncio.to_thread( diff --git a/pyathena/aio/polars/cursor.py b/pyathena/aio/polars/cursor.py index 82227822..9c17becf 100644 --- a/pyathena/aio/polars/cursor.py +++ b/pyathena/aio/polars/cursor.py @@ -2,6 +2,7 @@ import asyncio import logging +from collections.abc import Callable from multiprocessing import cpu_count from typing import TYPE_CHECKING, Any, cast @@ -95,6 +96,7 @@ async def execute( # type: ignore[override] 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, @@ -112,6 +114,7 @@ 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. @@ -124,7 +127,8 @@ async def execute( # type: ignore[override] Self reference for method chaining. """ self._reset_state() - options = (options if options is not None else ExecuteOptions()).merge( + options = ExecuteOptions.resolve( + options, work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -132,6 +136,7 @@ async def execute( # type: ignore[override] 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) @@ -141,6 +146,13 @@ async def execute( # type: ignore[override] options=options, ) + # Call user callbacks immediately after start_query_execution + # Both connection-level and execute-level callbacks are invoked if set + if self._on_start_query_execution: + self._on_start_query_execution(self.query_id) + if options.on_start_query_execution: + options.on_start_query_execution(self.query_id) + query_execution = await self._poll(self.query_id) if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED: self.result_set = await asyncio.to_thread( diff --git a/pyathena/aio/s3fs/cursor.py b/pyathena/aio/s3fs/cursor.py index 316a20a5..81ec028b 100644 --- a/pyathena/aio/s3fs/cursor.py +++ b/pyathena/aio/s3fs/cursor.py @@ -2,6 +2,7 @@ import asyncio import logging +from collections.abc import Callable from typing import Any, cast from pyathena.aio.common import WithAsyncFetch @@ -87,6 +88,7 @@ async def execute( # type: ignore[override] 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, @@ -104,6 +106,7 @@ 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. @@ -116,7 +119,8 @@ async def execute( # type: ignore[override] Self reference for method chaining. """ self._reset_state() - options = (options if options is not None else ExecuteOptions()).merge( + options = ExecuteOptions.resolve( + options, work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -124,6 +128,7 @@ async def execute( # type: ignore[override] 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( @@ -132,6 +137,13 @@ async def execute( # type: ignore[override] options=options, ) + # Call user callbacks immediately after start_query_execution + # Both connection-level and execute-level callbacks are invoked if set + if self._on_start_query_execution: + self._on_start_query_execution(self.query_id) + if options.on_start_query_execution: + options.on_start_query_execution(self.query_id) + query_execution = await self._poll(self.query_id) if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED: self.result_set = await asyncio.to_thread( diff --git a/pyathena/arrow/async_cursor.py b/pyathena/arrow/async_cursor.py index 386da908..3deeac12 100644 --- a/pyathena/arrow/async_cursor.py +++ b/pyathena/arrow/async_cursor.py @@ -187,7 +187,31 @@ def execute( options: ExecuteOptions | None = None, **kwargs, ) -> tuple[str, Future[AthenaArrowResultSet | Any]]: - options = (options if options is not None else ExecuteOptions()).merge( + """Execute a SQL query asynchronously and return results as Arrow Tables. + + Args: + operation: SQL query string to execute. + parameters: Query parameters for parameterized queries. + work_group: Athena workgroup to use for this query. + s3_staging_dir: S3 location for query results. + cache_size: Number of queries to check for result caching. + cache_expiration_time: Cache expiration time in seconds. + result_reuse_enable: Enable Athena result reuse for this query. + result_reuse_minutes: Minutes to reuse cached results. + paramstyle: Parameter style ('qmark' or 'pyformat'). + 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: + Tuple of (query_id, future) where future resolves to AthenaArrowResultSet. + """ + options = ExecuteOptions.resolve( + options, work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, diff --git a/pyathena/arrow/cursor.py b/pyathena/arrow/cursor.py index 9c566684..0546d7c9 100644 --- a/pyathena/arrow/cursor.py +++ b/pyathena/arrow/cursor.py @@ -173,7 +173,8 @@ def execute( >>> table = cursor.as_arrow() # Returns Apache Arrow Table """ self._reset_state() - options = (options if options is not None else ExecuteOptions()).merge( + options = ExecuteOptions.resolve( + options, work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, diff --git a/pyathena/async_cursor.py b/pyathena/async_cursor.py index 7ede4da6..c761e299 100644 --- a/pyathena/async_cursor.py +++ b/pyathena/async_cursor.py @@ -211,7 +211,8 @@ def execute( >>> # Do other work while query runs... >>> result_set = future.result() # Wait for completion """ - options = (options if options is not None else ExecuteOptions()).merge( + options = ExecuteOptions.resolve( + options, work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, diff --git a/pyathena/common.py b/pyathena/common.py index 50d78656..c5a96891 100644 --- a/pyathena/common.py +++ b/pyathena/common.py @@ -193,9 +193,9 @@ def __init__( 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. + # ``on_start_query_execution`` is invoked by cursors whose ``execute()`` + # supports it (the synchronous and aio cursors). Async/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 @@ -718,10 +718,8 @@ def _execute( query_id = self._find_previous_query_id( query, options.work_group, - cache_size=options.cache_size if options.cache_size else 0, - cache_expiration_time=options.cache_expiration_time - if options.cache_expiration_time - else 0, + cache_size=options.cache_size, + cache_expiration_time=options.cache_expiration_time, ) if query_id is None: try: diff --git a/pyathena/cursor.py b/pyathena/cursor.py index b279d83a..4ff883fd 100644 --- a/pyathena/cursor.py +++ b/pyathena/cursor.py @@ -131,7 +131,8 @@ def execute( ... ) """ self._reset_state() - options = (options if options is not None else ExecuteOptions()).merge( + options = ExecuteOptions.resolve( + options, work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, diff --git a/pyathena/options.py b/pyathena/options.py index 066a2cd7..bb058201 100644 --- a/pyathena/options.py +++ b/pyathena/options.py @@ -1,19 +1,16 @@ from __future__ import annotations -import logging from collections.abc import Callable -from dataclasses import dataclass, fields, replace +from dataclasses import dataclass, replace from typing import Any -_logger = logging.getLogger(__name__) - @dataclass(frozen=True) class ExecuteOptions: """Shared options for ``Cursor.execute()`` across all cursor implementations. This dataclass is the single source of truth for the query-execution - arguments shared by every cursor type (sync/async/aio and + arguments shared by every SQL cursor type (sync/async/aio and pandas/arrow/polars/s3fs variants). It can be passed to ``execute()`` via the ``options`` keyword argument as an alternative to individual keyword arguments: @@ -28,13 +25,19 @@ class ExecuteOptions: >>> cursor.execute("SELECT ...", options=options, work_group="adhoc") + Passing None for an individual keyword argument is treated as "not + provided" and leaves the corresponding ``options`` field unchanged; to + reset a field, use :meth:`merge` or construct a new instance. + Attributes: work_group: Athena workgroup to use for this query. Overrides the connection-level workgroup. s3_staging_dir: S3 location for query results. Overrides the connection-level staging directory. cache_size: Number of recent queries to scan for client-side result - caching. 0 (default) disables the cache lookup. + caching. 0 (default) disables the cache lookup, unless + ``cache_expiration_time`` is set to a positive value, in which + case all queries within the expiration window are scanned. cache_expiration_time: Maximum age in seconds of a cached query result to consider for reuse. 0 (default) means no age limit. result_reuse_enable: Enable Athena server-side result reuse for this @@ -45,9 +48,10 @@ class ExecuteOptions: paramstyle: Parameter style for this query ('qmark' or 'pyformat'). None (default) uses the module-level ``pyathena.paramstyle``. on_start_query_execution: Callback invoked with the query ID - immediately after the StartQueryExecution API call. Only invoked - by synchronous cursors; asynchronous cursors return the query ID - directly through their execution model. + immediately after the StartQueryExecution API call. Invoked by + synchronous and aio cursors; ``AsyncCursor``-based cursors + return the query ID directly through their execution model and + do not invoke it. result_set_type_hints: Mapping of column names (or indices) to Athena DDL type signatures for precise type conversion within complex types. For example: @@ -56,14 +60,31 @@ class ExecuteOptions: 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 = 0 + cache_expiration_time: int = 0 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 + @classmethod + def resolve(cls, options: ExecuteOptions | None, **overrides: Any) -> ExecuteOptions: + """Return ``options`` (or a default instance) with ``overrides`` applied. + + This is the canonical way for ``execute()`` implementations to combine + the ``options`` argument with the individual keyword arguments. + + Args: + options: Base options, or None to start from the defaults. + **overrides: Field values to apply on top of ``options``. + None values are ignored. + + Returns: + The effective ``ExecuteOptions`` for the call. + """ + return (options if options is not None else cls()).merge(**overrides) + def merge(self, **overrides: Any) -> ExecuteOptions: """Return a new instance with non-None ``overrides`` applied. @@ -78,13 +99,4 @@ def merge(self, **overrides: Any) -> ExecuteOptions: Raises: TypeError: If an override name is not a field of this class. """ - valid = {f.name for f in fields(self)} - unknown = set(overrides) - valid - if unknown: - raise TypeError( - f"Unknown {self.__class__.__name__} fields: {', '.join(sorted(unknown))}" - ) - applied = {k: v for k, v in overrides.items() if v is not None} - if not applied: - return self - return replace(self, **applied) + return replace(self, **{k: v for k, v in overrides.items() if v is not None}) diff --git a/pyathena/pandas/async_cursor.py b/pyathena/pandas/async_cursor.py index b7bcb535..4cd443b9 100644 --- a/pyathena/pandas/async_cursor.py +++ b/pyathena/pandas/async_cursor.py @@ -165,7 +165,34 @@ def execute( options: ExecuteOptions | None = None, **kwargs, ) -> tuple[str, Future[AthenaPandasResultSet | Any]]: - options = (options if options is not None else ExecuteOptions()).merge( + """Execute a SQL query asynchronously and return results as pandas DataFrames. + + Args: + operation: SQL query string to execute. + parameters: Query parameters for parameterized queries. + work_group: Athena workgroup to use for this query. + s3_staging_dir: S3 location for query results. + cache_size: Number of queries to check for result caching. + cache_expiration_time: Cache expiration time in seconds. + result_reuse_enable: Enable Athena result reuse for this query. + result_reuse_minutes: Minutes to reuse cached results. + paramstyle: Parameter style ('qmark' or 'pyformat'). + result_set_type_hints: Optional dictionary mapping column names to + Athena DDL type signatures for precise type conversion within + complex types. + 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). + 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: + Tuple of (query_id, future) where future resolves to AthenaPandasResultSet. + """ + options = ExecuteOptions.resolve( + options, work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, diff --git a/pyathena/pandas/cursor.py b/pyathena/pandas/cursor.py index 89a73c16..92dd81f1 100644 --- a/pyathena/pandas/cursor.py +++ b/pyathena/pandas/cursor.py @@ -193,7 +193,8 @@ def execute( >>> df = cursor.fetchall() # Returns pandas DataFrame """ self._reset_state() - options = (options if options is not None else ExecuteOptions()).merge( + options = ExecuteOptions.resolve( + options, work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, diff --git a/pyathena/polars/async_cursor.py b/pyathena/polars/async_cursor.py index 5c37857c..cab28836 100644 --- a/pyathena/polars/async_cursor.py +++ b/pyathena/polars/async_cursor.py @@ -232,7 +232,8 @@ def execute( >>> result_set = future.result() >>> df = result_set.as_polars() # Returns Polars DataFrame """ - options = (options if options is not None else ExecuteOptions()).merge( + options = ExecuteOptions.resolve( + options, work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, diff --git a/pyathena/polars/cursor.py b/pyathena/polars/cursor.py index 632fb44a..838761e9 100644 --- a/pyathena/polars/cursor.py +++ b/pyathena/polars/cursor.py @@ -192,7 +192,8 @@ def execute( >>> df = cursor.as_polars() # Returns Polars DataFrame """ self._reset_state() - options = (options if options is not None else ExecuteOptions()).merge( + options = ExecuteOptions.resolve( + options, work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, diff --git a/pyathena/s3fs/async_cursor.py b/pyathena/s3fs/async_cursor.py index 79f4defd..8fd60065 100644 --- a/pyathena/s3fs/async_cursor.py +++ b/pyathena/s3fs/async_cursor.py @@ -219,7 +219,8 @@ def execute( >>> result_set = future.result() >>> rows = result_set.fetchall() """ - options = (options if options is not None else ExecuteOptions()).merge( + options = ExecuteOptions.resolve( + options, work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, diff --git a/pyathena/s3fs/cursor.py b/pyathena/s3fs/cursor.py index c5525baf..9b508728 100644 --- a/pyathena/s3fs/cursor.py +++ b/pyathena/s3fs/cursor.py @@ -168,7 +168,8 @@ def execute( >>> rows = cursor.fetchall() """ self._reset_state() - options = (options if options is not None else ExecuteOptions()).merge( + options = ExecuteOptions.resolve( + options, work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, diff --git a/tests/pyathena/aio/test_cursor.py b/tests/pyathena/aio/test_cursor.py index e8ea6cf7..87bd7cf9 100644 --- a/tests/pyathena/aio/test_cursor.py +++ b/tests/pyathena/aio/test_cursor.py @@ -3,6 +3,7 @@ import pytest +from pyathena import ExecuteOptions from pyathena.error import DatabaseError, ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.result_set import AthenaResultSet @@ -65,6 +66,19 @@ async def test_execute_returns_self(self, aio_cursor): result = await aio_cursor.execute("SELECT * FROM one_row") assert result is aio_cursor + async def test_execute_with_callback(self, aio_cursor): + callback_results = [] + await aio_cursor.execute("SELECT 1", on_start_query_execution=callback_results.append) + assert callback_results == [aio_cursor.query_id] + assert await aio_cursor.fetchone() == (1,) + + async def test_execute_with_options(self, aio_cursor): + callback_results = [] + options = ExecuteOptions(on_start_query_execution=callback_results.append) + await aio_cursor.execute("SELECT 1", options=options) + assert callback_results == [aio_cursor.query_id] + assert await aio_cursor.fetchone() == (1,) + async def test_no_result_set_raises(self, aio_cursor): with pytest.raises(ProgrammingError): await aio_cursor.fetchone() diff --git a/tests/pyathena/test_options.py b/tests/pyathena/test_options.py index df5c096a..d929fcfb 100644 --- a/tests/pyathena/test_options.py +++ b/tests/pyathena/test_options.py @@ -58,10 +58,21 @@ def test_merge_does_not_mutate_original(): assert merged is not options -def test_merge_without_applied_overrides_returns_self(): +def test_merge_without_applied_overrides_is_equal(): options = ExecuteOptions(work_group="primary") - assert options.merge() is options - assert options.merge(work_group=None) is options + assert options.merge() == options + assert options.merge(work_group=None) == options + + +def test_resolve_with_none_returns_defaults_with_overrides(): + resolved = ExecuteOptions.resolve(None, work_group="wg") + assert resolved == ExecuteOptions(work_group="wg") + + +def test_resolve_applies_overrides_to_given_options(): + options = ExecuteOptions(work_group="primary", cache_size=100) + resolved = ExecuteOptions.resolve(options, work_group="adhoc", cache_size=None) + assert resolved == ExecuteOptions(work_group="adhoc", cache_size=100) def test_merge_raises_on_unknown_field(): From a39cb92fbdc643b46ccc2c142421d36d27428715 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Mon, 6 Jul 2026 23:45:38 +0900 Subject: [PATCH 3/4] Deduplicate callback invocation and options fallback after simplify review - Extract the connection-level and execute-level on_start_query_execution invocation, previously repeated in 10 cursor execute() methods, into BaseCursor._call_on_start_query_execution - Use ExecuteOptions.resolve() for the None fallback in both _execute implementations instead of an inline conditional Co-Authored-By: Claude Fable 5 --- pyathena/aio/arrow/cursor.py | 6 +----- pyathena/aio/common.py | 2 +- pyathena/aio/cursor.py | 6 +----- pyathena/aio/pandas/cursor.py | 6 +----- pyathena/aio/polars/cursor.py | 6 +----- pyathena/aio/s3fs/cursor.py | 6 +----- pyathena/arrow/cursor.py | 6 +----- pyathena/common.py | 14 +++++++++++++- pyathena/cursor.py | 6 +----- pyathena/pandas/cursor.py | 6 +----- pyathena/polars/cursor.py | 6 +----- pyathena/s3fs/cursor.py | 5 +---- 12 files changed, 24 insertions(+), 51 deletions(-) diff --git a/pyathena/aio/arrow/cursor.py b/pyathena/aio/arrow/cursor.py index 76276843..d5af61fb 100644 --- a/pyathena/aio/arrow/cursor.py +++ b/pyathena/aio/arrow/cursor.py @@ -141,11 +141,7 @@ async def execute( # type: ignore[override] ) # Call user callbacks immediately after start_query_execution - # Both connection-level and execute-level callbacks are invoked if set - if self._on_start_query_execution: - self._on_start_query_execution(self.query_id) - if options.on_start_query_execution: - options.on_start_query_execution(self.query_id) + 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: diff --git a/pyathena/aio/common.py b/pyathena/aio/common.py index 9c5c3e43..7cfd7c4b 100644 --- a/pyathena/aio/common.py +++ b/pyathena/aio/common.py @@ -30,7 +30,7 @@ async def _execute( # type: ignore[override] parameters: dict[str, Any] | list[str] | None = None, options: ExecuteOptions | None = None, ) -> str: - options = options if options is not None else ExecuteOptions() + options = ExecuteOptions.resolve(options) query, execution_parameters = self._prepare_query(operation, parameters, options.paramstyle) request = self._build_start_query_execution_request( diff --git a/pyathena/aio/cursor.py b/pyathena/aio/cursor.py index 27b3301d..70ef3aea 100644 --- a/pyathena/aio/cursor.py +++ b/pyathena/aio/cursor.py @@ -131,11 +131,7 @@ async def execute( # type: ignore[override] ) # Call user callbacks immediately after start_query_execution - # Both connection-level and execute-level callbacks are invoked if set - if self._on_start_query_execution: - self._on_start_query_execution(self.query_id) - if options.on_start_query_execution: - options.on_start_query_execution(self.query_id) + 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: diff --git a/pyathena/aio/pandas/cursor.py b/pyathena/aio/pandas/cursor.py index b02f729e..deca9cff 100644 --- a/pyathena/aio/pandas/cursor.py +++ b/pyathena/aio/pandas/cursor.py @@ -160,11 +160,7 @@ async def execute( # type: ignore[override] ) # Call user callbacks immediately after start_query_execution - # Both connection-level and execute-level callbacks are invoked if set - if self._on_start_query_execution: - self._on_start_query_execution(self.query_id) - if options.on_start_query_execution: - options.on_start_query_execution(self.query_id) + 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: diff --git a/pyathena/aio/polars/cursor.py b/pyathena/aio/polars/cursor.py index 9c17becf..fcc1049a 100644 --- a/pyathena/aio/polars/cursor.py +++ b/pyathena/aio/polars/cursor.py @@ -147,11 +147,7 @@ async def execute( # type: ignore[override] ) # Call user callbacks immediately after start_query_execution - # Both connection-level and execute-level callbacks are invoked if set - if self._on_start_query_execution: - self._on_start_query_execution(self.query_id) - if options.on_start_query_execution: - options.on_start_query_execution(self.query_id) + 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: diff --git a/pyathena/aio/s3fs/cursor.py b/pyathena/aio/s3fs/cursor.py index 81ec028b..3c58914e 100644 --- a/pyathena/aio/s3fs/cursor.py +++ b/pyathena/aio/s3fs/cursor.py @@ -138,11 +138,7 @@ async def execute( # type: ignore[override] ) # Call user callbacks immediately after start_query_execution - # Both connection-level and execute-level callbacks are invoked if set - if self._on_start_query_execution: - self._on_start_query_execution(self.query_id) - if options.on_start_query_execution: - options.on_start_query_execution(self.query_id) + 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: diff --git a/pyathena/arrow/cursor.py b/pyathena/arrow/cursor.py index 0546d7c9..5ad7cd60 100644 --- a/pyathena/arrow/cursor.py +++ b/pyathena/arrow/cursor.py @@ -193,11 +193,7 @@ def execute( ) # Call user callbacks immediately after start_query_execution - # Both connection-level and execute-level callbacks are invoked if set - if self._on_start_query_execution: - self._on_start_query_execution(self.query_id) - if options.on_start_query_execution: - options.on_start_query_execution(self.query_id) + self._call_on_start_query_execution(self.query_id, options) query_execution = cast(AthenaQueryExecution, self._poll(self.query_id)) if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED: self.result_set = AthenaArrowResultSet( diff --git a/pyathena/common.py b/pyathena/common.py index c5a96891..4ae3836f 100644 --- a/pyathena/common.py +++ b/pyathena/common.py @@ -698,13 +698,25 @@ def _prepare_unload( compression=AthenaCompression.COMPRESSION_SNAPPY, ) + def _call_on_start_query_execution(self, query_id: str, options: ExecuteOptions) -> None: + """Invoke the connection-level and execute-level query-start callbacks. + + Both callbacks are invoked if set. Called by cursors whose execution + model supports early access to the query ID (the synchronous and aio + cursors) immediately after the StartQueryExecution API call. + """ + if self._on_start_query_execution: + self._on_start_query_execution(query_id) + if options.on_start_query_execution: + options.on_start_query_execution(query_id) + def _execute( self, operation: str, parameters: dict[str, Any] | list[str] | None = None, options: ExecuteOptions | None = None, ) -> str: - options = options if options is not None else ExecuteOptions() + options = ExecuteOptions.resolve(options) query, execution_parameters = self._prepare_query(operation, parameters, options.paramstyle) request = self._build_start_query_execution_request( diff --git a/pyathena/cursor.py b/pyathena/cursor.py index 4ff883fd..eb302b30 100644 --- a/pyathena/cursor.py +++ b/pyathena/cursor.py @@ -150,11 +150,7 @@ def execute( ) # Call user callbacks immediately after start_query_execution - # Both connection-level and execute-level callbacks are invoked if set - if self._on_start_query_execution: - self._on_start_query_execution(self.query_id) - if options.on_start_query_execution: - options.on_start_query_execution(self.query_id) + self._call_on_start_query_execution(self.query_id, options) query_execution = cast(AthenaQueryExecution, self._poll(self.query_id)) if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED: diff --git a/pyathena/pandas/cursor.py b/pyathena/pandas/cursor.py index 92dd81f1..edb0f356 100644 --- a/pyathena/pandas/cursor.py +++ b/pyathena/pandas/cursor.py @@ -213,11 +213,7 @@ def execute( ) # Call user callbacks immediately after start_query_execution - # Both connection-level and execute-level callbacks are invoked if set - if self._on_start_query_execution: - self._on_start_query_execution(self.query_id) - if options.on_start_query_execution: - options.on_start_query_execution(self.query_id) + self._call_on_start_query_execution(self.query_id, options) query_execution = cast(AthenaQueryExecution, self._poll(self.query_id)) if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED: self.result_set = AthenaPandasResultSet( diff --git a/pyathena/polars/cursor.py b/pyathena/polars/cursor.py index 838761e9..a92ed580 100644 --- a/pyathena/polars/cursor.py +++ b/pyathena/polars/cursor.py @@ -212,11 +212,7 @@ def execute( ) # Call user callbacks immediately after start_query_execution - # Both connection-level and execute-level callbacks are invoked if set - if self._on_start_query_execution: - self._on_start_query_execution(self.query_id) - if options.on_start_query_execution: - options.on_start_query_execution(self.query_id) + self._call_on_start_query_execution(self.query_id, options) query_execution = cast(AthenaQueryExecution, self._poll(self.query_id)) if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED: self.result_set = AthenaPolarsResultSet( diff --git a/pyathena/s3fs/cursor.py b/pyathena/s3fs/cursor.py index 9b508728..6052b66e 100644 --- a/pyathena/s3fs/cursor.py +++ b/pyathena/s3fs/cursor.py @@ -187,10 +187,7 @@ def execute( ) # Call user callbacks immediately after start_query_execution - if self._on_start_query_execution: - self._on_start_query_execution(self.query_id) - if options.on_start_query_execution: - options.on_start_query_execution(self.query_id) + self._call_on_start_query_execution(self.query_id, options) query_execution = cast(AthenaQueryExecution, self._poll(self.query_id)) if query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED: From 1ad8a1f32d583036ed8205591ec94479fac82575 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Tue, 7 Jul 2026 09:54:02 +0900 Subject: [PATCH 4/4] Group ExecuteOptions unit tests into a TestExecuteOptions class Co-Authored-By: Claude Fable 5 --- tests/pyathena/test_options.py | 198 ++++++++++++++++----------------- 1 file changed, 95 insertions(+), 103 deletions(-) diff --git a/tests/pyathena/test_options.py b/tests/pyathena/test_options.py index d929fcfb..e3ae791f 100644 --- a/tests/pyathena/test_options.py +++ b/tests/pyathena/test_options.py @@ -5,107 +5,99 @@ from pyathena.options import ExecuteOptions -def test_execute_options_defaults(): - options = ExecuteOptions() - assert options.work_group is None - assert options.s3_staging_dir is None - assert options.cache_size == 0 - assert options.cache_expiration_time == 0 - assert options.result_reuse_enable is None - assert options.result_reuse_minutes is None - assert options.paramstyle is None - assert options.on_start_query_execution is None - assert options.result_set_type_hints is None - - -def test_execute_options_is_frozen(): - options = ExecuteOptions() - with pytest.raises(dataclasses.FrozenInstanceError): - options.work_group = "primary" # type: ignore[misc] - - -def test_merge_ignores_none_overrides(): - options = ExecuteOptions(work_group="primary", cache_size=100) - merged = options.merge(work_group=None, cache_size=None, s3_staging_dir=None) - assert merged.work_group == "primary" - assert merged.cache_size == 100 - assert merged.s3_staging_dir is None - - -def test_merge_overrides_take_precedence(): - options = ExecuteOptions( - work_group="primary", - cache_size=100, - result_reuse_enable=True, - paramstyle="pyformat", - ) - merged = options.merge( - work_group="adhoc", - cache_size=0, - result_reuse_enable=False, - paramstyle="qmark", +class TestExecuteOptions: + def test_defaults(self): + options = ExecuteOptions() + assert options.work_group is None + assert options.s3_staging_dir is None + assert options.cache_size == 0 + assert options.cache_expiration_time == 0 + assert options.result_reuse_enable is None + assert options.result_reuse_minutes is None + assert options.paramstyle is None + assert options.on_start_query_execution is None + assert options.result_set_type_hints is None + + def test_is_frozen(self): + options = ExecuteOptions() + with pytest.raises(dataclasses.FrozenInstanceError): + options.work_group = "primary" # type: ignore[misc] + + def test_merge_ignores_none_overrides(self): + options = ExecuteOptions(work_group="primary", cache_size=100) + merged = options.merge(work_group=None, cache_size=None, s3_staging_dir=None) + assert merged.work_group == "primary" + assert merged.cache_size == 100 + assert merged.s3_staging_dir is None + + def test_merge_overrides_take_precedence(self): + options = ExecuteOptions( + work_group="primary", + cache_size=100, + result_reuse_enable=True, + paramstyle="pyformat", + ) + merged = options.merge( + work_group="adhoc", + cache_size=0, + result_reuse_enable=False, + paramstyle="qmark", + ) + assert merged.work_group == "adhoc" + assert merged.cache_size == 0 + assert merged.result_reuse_enable is False + assert merged.paramstyle == "qmark" + + def test_merge_does_not_mutate_original(self): + options = ExecuteOptions(work_group="primary") + merged = options.merge(work_group="adhoc") + assert options.work_group == "primary" + assert merged is not options + + def test_merge_without_applied_overrides_is_equal(self): + options = ExecuteOptions(work_group="primary") + assert options.merge() == options + assert options.merge(work_group=None) == options + + def test_merge_raises_on_unknown_field(self): + with pytest.raises(TypeError, match="unknown_field"): + ExecuteOptions().merge(unknown_field="value") + + @pytest.mark.parametrize( + ("base", "overrides", "expected"), + [ + # Overrides fill in unset fields + ( + {}, + {"work_group": "wg", "result_reuse_minutes": 5}, + {"work_group": "wg", "result_reuse_minutes": 5}, + ), + # False is a real value and must override + ( + {"result_reuse_enable": True}, + {"result_reuse_enable": False}, + {"result_reuse_enable": False}, + ), + # 0 is a real value and must override + ({"cache_size": 100}, {"cache_size": 0}, {"cache_size": 0}), + # Callbacks and hints pass through + ( + {"result_set_type_hints": {"tags": "array(varchar)"}}, + {"result_set_type_hints": {"tags": "map(varchar, integer)"}}, + {"result_set_type_hints": {"tags": "map(varchar, integer)"}}, + ), + ], ) - assert merged.work_group == "adhoc" - assert merged.cache_size == 0 - assert merged.result_reuse_enable is False - assert merged.paramstyle == "qmark" - - -def test_merge_does_not_mutate_original(): - options = ExecuteOptions(work_group="primary") - merged = options.merge(work_group="adhoc") - assert options.work_group == "primary" - assert merged is not options - - -def test_merge_without_applied_overrides_is_equal(): - options = ExecuteOptions(work_group="primary") - assert options.merge() == options - assert options.merge(work_group=None) == options - - -def test_resolve_with_none_returns_defaults_with_overrides(): - resolved = ExecuteOptions.resolve(None, work_group="wg") - assert resolved == ExecuteOptions(work_group="wg") - - -def test_resolve_applies_overrides_to_given_options(): - options = ExecuteOptions(work_group="primary", cache_size=100) - resolved = ExecuteOptions.resolve(options, work_group="adhoc", cache_size=None) - assert resolved == ExecuteOptions(work_group="adhoc", cache_size=100) - - -def test_merge_raises_on_unknown_field(): - with pytest.raises(TypeError, match="unknown_field"): - ExecuteOptions().merge(unknown_field="value") - - -@pytest.mark.parametrize( - ("base", "overrides", "expected"), - [ - # Overrides fill in unset fields - ( - {}, - {"work_group": "wg", "result_reuse_minutes": 5}, - {"work_group": "wg", "result_reuse_minutes": 5}, - ), - # False is a real value and must override - ( - {"result_reuse_enable": True}, - {"result_reuse_enable": False}, - {"result_reuse_enable": False}, - ), - # 0 is a real value and must override - ({"cache_size": 100}, {"cache_size": 0}, {"cache_size": 0}), - # Callbacks and hints pass through - ( - {"result_set_type_hints": {"tags": "array(varchar)"}}, - {"result_set_type_hints": {"tags": "map(varchar, integer)"}}, - {"result_set_type_hints": {"tags": "map(varchar, integer)"}}, - ), - ], -) -def test_merge_precedence(base, overrides, expected): - merged = ExecuteOptions(**base).merge(**overrides) - for name, value in expected.items(): - assert getattr(merged, name) == value + def test_merge_precedence(self, base, overrides, expected): + merged = ExecuteOptions(**base).merge(**overrides) + for name, value in expected.items(): + assert getattr(merged, name) == value + + def test_resolve_with_none_returns_defaults_with_overrides(self): + resolved = ExecuteOptions.resolve(None, work_group="wg") + assert resolved == ExecuteOptions(work_group="wg") + + def test_resolve_applies_overrides_to_given_options(self): + options = ExecuteOptions(work_group="primary", cache_size=100) + resolved = ExecuteOptions.resolve(options, work_group="adhoc", cache_size=None) + assert resolved == ExecuteOptions(work_group="adhoc", cache_size=100)