From 09f7fb325745087de044cc06491ddb261a753637 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Tue, 6 Oct 2026 01:12:53 +0900 Subject: [PATCH 01/15] Reject unknown cursor keyword arguments for PyAthena 4.0 --- docs/cursor.md | 21 +++ pyathena/_kwargs.py | 60 ++++++ pyathena/aio/arrow/cursor.py | 9 +- pyathena/aio/cursor.py | 4 +- pyathena/aio/pandas/cursor.py | 7 +- pyathena/aio/polars/cursor.py | 8 +- pyathena/aio/s3fs/cursor.py | 4 +- pyathena/aio/spark/cursor.py | 4 +- pyathena/aio/sqlalchemy/base.py | 5 + pyathena/arrow/async_cursor.py | 9 +- pyathena/arrow/cursor.py | 9 +- pyathena/async_cursor.py | 4 +- pyathena/common.py | 2 - pyathena/connection.py | 17 +- pyathena/cursor.py | 4 +- pyathena/pandas/async_cursor.py | 3 +- pyathena/pandas/cursor.py | 7 +- pyathena/pandas/result_set.py | 33 ++++ pyathena/polars/async_cursor.py | 8 +- pyathena/polars/cursor.py | 8 +- pyathena/polars/result_set.py | 28 +++ pyathena/s3fs/async_cursor.py | 9 +- pyathena/s3fs/cursor.py | 9 +- pyathena/spark/async_cursor.py | 2 + pyathena/spark/cursor.py | 2 + pyathena/sqlalchemy/base.py | 4 +- tests/pyathena/test_connection.py | 6 +- tests/pyathena/test_kwargs.py | 292 ++++++++++++++++++++++++++++++ 28 files changed, 553 insertions(+), 25 deletions(-) create mode 100644 pyathena/_kwargs.py create mode 100644 tests/pyathena/test_kwargs.py diff --git a/docs/cursor.md b/docs/cursor.md index a3db13a2a..dd3cad76d 100644 --- a/docs/cursor.md +++ b/docs/cursor.md @@ -1,5 +1,26 @@ # Cursor +## Keyword arguments in 4.0.0 + +Starting in 4.0.0, cursor constructors and `execute()` raise `TypeError` for unknown keyword arguments. +Earlier versions silently ignored some extra arguments. +Correct misspelled names and pass backend options to the cursor that supports them: + +```python +cursor = connection.cursor(work_group="analytics") +cursor.execute("SELECT 1", work_group="analytics") +``` + +The validation happens before starting a query or a Spark calculation. +For native asyncio cursors, `execute()` validates arguments when awaited. +Pandas and Polars cursors continue to accept their reader options, with names checked against the installed library and the selected CSV, UNLOAD, and chunking mode. +Invalid values and combinations remain subject to the reader's validation. +Arrow and S3FS cursors accept their supported result-set options; arbitrary reader options are rejected. + +Connection-level `on_start_query_execution` and `on_poll` remain accepted by every cursor constructor. +The Future-based Async cursors and Spark cursors do not invoke `on_start_query_execution`. +For reader options, such as pandas `parse_dates`, pass the option to `execute()` rather than the constructor. + (default_cursor)= ## DefaultCursor diff --git a/pyathena/_kwargs.py b/pyathena/_kwargs.py new file mode 100644 index 000000000..03853fa25 --- /dev/null +++ b/pyathena/_kwargs.py @@ -0,0 +1,60 @@ +# Copyright 2026 The PyAthena authors +# +# Licensed under the MIT License. +# See LICENSE or https://opensource.org/licenses/MIT. +# +# SPDX-License-Identifier: MIT + +"""Validation of keyword names forwarded by cursors.""" + +from collections.abc import Callable, Collection +from inspect import Parameter, signature +from typing import Any + + +def keyword_parameters(func: Callable[..., Any]) -> set[str]: + """Return the explicitly named parameters that accept keyword arguments. + + Args: + func: The callable whose signature to inspect. + + Returns: + Its positional-or-keyword and keyword-only parameter names. + """ + return { + name + for name, parameter in signature(func).parameters.items() + if parameter.kind in (Parameter.POSITIONAL_OR_KEYWORD, Parameter.KEYWORD_ONLY) + } + + +def constructor_keyword_parameters(cls: type[Any]) -> set[str]: + """Return explicit constructor keyword names across a class's MRO. + + Args: + cls: The class whose forwarding constructors to inspect. + + Returns: + The named constructor parameters, excluding ``self``. + """ + names: set[str] = set() + for base in cls.__mro__: + if "__init__" in vars(base): + names.update(keyword_parameters(vars(base)["__init__"])) + return names - {"self"} + + +def validate_kwargs(method: str, kwargs: dict[str, Any], allowed: Collection[str] = ()) -> None: + """Reject the first keyword name that a cursor does not support. + + Args: + method: The method name included in the error. + kwargs: Extra keyword arguments given to the method. + allowed: The extra keyword names the method supports. + + Raises: + TypeError: If a keyword name is not allowed. + """ + for name in kwargs: + if name not in allowed: + raise TypeError(f"{method}() got an unexpected keyword argument '{name}'") diff --git a/pyathena/aio/arrow/cursor.py b/pyathena/aio/arrow/cursor.py index b66778bc9..8977e20a4 100644 --- a/pyathena/aio/arrow/cursor.py +++ b/pyathena/aio/arrow/cursor.py @@ -7,6 +7,7 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Any, cast +from pyathena._kwargs import validate_kwargs from pyathena.aio.common import WithAsyncFetch from pyathena.arrow.converter import ( DefaultArrowTypeConverter, @@ -143,13 +144,19 @@ async def execute( options: Shared execution options as an :class:`~pyathena.options.ExecuteOptions` instance. Individual keyword arguments take precedence over ``options`` fields. - **kwargs: Additional execution parameters. + **kwargs: Result-set overrides: ``block_size``, ``connect_timeout``, + and ``request_timeout``. Unknown names raise TypeError. ``block_size`` sets the read block size for this query, and ``connect_timeout`` and ``request_timeout`` override the cursor's values. Returns: Self reference for method chaining. """ + validate_kwargs( + f"{type(self).__name__}.execute", + kwargs, + ("block_size", "connect_timeout", "request_timeout"), + ) self._reset_state() options = ExecuteOptions.resolve( options, diff --git a/pyathena/aio/cursor.py b/pyathena/aio/cursor.py index 037addb33..6a6e0b53e 100644 --- a/pyathena/aio/cursor.py +++ b/pyathena/aio/cursor.py @@ -13,6 +13,7 @@ from collections.abc import Callable from typing import Any, cast +from pyathena._kwargs import validate_kwargs from pyathena.aio.common import WithAsyncFetch from pyathena.aio.result_set import AthenaAioDictResultSet, AthenaAioResultSet from pyathena.common import CursorIterator @@ -138,11 +139,12 @@ async def execute( options: Shared execution options as an :class:`~pyathena.options.ExecuteOptions` instance. Individual keyword arguments take precedence over ``options`` fields. - **kwargs: Additional execution parameters. + **kwargs: Unknown keyword arguments raise TypeError. Returns: Self reference for method chaining. """ + validate_kwargs(f"{type(self).__name__}.execute", kwargs) self._reset_state() options = ExecuteOptions.resolve( options, diff --git a/pyathena/aio/pandas/cursor.py b/pyathena/aio/pandas/cursor.py index bca2824b2..dce49586c 100644 --- a/pyathena/aio/pandas/cursor.py +++ b/pyathena/aio/pandas/cursor.py @@ -21,7 +21,11 @@ DefaultPandasTypeConverter, DefaultPandasUnloadTypeConverter, ) -from pyathena.pandas.result_set import AthenaPandasResultSet, PandasDataFrameIterator +from pyathena.pandas.result_set import ( + AthenaPandasResultSet, + PandasDataFrameIterator, + validate_execute_kwargs, +) from pyathena.util import override if TYPE_CHECKING: @@ -176,6 +180,7 @@ async def execute( Returns: Self reference for method chaining. """ + validate_execute_kwargs(f"{type(self).__name__}.execute", kwargs, self._unload) self._reset_state() options = ExecuteOptions.resolve( options, diff --git a/pyathena/aio/polars/cursor.py b/pyathena/aio/polars/cursor.py index 12cd3a64e..451a80d1a 100644 --- a/pyathena/aio/polars/cursor.py +++ b/pyathena/aio/polars/cursor.py @@ -17,7 +17,7 @@ DefaultPolarsTypeConverter, DefaultPolarsUnloadTypeConverter, ) -from pyathena.polars.result_set import AthenaPolarsResultSet +from pyathena.polars.result_set import AthenaPolarsResultSet, validate_execute_kwargs from pyathena.util import override if TYPE_CHECKING: @@ -160,6 +160,12 @@ async def execute( Returns: Self reference for method chaining. """ + validate_execute_kwargs( + f"{type(self).__name__}.execute", + kwargs, + self._unload, + kwargs.get("chunksize", self._chunksize), + ) self._reset_state() options = ExecuteOptions.resolve( options, diff --git a/pyathena/aio/s3fs/cursor.py b/pyathena/aio/s3fs/cursor.py index 93ef181d6..e61fcaa24 100644 --- a/pyathena/aio/s3fs/cursor.py +++ b/pyathena/aio/s3fs/cursor.py @@ -14,6 +14,7 @@ from collections.abc import Callable from typing import Any +from pyathena._kwargs import validate_kwargs from pyathena.aio.common import WithAsyncFetch from pyathena.common import CursorIterator from pyathena.error import OperationalError @@ -146,13 +147,14 @@ async def execute( options: Shared execution options as an :class:`~pyathena.options.ExecuteOptions` instance. Individual keyword arguments take precedence over ``options`` fields. - **kwargs: Additional execution parameters. + **kwargs: Supported S3FS result-set overrides. Unknown names raise TypeError. ``block_size`` sets the read block size for this query, and ``csv_reader`` overrides the cursor's value. Returns: Self reference for method chaining. """ + validate_kwargs(f"{type(self).__name__}.execute", kwargs, ("block_size", "csv_reader")) self._reset_state() options = ExecuteOptions.resolve( options, diff --git a/pyathena/aio/spark/cursor.py b/pyathena/aio/spark/cursor.py index effeceb02..677103daf 100644 --- a/pyathena/aio/spark/cursor.py +++ b/pyathena/aio/spark/cursor.py @@ -14,6 +14,7 @@ import uuid from typing import Any, cast +from pyathena._kwargs import validate_kwargs from pyathena.aio.util import async_retry_api_call from pyathena.error import DatabaseError, NotSupportedError, OperationalError, ProgrammingError from pyathena.model import ( @@ -352,11 +353,12 @@ async def execute( description: Calculation description. client_request_token: Idempotency token. work_group: Unused, kept for API compatibility. - **kwargs: Additional parameters. + **kwargs: Unknown keyword arguments raise TypeError. Returns: Self reference for method chaining. """ + validate_kwargs(f"{type(self).__name__}.execute", kwargs) # A failure below must not leave the previous calculation on the cursor. self._calculation_id = None self._calculation_execution = None diff --git a/pyathena/aio/sqlalchemy/base.py b/pyathena/aio/sqlalchemy/base.py index abb600c53..8c114b7c4 100644 --- a/pyathena/aio/sqlalchemy/base.py +++ b/pyathena/aio/sqlalchemy/base.py @@ -276,6 +276,11 @@ def cursor(self, cursor: Any = None, **kwargs: Any) -> AsyncAdapt_pyathena_curso raw_cursor = self._connection.cursor(_ASYNC_CURSOR_CLASSES.get(cursor, cursor), **kwargs) return AsyncAdapt_pyathena_cursor(raw_cursor) + def _internal_cursor(self, cursor: Any) -> AsyncAdapt_pyathena_cursor: + """Adapt an API cursor with the default backend's result settings excluded.""" + raw_cursor = self._connection._internal_cursor(_ASYNC_CURSOR_CLASSES.get(cursor, cursor)) + return AsyncAdapt_pyathena_cursor(raw_cursor) + def close(self) -> None: """Close the wrapped connection.""" self._connection.close() diff --git a/pyathena/arrow/async_cursor.py b/pyathena/arrow/async_cursor.py index 340147bc7..9ba794f22 100644 --- a/pyathena/arrow/async_cursor.py +++ b/pyathena/arrow/async_cursor.py @@ -8,6 +8,7 @@ from typing import Any, cast from pyathena import ProgrammingError +from pyathena._kwargs import validate_kwargs from pyathena.arrow.converter import ( DefaultArrowTypeConverter, DefaultArrowUnloadTypeConverter, @@ -211,13 +212,19 @@ def execute( options: Shared execution options as an :class:`~pyathena.options.ExecuteOptions` instance. Individual keyword arguments take precedence over ``options`` fields. - **kwargs: Additional execution parameters. + **kwargs: Result-set overrides: ``block_size``, ``connect_timeout``, + and ``request_timeout``. Unknown names raise TypeError. ``block_size`` sets the read block size for this query, and ``connect_timeout`` and ``request_timeout`` override the cursor's values. Returns: Tuple of (query_id, future) where future resolves to AthenaArrowResultSet. """ + validate_kwargs( + f"{type(self).__name__}.execute", + kwargs, + ("block_size", "connect_timeout", "request_timeout"), + ) options = ExecuteOptions.resolve( options, work_group=work_group, diff --git a/pyathena/arrow/cursor.py b/pyathena/arrow/cursor.py index 081350343..de7efb4cb 100644 --- a/pyathena/arrow/cursor.py +++ b/pyathena/arrow/cursor.py @@ -6,6 +6,7 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Any, cast +from pyathena._kwargs import validate_kwargs from pyathena.arrow.converter import ( DefaultArrowTypeConverter, DefaultArrowUnloadTypeConverter, @@ -169,7 +170,8 @@ def execute( options: Shared execution options as an :class:`~pyathena.options.ExecuteOptions` instance. Individual keyword arguments take precedence over ``options`` fields. - **kwargs: Additional execution parameters. + **kwargs: Result-set overrides: ``block_size``, ``connect_timeout``, + and ``request_timeout``. Unknown names raise TypeError. ``block_size`` sets the read block size for this query, and ``connect_timeout`` and ``request_timeout`` override the cursor's values. @@ -180,6 +182,11 @@ def execute( >>> cursor.execute("SELECT * FROM sales WHERE year = 2023") >>> table = cursor.as_arrow() # Returns Apache Arrow Table """ + validate_kwargs( + f"{type(self).__name__}.execute", + kwargs, + ("block_size", "connect_timeout", "request_timeout"), + ) self._reset_state() options = ExecuteOptions.resolve( options, diff --git a/pyathena/async_cursor.py b/pyathena/async_cursor.py index 076cee37c..6bd823125 100644 --- a/pyathena/async_cursor.py +++ b/pyathena/async_cursor.py @@ -8,6 +8,7 @@ from multiprocessing import cpu_count from typing import Any, cast +from pyathena._kwargs import validate_kwargs from pyathena.common import BaseCursor, CursorIterator from pyathena.error import NotSupportedError, ProgrammingError from pyathena.model import AthenaQueryExecution @@ -240,7 +241,7 @@ def execute( options: Shared execution options as an :class:`~pyathena.options.ExecuteOptions` instance. Individual keyword arguments take precedence over ``options`` fields. - **kwargs: Additional execution parameters. + **kwargs: Unknown keyword arguments raise TypeError. Returns: Tuple of (query_id, future) where: @@ -253,6 +254,7 @@ def execute( >>> # Do other work while query runs... >>> result_set = future.result() # Wait for completion """ + validate_kwargs(f"{type(self).__name__}.execute", kwargs) options = ExecuteOptions.resolve( options, work_group=work_group, diff --git a/pyathena/common.py b/pyathena/common.py index 6473d4dae..245011c34 100644 --- a/pyathena/common.py +++ b/pyathena/common.py @@ -211,7 +211,6 @@ def __init__( result_reuse_minutes: int, on_start_query_execution: Callable[[str], None] | None = None, on_poll: OnPollCallback | None = None, - **kwargs, ) -> None: """Initialize the cursor with the settings it uses to run queries. @@ -235,7 +234,6 @@ def __init__( waits for the query, by cursors whose ``execute()`` supports it. on_poll: Callback invoked once per poll iteration with the current execution object. - **kwargs: Ignored. """ super().__init__() self._connection = connection diff --git a/pyathena/connection.py b/pyathena/connection.py index 6d52074be..aa3c08204 100644 --- a/pyathena/connection.py +++ b/pyathena/connection.py @@ -21,6 +21,7 @@ from botocore.config import Config import pyathena +from pyathena._kwargs import constructor_keyword_parameters from pyathena.common import BaseCursor, CursorIterator, OnPollCallback from pyathena.converter import Converter from pyathena.cursor import Cursor @@ -609,7 +610,12 @@ def cursor( >>> pandas_cursor = connection.cursor(PandasCursor) >>> df = pandas_cursor.execute("SELECT * FROM my_table").fetchall() """ - kwargs = {**self.cursor_kwargs, **kwargs} + return self._cursor(cursor, **{**self.cursor_kwargs, **kwargs}) + + def _cursor( + self, cursor: type[FunctionalCursor] | None = None, **kwargs: Any + ) -> FunctionalCursor | ConnectionCursor: + """Create a cursor from already resolved cursor defaults.""" _cursor = cursor or self.cursor_class converter = kwargs.pop("converter", self._converter) if not converter: @@ -636,6 +642,15 @@ def cursor( **kwargs, ) + def _internal_cursor(self, cursor: type[FunctionalCursor]) -> FunctionalCursor: + """Create an API cursor without the default backend's result settings.""" + backend_options = constructor_keyword_parameters(self.cursor_class) - ( + constructor_keyword_parameters(cursor) + ) + kwargs = {k: v for k, v in self.cursor_kwargs.items() if k not in backend_options} + kwargs["converter"] = cursor.get_default_converter() + return cast(FunctionalCursor, self._cursor(cursor, **kwargs)) + def close(self) -> None: """Close the connection. diff --git a/pyathena/cursor.py b/pyathena/cursor.py index d29ff5105..4c36e8a5e 100644 --- a/pyathena/cursor.py +++ b/pyathena/cursor.py @@ -6,6 +6,7 @@ from collections.abc import Callable from typing import Any, cast +from pyathena._kwargs import validate_kwargs from pyathena.common import CursorIterator from pyathena.error import OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution @@ -148,7 +149,7 @@ def execute( options: Shared execution options as an :class:`~pyathena.options.ExecuteOptions` instance. Individual keyword arguments take precedence over ``options`` fields. - **kwargs: Additional execution parameters. + **kwargs: Unknown keyword arguments raise TypeError. Returns: Self reference for method chaining. @@ -162,6 +163,7 @@ def execute( ... } ... ) """ + validate_kwargs(f"{type(self).__name__}.execute", kwargs) self._reset_state() options = ExecuteOptions.resolve( options, diff --git a/pyathena/pandas/async_cursor.py b/pyathena/pandas/async_cursor.py index a23cc9cfd..045631b15 100644 --- a/pyathena/pandas/async_cursor.py +++ b/pyathena/pandas/async_cursor.py @@ -17,7 +17,7 @@ DefaultPandasTypeConverter, DefaultPandasUnloadTypeConverter, ) -from pyathena.pandas.result_set import AthenaPandasResultSet +from pyathena.pandas.result_set import AthenaPandasResultSet, validate_execute_kwargs from pyathena.util import override _logger = logging.getLogger(__name__) @@ -241,6 +241,7 @@ def execute( Returns: Tuple of (query_id, future) where future resolves to AthenaPandasResultSet. """ + validate_execute_kwargs(f"{type(self).__name__}.execute", kwargs, self._unload) options = ExecuteOptions.resolve( options, work_group=work_group, diff --git a/pyathena/pandas/cursor.py b/pyathena/pandas/cursor.py index afcd12cec..e7fbd2e58 100644 --- a/pyathena/pandas/cursor.py +++ b/pyathena/pandas/cursor.py @@ -19,7 +19,11 @@ DefaultPandasTypeConverter, DefaultPandasUnloadTypeConverter, ) -from pyathena.pandas.result_set import AthenaPandasResultSet, PandasDataFrameIterator +from pyathena.pandas.result_set import ( + AthenaPandasResultSet, + PandasDataFrameIterator, + validate_execute_kwargs, +) from pyathena.result_set import WithFetch from pyathena.util import override @@ -207,6 +211,7 @@ def execute( ... {"year": 2023}) >>> df = cursor.as_pandas() # Returns pandas DataFrame """ + validate_execute_kwargs(f"{type(self).__name__}.execute", kwargs, self._unload) self._reset_state() options = ExecuteOptions.resolve( options, diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 669683d06..881ec2a68 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -8,6 +8,7 @@ from collections.abc import Callable, Iterable, Iterator from contextlib import ExitStack from functools import partial +from importlib import import_module from io import BufferedReader, BytesIO, IOBase, StringIO, TextIOWrapper from multiprocessing import cpu_count from typing import ( @@ -19,6 +20,7 @@ from fsspec import open as filesystem_open from pyathena import OperationalError +from pyathena._kwargs import keyword_parameters, validate_kwargs from pyathena.converter import Converter from pyathena.error import ProgrammingError from pyathena.model import AthenaQueryExecution @@ -35,6 +37,37 @@ _logger = logging.getLogger(__name__) +def validate_execute_kwargs(method: str, kwargs: dict[str, Any], unload: bool) -> None: + """Validate extra execution keyword names before starting a query. + + Args: + method: The cursor method name included in an error. + kwargs: Extra execution arguments, including result-set overrides. + unload: Whether the cursor reads UNLOAD Parquet results. + + Raises: + TypeError: If an argument is not supported by the selected reader. + """ + if not kwargs: + return + pandas = import_module("pandas") + if unload: + parquet = import_module("pyarrow.parquet") + allowed = keyword_parameters(pandas.read_parquet) | keyword_parameters(parquet.read_table) + else: + allowed = keyword_parameters(pandas.read_csv) + allowed -= {"filepath_or_buffer", "path", "source"} + allowed |= { + "engine", + "chunksize", + "block_size", + "cache_type", + "max_workers", + "auto_optimize_chunksize", + } + validate_kwargs(method, kwargs, allowed) + + def _convert_binary_csv(converter: Callable[[str | None], Any], value: str) -> Any: return converter(None if value == _BINARY_NULL else value) diff --git a/pyathena/polars/async_cursor.py b/pyathena/polars/async_cursor.py index 96bb6c600..1f8ade6ae 100644 --- a/pyathena/polars/async_cursor.py +++ b/pyathena/polars/async_cursor.py @@ -16,7 +16,7 @@ DefaultPolarsTypeConverter, DefaultPolarsUnloadTypeConverter, ) -from pyathena.polars.result_set import AthenaPolarsResultSet +from pyathena.polars.result_set import AthenaPolarsResultSet, validate_execute_kwargs from pyathena.util import override _logger = logging.getLogger(__name__) @@ -244,6 +244,12 @@ def execute( >>> result_set = future.result() >>> df = result_set.as_polars() # Returns Polars DataFrame """ + validate_execute_kwargs( + f"{type(self).__name__}.execute", + kwargs, + self._unload, + kwargs.get("chunksize", self._chunksize), + ) options = ExecuteOptions.resolve( options, work_group=work_group, diff --git a/pyathena/polars/cursor.py b/pyathena/polars/cursor.py index 05c91b307..574f62b9f 100644 --- a/pyathena/polars/cursor.py +++ b/pyathena/polars/cursor.py @@ -19,7 +19,7 @@ DefaultPolarsTypeConverter, DefaultPolarsUnloadTypeConverter, ) -from pyathena.polars.result_set import AthenaPolarsResultSet +from pyathena.polars.result_set import AthenaPolarsResultSet, validate_execute_kwargs from pyathena.result_set import WithFetch from pyathena.util import override @@ -202,6 +202,12 @@ def execute( >>> cursor.execute("SELECT * FROM sales WHERE year = 2023") >>> df = cursor.as_polars() # Returns Polars DataFrame """ + validate_execute_kwargs( + f"{type(self).__name__}.execute", + kwargs, + self._unload, + kwargs.get("chunksize", self._chunksize), + ) self._reset_state() options = ExecuteOptions.resolve( options, diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 1ee5bc0bc..3cc4ebd45 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -13,6 +13,7 @@ import logging from collections import abc from collections.abc import Callable, Iterator +from importlib import import_module from io import BytesIO, StringIO from multiprocessing import cpu_count from typing import ( @@ -22,6 +23,7 @@ ) from pyathena import OperationalError +from pyathena._kwargs import keyword_parameters, validate_kwargs from pyathena.converter import Converter from pyathena.error import ProgrammingError from pyathena.model import AthenaQueryExecution @@ -39,6 +41,32 @@ _logger = logging.getLogger(__name__) +def validate_execute_kwargs( + method: str, kwargs: dict[str, Any], unload: bool, chunksize: int | None +) -> None: + """Validate extra execution keyword names before starting a query. + + Args: + method: The cursor method name included in an error. + kwargs: Extra execution arguments, including result-set overrides. + unload: Whether the cursor reads UNLOAD Parquet results. + chunksize: The effective chunk size, or None for eager reading. + + Raises: + TypeError: If an argument is not supported by the selected reader. + """ + if not kwargs: + return + polars = import_module("polars") + if unload: + reader = polars.scan_parquet if chunksize is not None else polars.read_parquet + else: + reader = polars.scan_csv if chunksize is not None else polars.read_csv + allowed = keyword_parameters(reader) - {"source"} + allowed |= {"block_size", "cache_type", "max_workers", "chunksize"} + validate_kwargs(method, kwargs, allowed) + + def _identity(x: Any) -> Any: """Identity function for use as default converter.""" return x diff --git a/pyathena/s3fs/async_cursor.py b/pyathena/s3fs/async_cursor.py index 02b7e61eb..ffbb28df3 100644 --- a/pyathena/s3fs/async_cursor.py +++ b/pyathena/s3fs/async_cursor.py @@ -7,6 +7,7 @@ from multiprocessing import cpu_count from typing import Any, cast +from pyathena._kwargs import validate_kwargs from pyathena.async_cursor import AsyncCursor from pyathena.common import CursorIterator from pyathena.error import ProgrammingError @@ -214,9 +215,10 @@ def execute( options: Shared execution options as an :class:`~pyathena.options.ExecuteOptions` instance. Individual keyword arguments take precedence over ``options`` fields. - **kwargs: Additional execution parameters. + **kwargs: Supported S3FS result-set overrides. Unknown names raise TypeError. ``block_size`` sets the read block size for this query, and ``csv_reader`` overrides the cursor's value. + ``filesystem_class`` overrides the filesystem used to read results. Returns: Tuple of (query_id, Future[AthenaS3FSResultSet]). @@ -226,6 +228,11 @@ def execute( >>> result_set = future.result() >>> rows = result_set.fetchall() """ + validate_kwargs( + f"{type(self).__name__}.execute", + kwargs, + ("block_size", "csv_reader", "filesystem_class"), + ) options = ExecuteOptions.resolve( options, work_group=work_group, diff --git a/pyathena/s3fs/cursor.py b/pyathena/s3fs/cursor.py index 96698fb4f..d7f67de44 100644 --- a/pyathena/s3fs/cursor.py +++ b/pyathena/s3fs/cursor.py @@ -6,6 +6,7 @@ from collections.abc import Callable from typing import Any, cast +from pyathena._kwargs import validate_kwargs from pyathena.common import CursorIterator from pyathena.error import OperationalError from pyathena.model import AthenaQueryExecution @@ -164,9 +165,10 @@ def execute( options: Shared execution options as an :class:`~pyathena.options.ExecuteOptions` instance. Individual keyword arguments take precedence over ``options`` fields. - **kwargs: Additional execution parameters. + **kwargs: Supported S3FS result-set overrides. Unknown names raise TypeError. ``block_size`` sets the read block size for this query, and ``csv_reader`` overrides the cursor's value. + ``filesystem_class`` overrides the filesystem used to read results. Returns: Self reference for method chaining. @@ -175,6 +177,11 @@ def execute( >>> cursor.execute("SELECT * FROM my_table WHERE id = %(id)s", {"id": 123}) >>> rows = cursor.fetchall() """ + validate_kwargs( + f"{type(self).__name__}.execute", + kwargs, + ("block_size", "csv_reader", "filesystem_class"), + ) self._reset_state() options = ExecuteOptions.resolve( options, diff --git a/pyathena/spark/async_cursor.py b/pyathena/spark/async_cursor.py index 1e34e617d..e7aac744e 100644 --- a/pyathena/spark/async_cursor.py +++ b/pyathena/spark/async_cursor.py @@ -12,6 +12,7 @@ from multiprocessing import cpu_count from typing import TYPE_CHECKING, Any, cast +from pyathena._kwargs import validate_kwargs from pyathena.model import AthenaCalculationExecution from pyathena.spark.common import SparkBaseCursor from pyathena.util import override @@ -210,6 +211,7 @@ def execute( work_group: str | None = None, **kwargs, ) -> tuple[str, "Future[AthenaQueryExecution | AthenaCalculationExecution]"]: + validate_kwargs(f"{type(self).__name__}.execute", kwargs) calculation_id = self._calculate( session_id=session_id if session_id else self._session_id, code_block=operation, diff --git a/pyathena/spark/cursor.py b/pyathena/spark/cursor.py index 87fe0ef00..11223e9ac 100644 --- a/pyathena/spark/cursor.py +++ b/pyathena/spark/cursor.py @@ -13,6 +13,7 @@ from typing import Any, cast from pyathena import OperationalError, ProgrammingError +from pyathena._kwargs import validate_kwargs from pyathena.model import AthenaCalculationExecution, AthenaCalculationExecutionStatus from pyathena.spark.common import SparkBaseCursor, WithCalculationExecution from pyathena.util import override @@ -110,6 +111,7 @@ def execute( work_group: str | None = None, **kwargs, ) -> SparkCursor: + validate_kwargs(f"{type(self).__name__}.execute", kwargs) # A failure below must not leave the previous calculation on the cursor. self._calculation_id = None self._calculation_execution = None diff --git a/pyathena/sqlalchemy/base.py b/pyathena/sqlalchemy/base.py index dae07e9f7..289f4749b 100644 --- a/pyathena/sqlalchemy/base.py +++ b/pyathena/sqlalchemy/base.py @@ -536,9 +536,7 @@ def _internal_cursor(raw_connection: PoolProxiedConnection) -> Any: returns its wrapper, so this is typed by the interface used here rather than by the class requested. """ - return raw_connection.driver_connection.cursor( # type: ignore[union-attr] - Cursor, converter=Cursor.get_default_converter() - ) + return raw_connection.driver_connection._internal_cursor(Cursor) # type: ignore[union-attr] def _column(self, name: str | None, type_: str, comment: str | None, partition: bool | None): return { diff --git a/tests/pyathena/test_connection.py b/tests/pyathena/test_connection.py index 7e6e42093..49876d506 100644 --- a/tests/pyathena/test_connection.py +++ b/tests/pyathena/test_connection.py @@ -47,8 +47,8 @@ class RecordingCursor(Cursor): - def __init__(self, **kwargs: Any) -> None: - self.kwargs = kwargs + def __init__(self, custom_option: Any = None, **kwargs: Any) -> None: + self.kwargs = {**kwargs, "custom_option": custom_option} super().__init__(**kwargs) @@ -85,7 +85,7 @@ class TestConnection: ("kill_on_interrupt", False, True), ("retry_config", RetryConfig(attempt=1), RetryConfig(attempt=2)), ("schema_name", "configured", "explicit"), - ("unload", True, False), + ("custom_option", True, False), ], ) def test_cursor_arguments_override_cursor_kwargs(self, key, configured, explicit): diff --git a/tests/pyathena/test_kwargs.py b/tests/pyathena/test_kwargs.py new file mode 100644 index 000000000..5c3d7e314 --- /dev/null +++ b/tests/pyathena/test_kwargs.py @@ -0,0 +1,292 @@ +# Copyright 2026 The PyAthena authors +# +# Licensed under the MIT License. +# See LICENSE or https://opensource.org/licenses/MIT. +# +# SPDX-License-Identifier: MIT + +from inspect import iscoroutinefunction +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from pyathena.aio.arrow.cursor import AioArrowCursor +from pyathena.aio.connection import AioConnection +from pyathena.aio.cursor import AioCursor, AioDictCursor +from pyathena.aio.pandas.cursor import AioPandasCursor +from pyathena.aio.polars.cursor import AioPolarsCursor +from pyathena.aio.s3fs.cursor import AioS3FSCursor +from pyathena.aio.spark.cursor import AioSparkCursor +from pyathena.aio.sqlalchemy.base import AsyncAdapt_pyathena_connection +from pyathena.arrow.async_cursor import AsyncArrowCursor +from pyathena.arrow.cursor import ArrowCursor +from pyathena.async_cursor import AsyncCursor, AsyncDictCursor +from pyathena.connection import Connection +from pyathena.cursor import Cursor, DictCursor +from pyathena.formatter import DefaultParameterFormatter +from pyathena.pandas.async_cursor import AsyncPandasCursor +from pyathena.pandas.cursor import PandasCursor +from pyathena.polars.async_cursor import AsyncPolarsCursor +from pyathena.polars.cursor import PolarsCursor +from pyathena.s3fs.async_cursor import AsyncS3FSCursor +from pyathena.s3fs.cursor import S3FSCursor +from pyathena.spark.async_cursor import AsyncSparkCursor +from pyathena.spark.common import SparkBaseCursor +from pyathena.spark.cursor import SparkCursor +from pyathena.sqlalchemy.base import AthenaDialect +from pyathena.util import RetryConfig + +CURSORS = [ + Cursor, + DictCursor, + AsyncCursor, + AsyncDictCursor, + AioCursor, + AioDictCursor, + PandasCursor, + AsyncPandasCursor, + AioPandasCursor, + ArrowCursor, + AsyncArrowCursor, + AioArrowCursor, + PolarsCursor, + AsyncPolarsCursor, + AioPolarsCursor, + S3FSCursor, + AsyncS3FSCursor, + AioS3FSCursor, + SparkCursor, + AsyncSparkCursor, + AioSparkCursor, +] + + +@pytest.fixture +def offline_connection(monkeypatch): + conn = Connection( + region_name="us-east-1", + s3_staging_dir="s3://bucket/path/", + aws_access_key_id="access_key", + aws_secret_access_key="secret_key", + ) + request = MagicMock(side_effect=AssertionError("An offline test made an AWS request")) + monkeypatch.setattr(conn.client, "_make_api_call", request) + monkeypatch.setattr(SparkBaseCursor, "_start_session", MagicMock(return_value="session")) + yield conn + request.assert_not_called() + conn.close() + + +def _constructor_kwargs(conn, cursor_class): + kwargs = { + "connection": conn, + "converter": cursor_class.get_default_converter(), + "formatter": DefaultParameterFormatter(), + "retry_config": RetryConfig(), + "s3_staging_dir": conn.s3_staging_dir, + "schema_name": conn.schema_name, + "catalog_name": conn.catalog_name, + "work_group": conn.work_group, + "poll_interval": 0, + "encryption_option": None, + "kms_key": None, + "kill_on_interrupt": True, + "result_reuse_enable": False, + "result_reuse_minutes": 60, + } + if issubclass(cursor_class, SparkBaseCursor): + kwargs["terminate_session_on_close"] = False + return kwargs + + +async def _close(cursor): + if iscoroutinefunction(cursor.close): + await cursor.close() + else: + cursor.close() + + +async def _execute(cursor, **kwargs): + if iscoroutinefunction(cursor.execute): + return await cursor.execute("SELECT 1", **kwargs) + return cursor.execute("SELECT 1", **kwargs) + + +@pytest.mark.parametrize("cursor_class", CURSORS) +@pytest.mark.parametrize("source", ["constructor", "cursor", "defaults"]) +def test_unknown_constructor_keyword(offline_connection, cursor_class, source): + if source == "constructor": + factory = cursor_class + kwargs = _constructor_kwargs(offline_connection, cursor_class) + else: + factory = offline_connection.cursor + kwargs = {"cursor": cursor_class} + if source == "defaults": + offline_connection.cursor_kwargs = {"work_gruop": "typo"} + if source != "defaults": + kwargs["work_gruop"] = "typo" + with pytest.raises(TypeError, match="unexpected keyword argument 'work_gruop'"): + factory(**kwargs) + + +@pytest.mark.parametrize("cursor_class", CURSORS) +async def test_shared_constructor_callbacks(offline_connection, cursor_class): + on_start, on_poll = MagicMock(), MagicMock() + offline_connection.on_start_query_execution = on_start + offline_connection.on_poll = on_poll + kwargs = ( + {"terminate_session_on_close": False} if issubclass(cursor_class, SparkBaseCursor) else {} + ) + cursor = offline_connection.cursor(cursor_class, **kwargs) + try: + assert cursor._on_start_query_execution is on_start + assert cursor._on_poll is on_poll + finally: + await _close(cursor) + + +@pytest.mark.parametrize("cursor_class", CURSORS) +async def test_unknown_execute_keyword_preserves_state(offline_connection, cursor_class): + cursor = cursor_class(**_constructor_kwargs(offline_connection, cursor_class)) + cursor._query_id = "previous" + cursor._result_set = MagicMock() + before = vars(cursor).copy() + try: + with pytest.raises(TypeError, match="unexpected keyword argument 'work_gruop'"): + await _execute(cursor, work_gruop="typo") + assert vars(cursor) == before + cursor._result_set.close.assert_not_called() + finally: + await _close(cursor) + + +@pytest.mark.parametrize( + ("cursor_class", "settings", "options"), + [(cls, {}, {"parse_dates": []}) for cls in (PandasCursor, AsyncPandasCursor, AioPandasCursor)] + + [ + (cls, {"unload": True}, {"use_threads": False}) + for cls in (PandasCursor, AsyncPandasCursor, AioPandasCursor) + ] + + [(cls, {}, {"separator": "|"}) for cls in (PolarsCursor, AsyncPolarsCursor, AioPolarsCursor)] + + [ + (cls, {"unload": True}, {"parallel": "none"}) + for cls in (PolarsCursor, AsyncPolarsCursor, AioPolarsCursor) + ] + + [ + (cls, {}, {"block_size": 64, "connect_timeout": None, "request_timeout": 2}) + for cls in (ArrowCursor, AsyncArrowCursor, AioArrowCursor) + ] + + [(cls, {}, {"block_size": 64}) for cls in (S3FSCursor, AsyncS3FSCursor, AioS3FSCursor)], +) +async def test_supported_execute_keywords(offline_connection, cursor_class, settings, options): + cursor = offline_connection.cursor(cursor_class, **settings) + try: + with ( + patch.object(cursor, "_execute", side_effect=RuntimeError("execution reached")), + pytest.raises(RuntimeError, match="execution reached"), + ): + await _execute(cursor, **options) + finally: + await _close(cursor) + + +@pytest.mark.parametrize( + "cursor_class", + [ + PandasCursor, + AsyncPandasCursor, + AioPandasCursor, + PolarsCursor, + AsyncPolarsCursor, + AioPolarsCursor, + ], +) +@pytest.mark.parametrize("unload", [False, True]) +async def test_keyword_rejected_for_other_reader_mode(offline_connection, cursor_class, unload): + cursor = offline_connection.cursor(cursor_class, unload=unload) + if "Pandas" in cursor_class.__name__: + key = "sep" if unload else "use_threads" + else: + key = "separator" if unload else "parallel" + try: + with pytest.raises(TypeError, match=f"unexpected keyword argument '{key}'"): + await _execute(cursor, **{key: "|"}) + finally: + await _close(cursor) + + +@pytest.mark.parametrize( + "cursor_class", + [ + PandasCursor, + AsyncPandasCursor, + AioPandasCursor, + PolarsCursor, + AsyncPolarsCursor, + AioPolarsCursor, + ], +) +async def test_reader_input_cannot_replace_query_results(offline_connection, cursor_class): + cursor = offline_connection.cursor(cursor_class) + key = "filepath_or_buffer" if "Pandas" in cursor_class.__name__ else "source" + try: + with pytest.raises(TypeError, match=f"unexpected keyword argument '{key}'"): + await _execute(cursor, **{key: "other.csv"}) + finally: + await _close(cursor) + + +@pytest.mark.parametrize("cursor_class", [PolarsCursor, AsyncPolarsCursor, AioPolarsCursor]) +@pytest.mark.parametrize("chunksize", [None, 10]) +async def test_polars_selects_effective_chunk_reader(offline_connection, cursor_class, chunksize): + cursor = offline_connection.cursor(cursor_class, chunksize=10) + try: + with patch.object(cursor, "_execute", side_effect=RuntimeError("execution reached")): + error = TypeError if chunksize is None else RuntimeError + message = "include_file_paths" if chunksize is None else "execution reached" + with pytest.raises(error, match=message): + await _execute(cursor, chunksize=chunksize, include_file_paths="path") + finally: + await _close(cursor) + + +@pytest.mark.parametrize( + "cursor_class", + [PandasCursor, ArrowCursor, PolarsCursor, AioPandasCursor, AioArrowCursor, AioPolarsCursor], +) +async def test_internal_cursor_isolates_backend_defaults(cursor_class): + is_async = cursor_class.__name__.startswith("Aio") + cls = AioConnection if is_async else Connection + defaults = {"unload": True, "schema_name": "configured", "arraysize": 25} + conn = cls( + region_name="us-east-1", + s3_staging_dir="s3://bucket/path/", + aws_access_key_id="access_key", + aws_secret_access_key="secret_key", + cursor_class=cursor_class, + cursor_kwargs=defaults, + ) + adapted = AsyncAdapt_pyathena_connection(MagicMock(), conn) if is_async else conn + raw_connection = SimpleNamespace(driver_connection=adapted) + internal = AthenaDialect._internal_cursor(raw_connection) + try: + cursor = internal._cursor if is_async else internal + assert type(cursor) is (AioCursor if is_async else Cursor) + assert cursor._schema_name == "configured" + assert cursor.arraysize == 25 + assert type(cursor._converter) is type(Cursor.get_default_converter()) + assert conn.cursor_kwargs == defaults + # Public creation still rejects settings unsupported by the requested class. + with pytest.raises(TypeError, match="unexpected keyword argument 'unload'"): + conn.cursor(AioCursor if is_async else Cursor) + finally: + await _close(cursor) + conn.close() + + +def test_internal_cursor_does_not_hide_unknown_defaults(offline_connection): + offline_connection.cursor_class = PandasCursor + offline_connection.cursor_kwargs = {"unload": True, "work_gruop": "typo"} + with pytest.raises(TypeError, match="unexpected keyword argument 'work_gruop'"): + offline_connection._internal_cursor(Cursor) From 0241fc231e56a2159d11b552ab3507b7b890056b Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Tue, 6 Oct 2026 01:39:39 +0900 Subject: [PATCH 02/15] Preserve reader mode compatibility and update reflection tests --- docs/cursor.md | 5 +- pyathena/aio/pandas/cursor.py | 8 +- pyathena/aio/polars/cursor.py | 14 ++-- pyathena/pandas/async_cursor.py | 6 +- pyathena/pandas/cursor.py | 8 +- pyathena/polars/async_cursor.py | 12 +-- pyathena/polars/cursor.py | 14 ++-- pyathena/polars/result_set.py | 24 +++++- tests/pyathena/polars/test_result_set.py | 93 ++++++++++++++++++++++-- tests/pyathena/sqlalchemy/test_base.py | 43 +++++++---- tests/pyathena/test_kwargs.py | 62 ++++++++++++++-- 11 files changed, 236 insertions(+), 53 deletions(-) diff --git a/docs/cursor.md b/docs/cursor.md index dd3cad76d..7b867e2d7 100644 --- a/docs/cursor.md +++ b/docs/cursor.md @@ -13,7 +13,10 @@ cursor.execute("SELECT 1", work_group="analytics") The validation happens before starting a query or a Spark calculation. For native asyncio cursors, `execute()` validates arguments when awaited. -Pandas and Polars cursors continue to accept their reader options, with names checked against the installed library and the selected CSV, UNLOAD, and chunking mode. +Pandas and Polars cursors continue to accept their reader options, with names checked against the installed library and the prepared query's CSV or UNLOAD mode. +Polars also checks the chunking mode. +For chunked CSV queries, both eager and scan options are accepted before execution because managed query results use the eager reader. +Options specific to the other CSV reader raise `TypeError` when the result is read. Invalid values and combinations remain subject to the reader's validation. Arrow and S3FS cursors accept their supported result-set options; arbitrary reader options are rejected. diff --git a/pyathena/aio/pandas/cursor.py b/pyathena/aio/pandas/cursor.py index dce49586c..57dbdf353 100644 --- a/pyathena/aio/pandas/cursor.py +++ b/pyathena/aio/pandas/cursor.py @@ -180,8 +180,6 @@ async def execute( Returns: Self reference for method chaining. """ - validate_execute_kwargs(f"{type(self).__name__}.execute", kwargs, self._unload) - self._reset_state() options = ExecuteOptions.resolve( options, work_group=work_group, @@ -195,6 +193,12 @@ async def execute( result_set_type_hints=result_set_type_hints, ) operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) + validate_execute_kwargs( + f"{type(self).__name__}.execute", + kwargs, + self._unload and operation.strip().upper().startswith("UNLOAD"), + ) + self._reset_state() self.query_id = await self._execute( operation, parameters=parameters, diff --git a/pyathena/aio/polars/cursor.py b/pyathena/aio/polars/cursor.py index 451a80d1a..4b0d44900 100644 --- a/pyathena/aio/polars/cursor.py +++ b/pyathena/aio/polars/cursor.py @@ -160,13 +160,6 @@ async def execute( Returns: Self reference for method chaining. """ - validate_execute_kwargs( - f"{type(self).__name__}.execute", - kwargs, - self._unload, - kwargs.get("chunksize", self._chunksize), - ) - self._reset_state() options = ExecuteOptions.resolve( options, work_group=work_group, @@ -180,6 +173,13 @@ async def execute( result_set_type_hints=result_set_type_hints, ) operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) + validate_execute_kwargs( + f"{type(self).__name__}.execute", + kwargs, + self._unload and operation.strip().upper().startswith("UNLOAD"), + kwargs.get("chunksize", self._chunksize), + ) + self._reset_state() self.query_id = await self._execute( operation, parameters=parameters, diff --git a/pyathena/pandas/async_cursor.py b/pyathena/pandas/async_cursor.py index 045631b15..44b240dc9 100644 --- a/pyathena/pandas/async_cursor.py +++ b/pyathena/pandas/async_cursor.py @@ -241,7 +241,6 @@ def execute( Returns: Tuple of (query_id, future) where future resolves to AthenaPandasResultSet. """ - validate_execute_kwargs(f"{type(self).__name__}.execute", kwargs, self._unload) options = ExecuteOptions.resolve( options, work_group=work_group, @@ -254,6 +253,11 @@ def execute( result_set_type_hints=result_set_type_hints, ) operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) + validate_execute_kwargs( + f"{type(self).__name__}.execute", + kwargs, + self._unload and operation.strip().upper().startswith("UNLOAD"), + ) query_id = self._execute( operation, parameters=parameters, diff --git a/pyathena/pandas/cursor.py b/pyathena/pandas/cursor.py index e7fbd2e58..86cef3031 100644 --- a/pyathena/pandas/cursor.py +++ b/pyathena/pandas/cursor.py @@ -211,8 +211,6 @@ def execute( ... {"year": 2023}) >>> df = cursor.as_pandas() # Returns pandas DataFrame """ - validate_execute_kwargs(f"{type(self).__name__}.execute", kwargs, self._unload) - self._reset_state() options = ExecuteOptions.resolve( options, work_group=work_group, @@ -226,6 +224,12 @@ def execute( result_set_type_hints=result_set_type_hints, ) operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) + validate_execute_kwargs( + f"{type(self).__name__}.execute", + kwargs, + self._unload and operation.strip().upper().startswith("UNLOAD"), + ) + self._reset_state() self.query_id = self._execute( operation, parameters=parameters, diff --git a/pyathena/polars/async_cursor.py b/pyathena/polars/async_cursor.py index 1f8ade6ae..b7c1aa15c 100644 --- a/pyathena/polars/async_cursor.py +++ b/pyathena/polars/async_cursor.py @@ -244,12 +244,6 @@ def execute( >>> result_set = future.result() >>> df = result_set.as_polars() # Returns Polars DataFrame """ - validate_execute_kwargs( - f"{type(self).__name__}.execute", - kwargs, - self._unload, - kwargs.get("chunksize", self._chunksize), - ) options = ExecuteOptions.resolve( options, work_group=work_group, @@ -262,6 +256,12 @@ def execute( result_set_type_hints=result_set_type_hints, ) operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) + validate_execute_kwargs( + f"{type(self).__name__}.execute", + kwargs, + self._unload and operation.strip().upper().startswith("UNLOAD"), + kwargs.get("chunksize", self._chunksize), + ) query_id = self._execute( operation, parameters=parameters, diff --git a/pyathena/polars/cursor.py b/pyathena/polars/cursor.py index 574f62b9f..3f8847fcb 100644 --- a/pyathena/polars/cursor.py +++ b/pyathena/polars/cursor.py @@ -202,13 +202,6 @@ def execute( >>> cursor.execute("SELECT * FROM sales WHERE year = 2023") >>> df = cursor.as_polars() # Returns Polars DataFrame """ - validate_execute_kwargs( - f"{type(self).__name__}.execute", - kwargs, - self._unload, - kwargs.get("chunksize", self._chunksize), - ) - self._reset_state() options = ExecuteOptions.resolve( options, work_group=work_group, @@ -222,6 +215,13 @@ def execute( result_set_type_hints=result_set_type_hints, ) operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) + validate_execute_kwargs( + f"{type(self).__name__}.execute", + kwargs, + self._unload and operation.strip().upper().startswith("UNLOAD"), + kwargs.get("chunksize", self._chunksize), + ) + self._reset_state() self.query_id = self._execute( operation, parameters=parameters, diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 3cc4ebd45..735aa00f8 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -42,7 +42,12 @@ def validate_execute_kwargs( - method: str, kwargs: dict[str, Any], unload: bool, chunksize: int | None + method: str, + kwargs: dict[str, Any], + unload: bool, + chunksize: int | None, + *, + allow_eager_fallback: bool = True, ) -> None: """Validate extra execution keyword names before starting a query. @@ -51,6 +56,7 @@ def validate_execute_kwargs( kwargs: Extra execution arguments, including result-set overrides. unload: Whether the cursor reads UNLOAD Parquet results. chunksize: The effective chunk size, or None for eager reading. + allow_eager_fallback: Include eager CSV options until the result location is known. Raises: TypeError: If an argument is not supported by the selected reader. @@ -63,6 +69,18 @@ def validate_execute_kwargs( else: reader = polars.scan_csv if chunksize is not None else polars.read_csv allowed = keyword_parameters(reader) - {"source"} + if not unload and chunksize is not None and allow_eager_fallback: + # Managed query results use read_csv even when chunking was requested. + allowed |= keyword_parameters(polars.read_csv) - {"source"} + # Polars still accepts these renamed parameters through its reader decorators. + for old, new in { + "dtypes": "schema_overrides", + "row_count_name": "row_index_name", + "row_count_offset": "row_index_offset", + "missing_utf8_is_empty_string": "empty_string_is_null", + }.items(): + if new in allowed: + allowed.add(old) allowed |= {"block_size", "cache_type", "max_workers", "chunksize"} validate_kwargs(method, kwargs, allowed) @@ -597,6 +615,7 @@ def _read_csv(self) -> pl.DataFrame: """ import polars as pl + validate_execute_kwargs("Polars CSV reader", self._kwargs, False, None) if not self._is_csv_readable(): return pl.DataFrame() @@ -794,6 +813,9 @@ def _iter_csv_chunks(self) -> Iterator[pl.DataFrame]: """ import polars as pl + validate_execute_kwargs( + "Polars CSV reader", self._kwargs, False, self._chunksize, allow_eager_fallback=False + ) if not self._is_csv_readable(): return diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 7cec99ea4..163ef3ef7 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -32,6 +32,87 @@ def _chunked_result_set() -> AthenaPolarsResultSet: class TestAthenaPolarsResultSet: + def test_managed_csv_accepts_eager_options_with_chunksize(self): + result_set = _chunked_result_set() + result_set._query_execution = None + result_set._converter = DefaultPolarsTypeConverter() + result_set._metadata = tuple( + {"Name": name, "Type": dtype, "Precision": 0, "Scale": 0, "Nullable": "UNKNOWN"} + for name, dtype in (("a", "integer"), ("b", "varchar")) + ) + result_set._kwargs = { + "columns": ["a"], + "n_threads": 1, + "use_pyarrow": False, + "batch_size": 1024, + } + with ( + patch.object( + AthenaPolarsResultSet, "_fetch_all_rows_as_csv", return_value=b"a,b\n1,x\n2,y\n" + ), + patch.object( + AthenaPolarsResultSet, + "_csv_storage_options", + new_callable=PropertyMock, + return_value={}, + ), + ): + frame = result_set._read_csv() + assert frame.to_dict(as_series=False) == {"a": [1, 2]} + + @pytest.mark.parametrize("reader", ["_read_csv", "_iter_csv_chunks"]) + def test_csv_rejects_options_for_other_actual_reader(self, reader): + result_set = _chunked_result_set() + key = "include_file_paths" if reader == "_read_csv" else "columns" + result_set._kwargs = {key: "a"} + with pytest.raises(TypeError, match=f"unexpected keyword argument '{key}'"): + ( + result_set._read_csv() + if reader == "_read_csv" + else list(result_set._iter_csv_chunks()) + ) + + @pytest.mark.parametrize("reader", ["_read_csv", "_iter_csv_chunks"]) + def test_csv_keeps_deprecated_row_index_options(self, tmp_path, reader): + path = tmp_path / "result.csv" + path.write_text("a\n1\n2\n") + result_set = _chunked_result_set() + result_set._kwargs = {"row_count_name": "row_number", "row_count_offset": 2} + with ( + patch.object( + AthenaPolarsResultSet, + "output_location", + new_callable=PropertyMock, + return_value=str(path), + ), + patch.object( + AthenaPolarsResultSet, + "_csv_dtypes", + new_callable=PropertyMock, + return_value={"a": pl.Int64}, + ), + patch.object( + AthenaPolarsResultSet, + "_csv_storage_options", + new_callable=PropertyMock, + return_value={}, + ), + patch.object( + AthenaPolarsResultSet, + "_parquet_storage_options", + new_callable=PropertyMock, + return_value={}, + ), + patch.object(AthenaPolarsResultSet, "_is_csv_readable", return_value=True), + pytest.warns(DeprecationWarning, match="row_count_(name|offset)"), + ): + frame = ( + result_set._read_csv() + if reader == "_read_csv" + else pl.concat(list(result_set._iter_csv_chunks())) + ) + assert frame.to_dict(as_series=False) == {"row_number": [2, 3], "a": [1, 2]} + def test_iter_csv_chunks_raises_when_read_fails_partway(self, tmp_path): """A CSV read that fails partway through the data raises instead of ending early.""" path = tmp_path / "result.csv" @@ -201,7 +282,7 @@ def test_storage_options_replace_defaults(self, reader, function): ), patch.object(AthenaPolarsResultSet, "_is_csv_readable", return_value=True), patch.object(AthenaPolarsResultSet, "_prepare_parquet_location", return_value=True), - patch(f"polars.{function}") as read, + patch(f"polars.{function}", autospec=True) as read, patch("pyathena.polars.result_set.to_column_info"), ): result = getattr(result_set, reader)() @@ -215,14 +296,14 @@ def test_storage_options_replace_defaults(self, reader, function): ({}, {"t": [datetime(2020, 1, 2, 3, 4, 5, 123456), None]}), ({"columns": ["v"]}, {}), ({"new_columns": ["t2", "v2"]}, {"t2": [datetime(2020, 1, 2, 3, 4, 5, 123456), None]}), - ({"with_column_names": lambda names: [n.upper() for n in names]}, OperationalError), + ({"with_column_names": lambda names: [n.upper() for n in names]}, TypeError), ], ) def test_read_csv_truncates_timestamps(self, kwargs, expected): """Timestamps that fail to parse are read again as text and truncated. - With ``with_column_names`` given to execute(), Polars renames the columns as it - reads them, so they are not, and the read fails as before. No AWS calls; the + ``with_column_names`` is a scan-only option and is rejected by the eager reader. + No AWS calls; the GetQueryResults rows are mocked. """ result_set = AthenaPolarsResultSet.__new__(AthenaPolarsResultSet) # bypass __init__ @@ -243,8 +324,8 @@ def test_read_csv_truncates_timestamps(self, kwargs, expected): return_value={}, ), ): - if expected is OperationalError: - with pytest.raises(OperationalError): + if expected is TypeError: + with pytest.raises(TypeError): result_set._read_csv() return df = result_set._read_csv() diff --git a/tests/pyathena/sqlalchemy/test_base.py b/tests/pyathena/sqlalchemy/test_base.py index 43ca5d511..cb52d2838 100644 --- a/tests/pyathena/sqlalchemy/test_base.py +++ b/tests/pyathena/sqlalchemy/test_base.py @@ -25,7 +25,7 @@ from sqlalchemy.util import PluginLoader from pyathena.aio.sqlalchemy.base import AthenaAioDialect -from pyathena.converter import DefaultTypeConverter +from pyathena.connection import Connection from pyathena.cursor import Cursor from pyathena.error import DatabaseError, OperationalError from pyathena.formatter import DefaultParameterFormatter @@ -178,11 +178,13 @@ def execute(operation, **kwargs): cursor = SimpleNamespace(execute=execute, fetchall=lambda: rows) - def open_cursor(cursor_class, converter=None): - opened.append((cursor_class, type(converter))) + def open_cursor(cursor_class): + opened.append(cursor_class) return contextlib.nullcontext(cursor) - raw_connection = SimpleNamespace(driver_connection=SimpleNamespace(cursor=open_cursor)) + raw_connection = SimpleNamespace( + driver_connection=SimpleNamespace(_internal_cursor=open_cursor) + ) columns = AthenaDialect()._columns_from_information_schema( raw_connection, "My_Schema", "O'Neil" @@ -190,7 +192,7 @@ def open_cursor(cursor_class, converter=None): # The dialect parses these rows itself, so it asks for an API cursor # rather than whatever result format the user configured. - assert opened == [(Cursor, DefaultTypeConverter)] + assert opened == [Cursor] assert [column["name"] for column in columns] == ["id", "payload", "label", "dt"] assert isinstance(columns[0]["type"], types.INTEGER) @@ -351,7 +353,8 @@ def get_table_metadata(table_name, **kwargs): schema_name=None, retry_config=RetryConfig(), driver_connection=SimpleNamespace( - cursor=lambda *args, **kwargs: contextlib.nullcontext(cursor) + cursor=lambda *args, **kwargs: contextlib.nullcontext(cursor), + _internal_cursor=lambda cursor_class: contextlib.nullcontext(cursor), ), ) connection = SimpleNamespace(connection=raw_connection) @@ -388,7 +391,8 @@ def get_table_metadata(table_name, **kwargs): schema_name="default", retry_config=RetryConfig(), driver_connection=SimpleNamespace( - cursor=lambda *args, **kwargs: contextlib.nullcontext(cursor) + cursor=lambda *args, **kwargs: contextlib.nullcontext(cursor), + _internal_cursor=lambda cursor_class: contextlib.nullcontext(cursor), ), ) connection = SimpleNamespace(connection=raw_connection) @@ -441,7 +445,8 @@ def get_table_metadata(table_name, **kwargs): schema_name="default", retry_config=RetryConfig(), driver_connection=SimpleNamespace( - cursor=lambda *args, **kwargs: contextlib.nullcontext(cursor) + cursor=lambda *args, **kwargs: contextlib.nullcontext(cursor), + _internal_cursor=lambda cursor_class: contextlib.nullcontext(cursor), ), ) connection = SimpleNamespace(connection=raw_connection) @@ -481,7 +486,8 @@ def get_table_metadata(table_name, **kwargs): catalog_name=catalog_name, schema_name="default", driver_connection=SimpleNamespace( - cursor=lambda *args, **kwargs: contextlib.nullcontext(cursor) + cursor=lambda *args, **kwargs: contextlib.nullcontext(cursor), + _internal_cursor=lambda cursor_class: contextlib.nullcontext(cursor), ), ) return SimpleNamespace(connection=raw_connection), error, executed @@ -561,8 +567,8 @@ def test_get_view_definition_keeps_blank_lines(self): executed = [] opened = [] - def open_cursor(cursor_class, converter=None): - opened.append((cursor_class, type(converter))) + def open_cursor(cursor_class): + opened.append(cursor_class) return contextlib.nullcontext( SimpleNamespace( execute=lambda operation, **kwargs: executed.append(operation), @@ -573,14 +579,14 @@ def open_cursor(cursor_class, converter=None): raw_connection = SimpleNamespace( cursor_kwargs={}, schema_name="default", - driver_connection=SimpleNamespace(cursor=open_cursor), + driver_connection=SimpleNamespace(_internal_cursor=open_cursor), ) connection = SimpleNamespace(connection=raw_connection) definition = AthenaDialect().get_view_definition(connection, "v") assert definition == "CREATE VIEW v AS\n\nSELECT 1" - assert opened == [(Cursor, DefaultTypeConverter)] + assert opened == [Cursor] assert executed == ['SHOW CREATE VIEW "default"."v";'] def test_get_view_definition_propagates_a_failed_request(self): @@ -605,7 +611,7 @@ def _view_connection(execute): cursor_kwargs={}, schema_name="default", driver_connection=SimpleNamespace( - cursor=lambda *args, **kwargs: contextlib.nullcontext( + _internal_cursor=lambda cursor_class: contextlib.nullcontext( SimpleNamespace(execute=execute, fetchall=list) ) ), @@ -1414,7 +1420,14 @@ def test_get_view_definition_across_cursor_types(self, engine): definition = sqlalchemy.inspect(conn).get_view_definition(view_name, schema=ENV.schema) # What Athena actually returned, row by row, independent of the # dialect. Comparing against this catches a partial loss too. - with raw_connection.cursor(Cursor) as cursor: + with ( + Connection( + session=raw_connection.session, + s3_staging_dir=raw_connection.s3_staging_dir, + work_group=raw_connection.work_group, + ) as baseline, + baseline.cursor() as cursor, + ): cursor.execute(f'SHOW CREATE VIEW "{ENV.schema}"."{view_name}";') rows = [row[0] for row in cursor.fetchall()] finally: diff --git a/tests/pyathena/test_kwargs.py b/tests/pyathena/test_kwargs.py index 5c3d7e314..dd21b29cc 100644 --- a/tests/pyathena/test_kwargs.py +++ b/tests/pyathena/test_kwargs.py @@ -107,10 +107,10 @@ async def _close(cursor): cursor.close() -async def _execute(cursor, **kwargs): +async def _execute(cursor, operation="SELECT 1", **kwargs): if iscoroutinefunction(cursor.execute): - return await cursor.execute("SELECT 1", **kwargs) - return cursor.execute("SELECT 1", **kwargs) + return await cursor.execute(operation, **kwargs) + return cursor.execute(operation, **kwargs) @pytest.mark.parametrize("cursor_class", CURSORS) @@ -151,6 +151,9 @@ async def test_unknown_execute_keyword_preserves_state(offline_connection, curso cursor = cursor_class(**_constructor_kwargs(offline_connection, cursor_class)) cursor._query_id = "previous" cursor._result_set = MagicMock() + if isinstance(cursor, SparkBaseCursor): + cursor._calculation_id = "previous-calculation" + cursor._calculation_execution = MagicMock() before = vars(cursor).copy() try: with pytest.raises(TypeError, match="unexpected keyword argument 'work_gruop'"): @@ -251,6 +254,53 @@ async def test_polars_selects_effective_chunk_reader(offline_connection, cursor_ await _close(cursor) +@pytest.mark.parametrize( + ("cursor_class", "options"), + [(cls, {"nrows": 10}) for cls in (PandasCursor, AsyncPandasCursor, AioPandasCursor)] + + [ + (cls, {"separator": "\t", "null_values": "NULL"}) + for cls in (PolarsCursor, AsyncPolarsCursor, AioPolarsCursor) + ], +) +async def test_unload_non_select_uses_csv_keywords(offline_connection, cursor_class, options): + cursor = offline_connection.cursor(cursor_class, unload=True) + try: + with ( + patch.object(cursor, "_execute", side_effect=RuntimeError("execution reached")), + pytest.raises(RuntimeError, match="execution reached"), + ): + await _execute(cursor, operation="DESCRIBE t", **options) + finally: + await _close(cursor) + + +@pytest.mark.parametrize("cursor_class", [PolarsCursor, AsyncPolarsCursor, AioPolarsCursor]) +async def test_polars_chunked_csv_allows_managed_reader_options(offline_connection, cursor_class): + cursor = offline_connection.cursor(cursor_class, chunksize=10) + try: + with ( + patch.object(cursor, "_execute", side_effect=RuntimeError("execution reached")), + pytest.raises(RuntimeError, match="execution reached"), + ): + await _execute(cursor, columns=["a"], n_threads=1, use_pyarrow=False, batch_size=1024) + finally: + await _close(cursor) + + +@pytest.mark.parametrize("cursor_class", [PolarsCursor, AsyncPolarsCursor, AioPolarsCursor]) +@pytest.mark.parametrize("unload", [False, True]) +async def test_polars_accepts_reader_aliases(offline_connection, cursor_class, unload): + cursor = offline_connection.cursor(cursor_class, unload=unload) + try: + with ( + patch.object(cursor, "_execute", side_effect=RuntimeError("execution reached")), + pytest.raises(RuntimeError, match="execution reached"), + ): + await _execute(cursor, row_count_name="row_number", row_count_offset=2) + finally: + await _close(cursor) + + @pytest.mark.parametrize( "cursor_class", [PandasCursor, ArrowCursor, PolarsCursor, AioPandasCursor, AioArrowCursor, AioPolarsCursor], @@ -269,8 +319,9 @@ async def test_internal_cursor_isolates_backend_defaults(cursor_class): ) adapted = AsyncAdapt_pyathena_connection(MagicMock(), conn) if is_async else conn raw_connection = SimpleNamespace(driver_connection=adapted) - internal = AthenaDialect._internal_cursor(raw_connection) + cursor = None try: + internal = AthenaDialect._internal_cursor(raw_connection) cursor = internal._cursor if is_async else internal assert type(cursor) is (AioCursor if is_async else Cursor) assert cursor._schema_name == "configured" @@ -281,7 +332,8 @@ async def test_internal_cursor_isolates_backend_defaults(cursor_class): with pytest.raises(TypeError, match="unexpected keyword argument 'unload'"): conn.cursor(AioCursor if is_async else Cursor) finally: - await _close(cursor) + if cursor is not None: + await _close(cursor) conn.close() From fbb2832d2c8db2e93080e74344eb319eb75b5e10 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Tue, 6 Oct 2026 02:03:40 +0900 Subject: [PATCH 03/15] Fix keyword validation failure paths and reader alias forwarding --- docs/cursor.md | 3 +- pyathena/aio/pandas/cursor.py | 39 ++++++++----- pyathena/aio/polars/cursor.py | 40 ++++++++----- pyathena/pandas/async_cursor.py | 36 ++++++++---- pyathena/pandas/cursor.py | 39 ++++++++----- pyathena/polars/async_cursor.py | 37 ++++++++---- pyathena/polars/cursor.py | 40 ++++++++----- pyathena/polars/result_set.py | 21 +++++-- tests/pyathena/polars/test_result_set.py | 65 +++++++++++++++++++-- tests/pyathena/test_kwargs.py | 74 +++++++++++++++++++++++- 10 files changed, 304 insertions(+), 90 deletions(-) diff --git a/docs/cursor.md b/docs/cursor.md index 7b867e2d7..3a2cb813d 100644 --- a/docs/cursor.md +++ b/docs/cursor.md @@ -16,7 +16,8 @@ For native asyncio cursors, `execute()` validates arguments when awaited. Pandas and Polars cursors continue to accept their reader options, with names checked against the installed library and the prepared query's CSV or UNLOAD mode. Polars also checks the chunking mode. For chunked CSV queries, both eager and scan options are accepted before execution because managed query results use the eager reader. -Options specific to the other CSV reader raise `TypeError` when the result is read. +Options specific to the other CSV reader raise `TypeError` when the result set is created. +For Future-based cursors, this error is reported by `future.result()`. Invalid values and combinations remain subject to the reader's validation. Arrow and S3FS cursors accept their supported result-set options; arbitrary reader options are rejected. diff --git a/pyathena/aio/pandas/cursor.py b/pyathena/aio/pandas/cursor.py index 57dbdf353..83805328f 100644 --- a/pyathena/aio/pandas/cursor.py +++ b/pyathena/aio/pandas/cursor.py @@ -180,19 +180,32 @@ async def execute( Returns: Self reference for method chaining. """ - options = ExecuteOptions.resolve( - options, - work_group=work_group, - s3_staging_dir=s3_staging_dir, - cache_size=cache_size, - cache_expiration_time=cache_expiration_time, - result_reuse_enable=result_reuse_enable, - result_reuse_minutes=result_reuse_minutes, - paramstyle=paramstyle, - on_start_query_execution=on_start_query_execution, - result_set_type_hints=result_set_type_hints, - ) - operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) + try: + options = ExecuteOptions.resolve( + options, + work_group=work_group, + s3_staging_dir=s3_staging_dir, + cache_size=cache_size, + cache_expiration_time=cache_expiration_time, + result_reuse_enable=result_reuse_enable, + result_reuse_minutes=result_reuse_minutes, + paramstyle=paramstyle, + on_start_query_execution=on_start_query_execution, + result_set_type_hints=result_set_type_hints, + ) + operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) + except Exception: + validate_execute_kwargs( + f"{type(self).__name__}.execute", + kwargs, + bool( + self._unload + and isinstance(operation, str) + and operation.strip().upper().startswith(("SELECT", "WITH", "UNLOAD")) + ), + ) + self._reset_state() + raise validate_execute_kwargs( f"{type(self).__name__}.execute", kwargs, diff --git a/pyathena/aio/polars/cursor.py b/pyathena/aio/polars/cursor.py index 4b0d44900..2f5a2ca7c 100644 --- a/pyathena/aio/polars/cursor.py +++ b/pyathena/aio/polars/cursor.py @@ -160,19 +160,33 @@ async def execute( Returns: Self reference for method chaining. """ - options = ExecuteOptions.resolve( - options, - work_group=work_group, - s3_staging_dir=s3_staging_dir, - cache_size=cache_size, - cache_expiration_time=cache_expiration_time, - result_reuse_enable=result_reuse_enable, - result_reuse_minutes=result_reuse_minutes, - paramstyle=paramstyle, - on_start_query_execution=on_start_query_execution, - result_set_type_hints=result_set_type_hints, - ) - operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) + try: + options = ExecuteOptions.resolve( + options, + work_group=work_group, + s3_staging_dir=s3_staging_dir, + cache_size=cache_size, + cache_expiration_time=cache_expiration_time, + result_reuse_enable=result_reuse_enable, + result_reuse_minutes=result_reuse_minutes, + paramstyle=paramstyle, + on_start_query_execution=on_start_query_execution, + result_set_type_hints=result_set_type_hints, + ) + operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) + except Exception: + validate_execute_kwargs( + f"{type(self).__name__}.execute", + kwargs, + bool( + self._unload + and isinstance(operation, str) + and operation.strip().upper().startswith(("SELECT", "WITH", "UNLOAD")) + ), + kwargs.get("chunksize", self._chunksize), + ) + self._reset_state() + raise validate_execute_kwargs( f"{type(self).__name__}.execute", kwargs, diff --git a/pyathena/pandas/async_cursor.py b/pyathena/pandas/async_cursor.py index 44b240dc9..ff8e33365 100644 --- a/pyathena/pandas/async_cursor.py +++ b/pyathena/pandas/async_cursor.py @@ -241,18 +241,30 @@ def execute( 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, - cache_expiration_time=cache_expiration_time, - 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) + try: + options = ExecuteOptions.resolve( + options, + work_group=work_group, + s3_staging_dir=s3_staging_dir, + cache_size=cache_size, + cache_expiration_time=cache_expiration_time, + result_reuse_enable=result_reuse_enable, + result_reuse_minutes=result_reuse_minutes, + paramstyle=paramstyle, + result_set_type_hints=result_set_type_hints, + ) + operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) + except Exception: + validate_execute_kwargs( + f"{type(self).__name__}.execute", + kwargs, + bool( + self._unload + and isinstance(operation, str) + and operation.strip().upper().startswith(("SELECT", "WITH", "UNLOAD")) + ), + ) + raise validate_execute_kwargs( f"{type(self).__name__}.execute", kwargs, diff --git a/pyathena/pandas/cursor.py b/pyathena/pandas/cursor.py index 86cef3031..35f6300ed 100644 --- a/pyathena/pandas/cursor.py +++ b/pyathena/pandas/cursor.py @@ -211,19 +211,32 @@ def execute( ... {"year": 2023}) >>> df = cursor.as_pandas() # Returns pandas DataFrame """ - options = ExecuteOptions.resolve( - options, - work_group=work_group, - s3_staging_dir=s3_staging_dir, - cache_size=cache_size, - cache_expiration_time=cache_expiration_time, - result_reuse_enable=result_reuse_enable, - result_reuse_minutes=result_reuse_minutes, - paramstyle=paramstyle, - on_start_query_execution=on_start_query_execution, - result_set_type_hints=result_set_type_hints, - ) - operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) + try: + options = ExecuteOptions.resolve( + options, + work_group=work_group, + s3_staging_dir=s3_staging_dir, + cache_size=cache_size, + cache_expiration_time=cache_expiration_time, + result_reuse_enable=result_reuse_enable, + result_reuse_minutes=result_reuse_minutes, + paramstyle=paramstyle, + on_start_query_execution=on_start_query_execution, + result_set_type_hints=result_set_type_hints, + ) + operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) + except Exception: + validate_execute_kwargs( + f"{type(self).__name__}.execute", + kwargs, + bool( + self._unload + and isinstance(operation, str) + and operation.strip().upper().startswith(("SELECT", "WITH", "UNLOAD")) + ), + ) + self._reset_state() + raise validate_execute_kwargs( f"{type(self).__name__}.execute", kwargs, diff --git a/pyathena/polars/async_cursor.py b/pyathena/polars/async_cursor.py index b7c1aa15c..1f8b76d85 100644 --- a/pyathena/polars/async_cursor.py +++ b/pyathena/polars/async_cursor.py @@ -244,18 +244,31 @@ def execute( >>> result_set = future.result() >>> df = result_set.as_polars() # Returns Polars DataFrame """ - options = ExecuteOptions.resolve( - options, - work_group=work_group, - s3_staging_dir=s3_staging_dir, - cache_size=cache_size, - cache_expiration_time=cache_expiration_time, - result_reuse_enable=result_reuse_enable, - result_reuse_minutes=result_reuse_minutes, - paramstyle=paramstyle, - result_set_type_hints=result_set_type_hints, - ) - operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) + try: + options = ExecuteOptions.resolve( + options, + work_group=work_group, + s3_staging_dir=s3_staging_dir, + cache_size=cache_size, + cache_expiration_time=cache_expiration_time, + result_reuse_enable=result_reuse_enable, + result_reuse_minutes=result_reuse_minutes, + paramstyle=paramstyle, + result_set_type_hints=result_set_type_hints, + ) + operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) + except Exception: + validate_execute_kwargs( + f"{type(self).__name__}.execute", + kwargs, + bool( + self._unload + and isinstance(operation, str) + and operation.strip().upper().startswith(("SELECT", "WITH", "UNLOAD")) + ), + kwargs.get("chunksize", self._chunksize), + ) + raise validate_execute_kwargs( f"{type(self).__name__}.execute", kwargs, diff --git a/pyathena/polars/cursor.py b/pyathena/polars/cursor.py index 3f8847fcb..15fcce838 100644 --- a/pyathena/polars/cursor.py +++ b/pyathena/polars/cursor.py @@ -202,19 +202,33 @@ def execute( >>> cursor.execute("SELECT * FROM sales WHERE year = 2023") >>> df = cursor.as_polars() # Returns Polars DataFrame """ - options = ExecuteOptions.resolve( - options, - work_group=work_group, - s3_staging_dir=s3_staging_dir, - cache_size=cache_size, - cache_expiration_time=cache_expiration_time, - result_reuse_enable=result_reuse_enable, - result_reuse_minutes=result_reuse_minutes, - paramstyle=paramstyle, - on_start_query_execution=on_start_query_execution, - result_set_type_hints=result_set_type_hints, - ) - operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) + try: + options = ExecuteOptions.resolve( + options, + work_group=work_group, + s3_staging_dir=s3_staging_dir, + cache_size=cache_size, + cache_expiration_time=cache_expiration_time, + result_reuse_enable=result_reuse_enable, + result_reuse_minutes=result_reuse_minutes, + paramstyle=paramstyle, + on_start_query_execution=on_start_query_execution, + result_set_type_hints=result_set_type_hints, + ) + operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) + except Exception: + validate_execute_kwargs( + f"{type(self).__name__}.execute", + kwargs, + bool( + self._unload + and isinstance(operation, str) + and operation.strip().upper().startswith(("SELECT", "WITH", "UNLOAD")) + ), + kwargs.get("chunksize", self._chunksize), + ) + self._reset_state() + raise validate_execute_kwargs( f"{type(self).__name__}.execute", kwargs, diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 735aa00f8..d2bb0fd70 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -374,6 +374,8 @@ def _read_kwargs( The arguments for the read function. A value given to ``execute()`` replaces the one the result set chose, including the whole ``storage_options``. """ + if "dtypes" in self._kwargs and "schema_overrides" not in self._kwargs: + defaults.pop("schema_overrides", None) return { **defaults, **self._kwargs, @@ -545,6 +547,13 @@ def _create_dataframe_iterator(self) -> PolarsDataFrameIterator: """ reader: Iterator[pl.DataFrame] if not self.is_unload: + validate_execute_kwargs( + "Polars CSV reader", + self._kwargs, + False, + self._chunksize, + allow_eager_fallback=False, + ) reader = self._iter_csv_chunks() elif self._prepare_parquet_location(): self._metadata = self._read_parquet_schema() @@ -626,7 +635,7 @@ def _read_csv(self) -> pl.DataFrame: has_header=has_header, schema_overrides=self._csv_dtypes, ) - if "schema_overrides" not in self._kwargs: + if not {"schema_overrides", "dtypes"}.intersection(self._kwargs): # Renamed after reading, so that Polars matches the types to the header. kwargs.pop("new_columns", None) source: str | bytes @@ -649,6 +658,7 @@ def _read_csv(self) -> pl.DataFrame: # Athena writes up to 12 fractional digits, which Polars does not parse # into a Datetime whose time unit holds fewer, so the data is read again # with the timestamp columns as text. + kwargs.pop("dtypes", None) kwargs["schema_overrides"] = { **self._csv_dtypes, **dict.fromkeys(timestamp_dtypes, pl.String), @@ -797,7 +807,9 @@ def _get_csv_params(self) -> tuple[str, bool, list[str] | None]: separator = "," has_header = True new_columns = ( - None if "schema_overrides" in self._kwargs else self._kwargs.get("new_columns") + None + if {"schema_overrides", "dtypes"}.intersection(self._kwargs) + else self._kwargs.get("new_columns") ) return separator, has_header, new_columns @@ -813,9 +825,6 @@ def _iter_csv_chunks(self) -> Iterator[pl.DataFrame]: """ import polars as pl - validate_execute_kwargs( - "Polars CSV reader", self._kwargs, False, self._chunksize, allow_eager_fallback=False - ) if not self._is_csv_readable(): return @@ -831,7 +840,7 @@ def _iter_csv_chunks(self) -> Iterator[pl.DataFrame]: has_header=has_header, schema_overrides=self._csv_dtypes, ) - if "schema_overrides" not in self._kwargs: + if not {"schema_overrides", "dtypes"}.intersection(self._kwargs): # Renamed after reading, so that Polars matches the types to the header. read_kwargs.pop("new_columns", None) diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 163ef3ef7..fef01793a 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -63,14 +63,62 @@ def test_managed_csv_accepts_eager_options_with_chunksize(self): @pytest.mark.parametrize("reader", ["_read_csv", "_iter_csv_chunks"]) def test_csv_rejects_options_for_other_actual_reader(self, reader): result_set = _chunked_result_set() + result_set._query_execution = None + result_set._unload = False key = "include_file_paths" if reader == "_read_csv" else "columns" result_set._kwargs = {key: "a"} - with pytest.raises(TypeError, match=f"unexpected keyword argument '{key}'"): - ( + for _ in range(2): + with pytest.raises(TypeError, match=f"unexpected keyword argument '{key}'"): + ( + result_set._read_csv() + if reader == "_read_csv" + else result_set._create_dataframe_iterator() + ) + + @pytest.mark.parametrize("reader", ["_read_csv", "_iter_csv_chunks"]) + @pytest.mark.parametrize("rename", [False, True]) + def test_csv_dtypes_alias_replaces_inferred_schema(self, tmp_path, reader, rename): + path = tmp_path / "result.csv" + path.write_text("a\n001\n002\n") + result_set = _chunked_result_set() + name = "z" if rename else "a" + result_set._kwargs = {"dtypes": {name: pl.String}} + if rename: + result_set._kwargs["new_columns"] = [name] + with ( + patch.object( + AthenaPolarsResultSet, + "output_location", + new_callable=PropertyMock, + return_value=str(path), + ), + patch.object( + AthenaPolarsResultSet, + "_csv_dtypes", + new_callable=PropertyMock, + return_value={"a": pl.Int64}, + ), + patch.object( + AthenaPolarsResultSet, + "_csv_storage_options", + new_callable=PropertyMock, + return_value={}, + ), + patch.object( + AthenaPolarsResultSet, + "_parquet_storage_options", + new_callable=PropertyMock, + return_value={}, + ), + patch.object(AthenaPolarsResultSet, "_is_csv_readable", return_value=True), + pytest.warns(DeprecationWarning, match="dtypes"), + ): + frame = ( result_set._read_csv() if reader == "_read_csv" - else list(result_set._iter_csv_chunks()) + else pl.concat(list(result_set._iter_csv_chunks())) ) + assert frame.to_dict(as_series=False) == {name: ["001", "002"]} @pytest.mark.parametrize("reader", ["_read_csv", "_iter_csv_chunks"]) def test_csv_keeps_deprecated_row_index_options(self, tmp_path, reader): @@ -294,6 +342,10 @@ def test_storage_options_replace_defaults(self, reader, function): ("kwargs", "expected"), [ ({}, {"t": [datetime(2020, 1, 2, 3, 4, 5, 123456), None]}), + ( + {"dtypes": {"t": pl.Datetime("us")}}, + {"t": [datetime(2020, 1, 2, 3, 4, 5, 123456), None]}, + ), ({"columns": ["v"]}, {}), ({"new_columns": ["t2", "v2"]}, {"t2": [datetime(2020, 1, 2, 3, 4, 5, 123456), None]}), ({"with_column_names": lambda names: [n.upper() for n in names]}, TypeError), @@ -303,8 +355,7 @@ def test_read_csv_truncates_timestamps(self, kwargs, expected): """Timestamps that fail to parse are read again as text and truncated. ``with_column_names`` is a scan-only option and is rejected by the eager reader. - No AWS calls; the - GetQueryResults rows are mocked. + No AWS calls; the GetQueryResults rows are mocked. """ result_set = AthenaPolarsResultSet.__new__(AthenaPolarsResultSet) # bypass __init__ result_set._query_execution = None @@ -325,7 +376,9 @@ def test_read_csv_truncates_timestamps(self, kwargs, expected): ), ): if expected is TypeError: - with pytest.raises(TypeError): + with pytest.raises( + TypeError, match="unexpected keyword argument 'with_column_names'" + ): result_set._read_csv() return df = result_set._read_csv() diff --git a/tests/pyathena/test_kwargs.py b/tests/pyathena/test_kwargs.py index dd21b29cc..9ec170249 100644 --- a/tests/pyathena/test_kwargs.py +++ b/tests/pyathena/test_kwargs.py @@ -24,6 +24,7 @@ from pyathena.async_cursor import AsyncCursor, AsyncDictCursor from pyathena.connection import Connection from pyathena.cursor import Cursor, DictCursor +from pyathena.error import ProgrammingError from pyathena.formatter import DefaultParameterFormatter from pyathena.pandas.async_cursor import AsyncPandasCursor from pyathena.pandas.cursor import PandasCursor @@ -150,7 +151,7 @@ async def test_shared_constructor_callbacks(offline_connection, cursor_class): async def test_unknown_execute_keyword_preserves_state(offline_connection, cursor_class): cursor = cursor_class(**_constructor_kwargs(offline_connection, cursor_class)) cursor._query_id = "previous" - cursor._result_set = MagicMock() + cursor._result_set = MagicMock(is_closed=False) if isinstance(cursor, SparkBaseCursor): cursor._calculation_id = "previous-calculation" cursor._calculation_execution = MagicMock() @@ -164,6 +165,76 @@ async def test_unknown_execute_keyword_preserves_state(offline_connection, curso await _close(cursor) +@pytest.mark.parametrize( + "cursor_class", + [ + PandasCursor, + AsyncPandasCursor, + AioPandasCursor, + PolarsCursor, + AsyncPolarsCursor, + AioPolarsCursor, + ], +) +@pytest.mark.parametrize( + "failure", ["empty", "missing_staging", "invalid_options", "invalid_operation"] +) +@pytest.mark.parametrize("unknown_keyword", [False, True]) +async def test_preparation_failure_resets_state_except_unknown_keyword( + offline_connection, cursor_class, failure, unknown_keyword +): + cursor = offline_connection.cursor(cursor_class, unload=True) + previous_result = MagicMock(is_closed=False) + cursor._query_id = "previous" + cursor._result_set = previous_result + operation, options = "SELECT 1", {} + if failure == "empty": + operation = " " + elif failure == "missing_staging": + cursor._s3_staging_dir = None + elif failure == "invalid_options": + options["options"] = object() + else: + operation = object() + if unknown_keyword: + options["work_gruop"] = "typo" + before = vars(cursor).copy() + error = ( + TypeError + if unknown_keyword + else ( + AttributeError + if failure in ("invalid_options", "invalid_operation") + else ProgrammingError + ) + ) + message = ( + "unexpected keyword argument 'work_gruop'" + if unknown_keyword + else ( + "merge" + if failure == "invalid_options" + else ( + "strip" + if failure == "invalid_operation" + else "Query is none|s3_staging_dir is required" + ) + ) + ) + try: + with pytest.raises(error, match=message): + await _execute(cursor, operation=operation, **options) + if unknown_keyword or cursor_class.__name__.startswith("Async"): + assert vars(cursor) == before + previous_result.close.assert_not_called() + else: + assert cursor.query_id is None + assert cursor.result_set is None + previous_result.close.assert_called_once() + finally: + await _close(cursor) + + @pytest.mark.parametrize( ("cursor_class", "settings", "options"), [(cls, {}, {"parse_dates": []}) for cls in (PandasCursor, AsyncPandasCursor, AioPandasCursor)] @@ -316,6 +387,7 @@ async def test_internal_cursor_isolates_backend_defaults(cursor_class): aws_secret_access_key="secret_key", cursor_class=cursor_class, cursor_kwargs=defaults, + converter=cursor_class.get_default_converter(True), ) adapted = AsyncAdapt_pyathena_connection(MagicMock(), conn) if is_async else conn raw_connection = SimpleNamespace(driver_connection=adapted) From 7496dcd7d5788ca02fcde3ad3628fc1b2a1a24bb Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Tue, 6 Oct 2026 02:06:00 +0900 Subject: [PATCH 04/15] Clarify when unknown keyword names are rejected --- docs/cursor.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/cursor.md b/docs/cursor.md index 3a2cb813d..c72b71124 100644 --- a/docs/cursor.md +++ b/docs/cursor.md @@ -11,7 +11,7 @@ cursor = connection.cursor(work_group="analytics") cursor.execute("SELECT 1", work_group="analytics") ``` -The validation happens before starting a query or a Spark calculation. +Unknown keyword names are rejected before starting a query or a Spark calculation. For native asyncio cursors, `execute()` validates arguments when awaited. Pandas and Polars cursors continue to accept their reader options, with names checked against the installed library and the prepared query's CSV or UNLOAD mode. Polars also checks the chunking mode. From 6d6fe2f63b90d0dbb385636cf5aa677a407a1a44 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Tue, 6 Oct 2026 02:24:53 +0900 Subject: [PATCH 05/15] Preserve explicit dtype aliases during timestamp reads --- pyathena/polars/result_set.py | 14 ++++++++------ tests/pyathena/polars/test_result_set.py | 20 +++++++++++++++++--- 2 files changed, 25 insertions(+), 9 deletions(-) diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index d2bb0fd70..97e7202a4 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -482,17 +482,20 @@ def _get_timestamp_dtypes(self, has_header: bool) -> dict[str, Any]: """Get the Datetime dtypes of the timestamp columns, which are read as text. Args: - has_header: Whether the CSV data has a header. Without one, or with - ``schema_overrides`` or ``with_column_names`` given to ``execute()``, - Polars does not name the columns by the header, and they are not read - as text. + has_header: Whether the CSV data has a header. The fallback is disabled + without a header or with ``schema_overrides``, its ``dtypes`` alias, + or ``with_column_names`` given to ``execute()``. Returns: The Datetime dtypes keyed by the header of a CSV file. """ import polars as pl - if not has_header or self._kwargs.keys() & {"schema_overrides", "with_column_names"}: + if not has_header or self._kwargs.keys() & { + "schema_overrides", + "dtypes", + "with_column_names", + }: return {} dtypes = self._csv_dtypes return { @@ -658,7 +661,6 @@ def _read_csv(self) -> pl.DataFrame: # Athena writes up to 12 fractional digits, which Polars does not parse # into a Datetime whose time unit holds fewer, so the data is read again # with the timestamp columns as text. - kwargs.pop("dtypes", None) kwargs["schema_overrides"] = { **self._csv_dtypes, **dict.fromkeys(timestamp_dtypes, pl.String), diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index fef01793a..38e84dbae 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -60,7 +60,7 @@ def test_managed_csv_accepts_eager_options_with_chunksize(self): frame = result_set._read_csv() assert frame.to_dict(as_series=False) == {"a": [1, 2]} - @pytest.mark.parametrize("reader", ["_read_csv", "_iter_csv_chunks"]) + @pytest.mark.parametrize("reader", ["_read_csv", "_create_dataframe_iterator"]) def test_csv_rejects_options_for_other_actual_reader(self, reader): result_set = _chunked_result_set() result_set._query_execution = None @@ -342,9 +342,17 @@ def test_storage_options_replace_defaults(self, reader, function): ("kwargs", "expected"), [ ({}, {"t": [datetime(2020, 1, 2, 3, 4, 5, 123456), None]}), + ({"dtypes": {"t": pl.Datetime("ms")}}, OperationalError), + ({"schema_overrides": {"t": pl.Datetime("ms")}}, OperationalError), + ({"dtypes": {"v": pl.Int64}}, OperationalError), + ({"schema_overrides": {"v": pl.Int64}}, OperationalError), + ( + {"dtypes": {"t2": pl.Datetime("us")}, "new_columns": ["t2", "v2"]}, + OperationalError, + ), ( - {"dtypes": {"t": pl.Datetime("us")}}, - {"t": [datetime(2020, 1, 2, 3, 4, 5, 123456), None]}, + {"schema_overrides": {"t2": pl.Datetime("us")}, "new_columns": ["t2", "v2"]}, + OperationalError, ), ({"columns": ["v"]}, {}), ({"new_columns": ["t2", "v2"]}, {"t2": [datetime(2020, 1, 2, 3, 4, 5, 123456), None]}), @@ -355,6 +363,8 @@ def test_read_csv_truncates_timestamps(self, kwargs, expected): """Timestamps that fail to parse are read again as text and truncated. ``with_column_names`` is a scan-only option and is rejected by the eager reader. + Explicit schema overrides, including the dtypes alias, retain reader errors + instead of retrying with inferred types. No AWS calls; the GetQueryResults rows are mocked. """ result_set = AthenaPolarsResultSet.__new__(AthenaPolarsResultSet) # bypass __init__ @@ -375,6 +385,10 @@ def test_read_csv_truncates_timestamps(self, kwargs, expected): return_value={}, ), ): + if expected is OperationalError: + with pytest.raises(OperationalError, match="could not parse"): + result_set._read_csv() + return if expected is TypeError: with pytest.raises( TypeError, match="unexpected keyword argument 'with_column_names'" From f2460c5c61787aa55203e1f70f3808b2ebb663ff Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Tue, 6 Oct 2026 02:40:49 +0900 Subject: [PATCH 06/15] Clarify the explicit schema override regression contract --- tests/pyathena/polars/test_result_set.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 38e84dbae..bc5736e83 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -363,8 +363,8 @@ def test_read_csv_truncates_timestamps(self, kwargs, expected): """Timestamps that fail to parse are read again as text and truncated. ``with_column_names`` is a scan-only option and is rejected by the eager reader. - Explicit schema overrides, including the dtypes alias, retain reader errors - instead of retrying with inferred types. + The ``schema_overrides`` option and its ``dtypes`` alias disable timestamp + retries and preserve reader errors. No AWS calls; the GetQueryResults rows are mocked. """ result_set = AthenaPolarsResultSet.__new__(AthenaPolarsResultSet) # bypass __init__ From 5b48a622667322689dee3fc973a6d1c26e3e8353 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Wed, 7 Oct 2026 09:15:36 +0900 Subject: [PATCH 07/15] Share UNLOAD reader classification across cursors --- pyathena/aio/pandas/cursor.py | 8 ++----- pyathena/aio/polars/cursor.py | 8 ++----- pyathena/common.py | 16 ++++++++++++++ pyathena/pandas/async_cursor.py | 8 ++----- pyathena/pandas/cursor.py | 8 ++----- pyathena/polars/async_cursor.py | 8 ++----- pyathena/polars/cursor.py | 8 ++----- tests/pyathena/test_kwargs.py | 39 +++++++++++++++++++++++++++++++++ 8 files changed, 67 insertions(+), 36 deletions(-) diff --git a/pyathena/aio/pandas/cursor.py b/pyathena/aio/pandas/cursor.py index 83805328f..234cd26c2 100644 --- a/pyathena/aio/pandas/cursor.py +++ b/pyathena/aio/pandas/cursor.py @@ -198,18 +198,14 @@ async def execute( validate_execute_kwargs( f"{type(self).__name__}.execute", kwargs, - bool( - self._unload - and isinstance(operation, str) - and operation.strip().upper().startswith(("SELECT", "WITH", "UNLOAD")) - ), + self._is_unload_query(operation, prepared=False), ) self._reset_state() raise validate_execute_kwargs( f"{type(self).__name__}.execute", kwargs, - self._unload and operation.strip().upper().startswith("UNLOAD"), + self._is_unload_query(operation), ) self._reset_state() self.query_id = await self._execute( diff --git a/pyathena/aio/polars/cursor.py b/pyathena/aio/polars/cursor.py index 2f5a2ca7c..92adf284a 100644 --- a/pyathena/aio/polars/cursor.py +++ b/pyathena/aio/polars/cursor.py @@ -178,11 +178,7 @@ async def execute( validate_execute_kwargs( f"{type(self).__name__}.execute", kwargs, - bool( - self._unload - and isinstance(operation, str) - and operation.strip().upper().startswith(("SELECT", "WITH", "UNLOAD")) - ), + self._is_unload_query(operation, prepared=False), kwargs.get("chunksize", self._chunksize), ) self._reset_state() @@ -190,7 +186,7 @@ async def execute( validate_execute_kwargs( f"{type(self).__name__}.execute", kwargs, - self._unload and operation.strip().upper().startswith("UNLOAD"), + self._is_unload_query(operation), kwargs.get("chunksize", self._chunksize), ) self._reset_state() diff --git a/pyathena/common.py b/pyathena/common.py index 245011c34..0749b33e9 100644 --- a/pyathena/common.py +++ b/pyathena/common.py @@ -1193,6 +1193,22 @@ def _prepare_query( _logger.debug(query) return query, execution_parameters + def _is_unload_query(self, operation: object, *, prepared: bool = True) -> bool: + """Check whether an operation selects the UNLOAD Parquet reader. + + Args: + operation: The SQL operation, or an invalid value to be checked later. + prepared: Whether ``_prepare_unload()`` completed. Before preparation, + SELECT and WITH operations also select the UNLOAD reader when enabled. + + Returns: + True if UNLOAD is enabled and the operation selects its reader. + """ + if not getattr(self, "_unload", False) or not isinstance(operation, str): + return False + prefixes = ("UNLOAD",) if prepared else ("SELECT", "WITH", "UNLOAD") + return operation.strip().upper().startswith(prefixes) + def _prepare_unload( self, operation: str, diff --git a/pyathena/pandas/async_cursor.py b/pyathena/pandas/async_cursor.py index ff8e33365..d391f1e5b 100644 --- a/pyathena/pandas/async_cursor.py +++ b/pyathena/pandas/async_cursor.py @@ -258,17 +258,13 @@ def execute( validate_execute_kwargs( f"{type(self).__name__}.execute", kwargs, - bool( - self._unload - and isinstance(operation, str) - and operation.strip().upper().startswith(("SELECT", "WITH", "UNLOAD")) - ), + self._is_unload_query(operation, prepared=False), ) raise validate_execute_kwargs( f"{type(self).__name__}.execute", kwargs, - self._unload and operation.strip().upper().startswith("UNLOAD"), + self._is_unload_query(operation), ) query_id = self._execute( operation, diff --git a/pyathena/pandas/cursor.py b/pyathena/pandas/cursor.py index 35f6300ed..f839f5956 100644 --- a/pyathena/pandas/cursor.py +++ b/pyathena/pandas/cursor.py @@ -229,18 +229,14 @@ def execute( validate_execute_kwargs( f"{type(self).__name__}.execute", kwargs, - bool( - self._unload - and isinstance(operation, str) - and operation.strip().upper().startswith(("SELECT", "WITH", "UNLOAD")) - ), + self._is_unload_query(operation, prepared=False), ) self._reset_state() raise validate_execute_kwargs( f"{type(self).__name__}.execute", kwargs, - self._unload and operation.strip().upper().startswith("UNLOAD"), + self._is_unload_query(operation), ) self._reset_state() self.query_id = self._execute( diff --git a/pyathena/polars/async_cursor.py b/pyathena/polars/async_cursor.py index 1f8b76d85..1a32329d5 100644 --- a/pyathena/polars/async_cursor.py +++ b/pyathena/polars/async_cursor.py @@ -261,18 +261,14 @@ def execute( validate_execute_kwargs( f"{type(self).__name__}.execute", kwargs, - bool( - self._unload - and isinstance(operation, str) - and operation.strip().upper().startswith(("SELECT", "WITH", "UNLOAD")) - ), + self._is_unload_query(operation, prepared=False), kwargs.get("chunksize", self._chunksize), ) raise validate_execute_kwargs( f"{type(self).__name__}.execute", kwargs, - self._unload and operation.strip().upper().startswith("UNLOAD"), + self._is_unload_query(operation), kwargs.get("chunksize", self._chunksize), ) query_id = self._execute( diff --git a/pyathena/polars/cursor.py b/pyathena/polars/cursor.py index 15fcce838..f9ec5402e 100644 --- a/pyathena/polars/cursor.py +++ b/pyathena/polars/cursor.py @@ -220,11 +220,7 @@ def execute( validate_execute_kwargs( f"{type(self).__name__}.execute", kwargs, - bool( - self._unload - and isinstance(operation, str) - and operation.strip().upper().startswith(("SELECT", "WITH", "UNLOAD")) - ), + self._is_unload_query(operation, prepared=False), kwargs.get("chunksize", self._chunksize), ) self._reset_state() @@ -232,7 +228,7 @@ def execute( validate_execute_kwargs( f"{type(self).__name__}.execute", kwargs, - self._unload and operation.strip().upper().startswith("UNLOAD"), + self._is_unload_query(operation), kwargs.get("chunksize", self._chunksize), ) self._reset_state() diff --git a/tests/pyathena/test_kwargs.py b/tests/pyathena/test_kwargs.py index 9ec170249..b2cc79012 100644 --- a/tests/pyathena/test_kwargs.py +++ b/tests/pyathena/test_kwargs.py @@ -235,6 +235,45 @@ async def test_preparation_failure_resets_state_except_unknown_keyword( await _close(cursor) +@pytest.mark.parametrize( + "cursor_class", + [ + PandasCursor, + AsyncPandasCursor, + AioPandasCursor, + PolarsCursor, + AsyncPolarsCursor, + AioPolarsCursor, + ], +) +@pytest.mark.parametrize( + ("operation", "unload_reader"), + [ + (" \n sElEcT 1", True), + (" \n wItH t AS (SELECT 1) SELECT * FROM t", True), + ( + " \n uNlOaD (SELECT 1) TO 's3://bucket/path/' WITH (format='PARQUET')", + True, + ), + (" \n dEsCrIbE t", False), + ], +) +async def test_preparation_failure_preserves_supported_reader_keywords( + offline_connection, cursor_class, operation, unload_reader +): + cursor = offline_connection.cursor(cursor_class, unload=True) + cursor._s3_staging_dir = None + if "Polars" in cursor_class.__name__: + options = {"parallel": "none"} if unload_reader else {"separator": "|"} + else: + options = {"use_threads": False} if unload_reader else {"sep": "|"} + try: + with pytest.raises(ProgrammingError, match="s3_staging_dir is required"): + await _execute(cursor, operation=operation, **options) + finally: + await _close(cursor) + + @pytest.mark.parametrize( ("cursor_class", "settings", "options"), [(cls, {}, {"parse_dates": []}) for cls in (PandasCursor, AsyncPandasCursor, AioPandasCursor)] From 5c96c9257f64f6fa5c7fa1bc1cee47d8f496ebee Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 9 Oct 2026 11:48:42 +0900 Subject: [PATCH 08/15] Cache cursor keyword signatures and skip unnecessary SQL classification --- pyathena/_kwargs.py | 28 +++++++-- pyathena/aio/pandas/cursor.py | 13 ++-- pyathena/aio/polars/cursor.py | 13 ++-- pyathena/common.py | 32 +++++++++- pyathena/pandas/async_cursor.py | 13 ++-- pyathena/pandas/cursor.py | 13 ++-- pyathena/polars/async_cursor.py | 13 ++-- pyathena/polars/cursor.py | 13 ++-- tests/pyathena/test_kwargs.py | 101 +++++++++++++++++++++++++++++++- 9 files changed, 196 insertions(+), 43 deletions(-) diff --git a/pyathena/_kwargs.py b/pyathena/_kwargs.py index 03853fa25..2c8ad531a 100644 --- a/pyathena/_kwargs.py +++ b/pyathena/_kwargs.py @@ -8,10 +8,20 @@ """Validation of keyword names forwarded by cursors.""" from collections.abc import Callable, Collection +from functools import lru_cache from inspect import Parameter, signature from typing import Any +@lru_cache(maxsize=64) +def _cached_keyword_parameters(func: Callable[..., Any]) -> frozenset[str]: + return frozenset( + name + for name, parameter in signature(func).parameters.items() + if parameter.kind in (Parameter.POSITIONAL_OR_KEYWORD, Parameter.KEYWORD_ONLY) + ) + + def keyword_parameters(func: Callable[..., Any]) -> set[str]: """Return the explicitly named parameters that accept keyword arguments. @@ -20,12 +30,20 @@ def keyword_parameters(func: Callable[..., Any]) -> set[str]: Returns: Its positional-or-keyword and keyword-only parameter names. + + Notes: + Cache up to 64 callable signatures, returning a fresh set on each call. + Replacing a callable selects a new cache entry. In-place signature changes + require clearing ``_cached_keyword_parameters.cache_clear()``. + Unhashable callables are inspected without caching. """ - return { - name - for name, parameter in signature(func).parameters.items() - if parameter.kind in (Parameter.POSITIONAL_OR_KEYWORD, Parameter.KEYWORD_ONLY) - } + try: + hash(func) + except TypeError: + names = _cached_keyword_parameters.__wrapped__(func) + else: + names = _cached_keyword_parameters(func) + return set(names) def constructor_keyword_parameters(cls: type[Any]) -> set[str]: diff --git a/pyathena/aio/pandas/cursor.py b/pyathena/aio/pandas/cursor.py index 234cd26c2..9c92684df 100644 --- a/pyathena/aio/pandas/cursor.py +++ b/pyathena/aio/pandas/cursor.py @@ -195,17 +195,18 @@ async def execute( ) operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) except Exception: - validate_execute_kwargs( - f"{type(self).__name__}.execute", + self._validate_reader_kwargs( + operation, kwargs, - self._is_unload_query(operation, prepared=False), + validate_execute_kwargs, + prepared=False, ) self._reset_state() raise - validate_execute_kwargs( - f"{type(self).__name__}.execute", + self._validate_reader_kwargs( + operation, kwargs, - self._is_unload_query(operation), + validate_execute_kwargs, ) self._reset_state() self.query_id = await self._execute( diff --git a/pyathena/aio/polars/cursor.py b/pyathena/aio/polars/cursor.py index 92adf284a..4a1234b49 100644 --- a/pyathena/aio/polars/cursor.py +++ b/pyathena/aio/polars/cursor.py @@ -175,18 +175,19 @@ async def execute( ) operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) except Exception: - validate_execute_kwargs( - f"{type(self).__name__}.execute", + self._validate_reader_kwargs( + operation, kwargs, - self._is_unload_query(operation, prepared=False), + validate_execute_kwargs, kwargs.get("chunksize", self._chunksize), + prepared=False, ) self._reset_state() raise - validate_execute_kwargs( - f"{type(self).__name__}.execute", + self._validate_reader_kwargs( + operation, kwargs, - self._is_unload_query(operation), + validate_execute_kwargs, kwargs.get("chunksize", self._chunksize), ) self._reset_state() diff --git a/pyathena/common.py b/pyathena/common.py index 0749b33e9..802df3f9e 100644 --- a/pyathena/common.py +++ b/pyathena/common.py @@ -3,6 +3,7 @@ from __future__ import annotations import logging +import re import sys import threading import time @@ -47,6 +48,7 @@ # a KeyboardInterrupt is raised promptly where an untimed lock wait cannot be # interrupted by signals (Windows before Python 3.14). _INTERRUPT_CHECK_INTERVAL = 0.1 +_OPERATION_START = re.compile(r"\S") OnPollCallback = Callable[[AthenaQueryExecution | AthenaCalculationExecutionStatus], None] """Type of the optional ``on_poll`` callback. @@ -1207,7 +1209,35 @@ def _is_unload_query(self, operation: object, *, prepared: bool = True) -> bool: if not getattr(self, "_unload", False) or not isinstance(operation, str): return False prefixes = ("UNLOAD",) if prepared else ("SELECT", "WITH", "UNLOAD") - return operation.strip().upper().startswith(prefixes) + start = _OPERATION_START.search(operation) + if start is None: + return False + return operation[start.start() : start.start() + 6].upper().startswith(prefixes) + + def _validate_reader_kwargs( + self, + operation: object, + kwargs: dict[str, Any], + validator: Callable[..., None], + *reader_args: Any, + prepared: bool = True, + ) -> None: + """Validate reader keyword names when extra arguments are present. + + Args: + operation: The SQL operation used to select the reader. + kwargs: Extra keyword arguments given to ``execute()``. + validator: The backend's reader keyword validator. + *reader_args: Additional reader settings needed by the validator. + prepared: Whether ``_prepare_unload()`` completed. + """ + if kwargs: + validator( + f"{type(self).__name__}.execute", + kwargs, + self._is_unload_query(operation, prepared=prepared), + *reader_args, + ) def _prepare_unload( self, diff --git a/pyathena/pandas/async_cursor.py b/pyathena/pandas/async_cursor.py index d391f1e5b..aaa167de6 100644 --- a/pyathena/pandas/async_cursor.py +++ b/pyathena/pandas/async_cursor.py @@ -255,16 +255,17 @@ def execute( ) operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) except Exception: - validate_execute_kwargs( - f"{type(self).__name__}.execute", + self._validate_reader_kwargs( + operation, kwargs, - self._is_unload_query(operation, prepared=False), + validate_execute_kwargs, + prepared=False, ) raise - validate_execute_kwargs( - f"{type(self).__name__}.execute", + self._validate_reader_kwargs( + operation, kwargs, - self._is_unload_query(operation), + validate_execute_kwargs, ) query_id = self._execute( operation, diff --git a/pyathena/pandas/cursor.py b/pyathena/pandas/cursor.py index f839f5956..586eeefb0 100644 --- a/pyathena/pandas/cursor.py +++ b/pyathena/pandas/cursor.py @@ -226,17 +226,18 @@ def execute( ) operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) except Exception: - validate_execute_kwargs( - f"{type(self).__name__}.execute", + self._validate_reader_kwargs( + operation, kwargs, - self._is_unload_query(operation, prepared=False), + validate_execute_kwargs, + prepared=False, ) self._reset_state() raise - validate_execute_kwargs( - f"{type(self).__name__}.execute", + self._validate_reader_kwargs( + operation, kwargs, - self._is_unload_query(operation), + validate_execute_kwargs, ) self._reset_state() self.query_id = self._execute( diff --git a/pyathena/polars/async_cursor.py b/pyathena/polars/async_cursor.py index 1a32329d5..d68ac6323 100644 --- a/pyathena/polars/async_cursor.py +++ b/pyathena/polars/async_cursor.py @@ -258,17 +258,18 @@ def execute( ) operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) except Exception: - validate_execute_kwargs( - f"{type(self).__name__}.execute", + self._validate_reader_kwargs( + operation, kwargs, - self._is_unload_query(operation, prepared=False), + validate_execute_kwargs, kwargs.get("chunksize", self._chunksize), + prepared=False, ) raise - validate_execute_kwargs( - f"{type(self).__name__}.execute", + self._validate_reader_kwargs( + operation, kwargs, - self._is_unload_query(operation), + validate_execute_kwargs, kwargs.get("chunksize", self._chunksize), ) query_id = self._execute( diff --git a/pyathena/polars/cursor.py b/pyathena/polars/cursor.py index f9ec5402e..26ba8b89f 100644 --- a/pyathena/polars/cursor.py +++ b/pyathena/polars/cursor.py @@ -217,18 +217,19 @@ def execute( ) operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) except Exception: - validate_execute_kwargs( - f"{type(self).__name__}.execute", + self._validate_reader_kwargs( + operation, kwargs, - self._is_unload_query(operation, prepared=False), + validate_execute_kwargs, kwargs.get("chunksize", self._chunksize), + prepared=False, ) self._reset_state() raise - validate_execute_kwargs( - f"{type(self).__name__}.execute", + self._validate_reader_kwargs( + operation, kwargs, - self._is_unload_query(operation), + validate_execute_kwargs, kwargs.get("chunksize", self._chunksize), ) self._reset_state() diff --git a/tests/pyathena/test_kwargs.py b/tests/pyathena/test_kwargs.py index b2cc79012..921a04c6e 100644 --- a/tests/pyathena/test_kwargs.py +++ b/tests/pyathena/test_kwargs.py @@ -5,12 +5,13 @@ # # SPDX-License-Identifier: MIT -from inspect import iscoroutinefunction +from inspect import iscoroutinefunction, signature from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest +from pyathena._kwargs import constructor_keyword_parameters, keyword_parameters from pyathena.aio.arrow.cursor import AioArrowCursor from pyathena.aio.connection import AioConnection from pyathena.aio.cursor import AioCursor, AioDictCursor @@ -22,6 +23,7 @@ from pyathena.arrow.async_cursor import AsyncArrowCursor from pyathena.arrow.cursor import ArrowCursor from pyathena.async_cursor import AsyncCursor, AsyncDictCursor +from pyathena.common import BaseCursor from pyathena.connection import Connection from pyathena.cursor import Cursor, DictCursor from pyathena.error import ProgrammingError @@ -63,6 +65,75 @@ ] +def test_keyword_parameters_reuses_inspection_without_sharing_mutable_names(): + def reader(source, /, value=None, *args, option=None, **kwargs): + pass + + with patch("pyathena._kwargs.signature", wraps=signature) as inspect_signature: + names = keyword_parameters(reader) + assert names == {"value", "option"} + names.remove("value") + names.add("unsupported") + assert keyword_parameters(reader) == {"value", "option"} + inspect_signature.assert_called_once_with(reader) + + +def test_keyword_parameters_accepts_unhashable_callable(): + class Reader: + __hash__ = None + + def __call__(self, *, setting=None): + pass + + reader = Reader() + assert keyword_parameters(reader) == {"setting"} + assert keyword_parameters(reader) == {"setting"} + + +def test_constructor_keyword_parameters_observes_replaced_constructor(): + class Parent: + def __init__(self, *, shared=None): + pass + + class Child(Parent): + def __init__(self, *, original=None): + pass + + def replacement(self, *, updated=None): + pass + + assert constructor_keyword_parameters(Child) == {"shared", "original"} + Child.__init__ = replacement + assert constructor_keyword_parameters(Child) == {"shared", "updated"} + + +@pytest.mark.parametrize( + ("operation", "before_preparation", "after_preparation"), + [ + (" \n sElEcT 1", True, False), + ("\u2003WITH t AS (SELECT 1) SELECT * FROM t", True, False), + ("\u001cUNLOAD (SELECT 1)", True, True), + ("\u017felect 1", True, False), + ("w\u0131th t AS (SELECT 1) SELECT * FROM t", True, False), + ("w\u0130th t AS (SELECT 1) SELECT * FROM t", False, False), + ("\ufeffSELECT 1", False, False), + ("SELECTED", True, False), + ("DESCRIBE t", False, False), + ("", False, False), + (" \n\t", False, False), + (None, False, False), + ("UNLOAD " + "x" * 100_000, True, True), + ], +) +def test_unload_reader_classification(operation, before_preparation, after_preparation): + cursor = SimpleNamespace(_unload=True) + assert BaseCursor._is_unload_query(cursor, operation, prepared=False) is before_preparation + assert BaseCursor._is_unload_query(cursor, operation) is after_preparation + cursor._unload = False + assert BaseCursor._is_unload_query(cursor, operation, prepared=False) is False + assert BaseCursor._is_unload_query(cursor, operation) is False + + @pytest.fixture def offline_connection(monkeypatch): conn = Connection( @@ -114,6 +185,34 @@ async def _execute(cursor, operation="SELECT 1", **kwargs): return cursor.execute(operation, **kwargs) +@pytest.mark.parametrize( + "cursor_class", + [ + PandasCursor, + AsyncPandasCursor, + AioPandasCursor, + PolarsCursor, + AsyncPolarsCursor, + AioPolarsCursor, + ], +) +@pytest.mark.parametrize("preparation_failure", [False, True]) +async def test_execute_without_reader_keywords_skips_classification( + offline_connection, cursor_class, preparation_failure +): + cursor = offline_connection.cursor(cursor_class, unload=True) + try: + target = "_prepare_unload" if preparation_failure else "_execute" + with ( + patch.object(cursor, "_is_unload_query", side_effect=AssertionError("SQL was scanned")), + patch.object(cursor, target, side_effect=ProgrammingError("execution stopped")), + pytest.raises(ProgrammingError, match="execution stopped"), + ): + await _execute(cursor) + finally: + await _close(cursor) + + @pytest.mark.parametrize("cursor_class", CURSORS) @pytest.mark.parametrize("source", ["constructor", "cursor", "defaults"]) def test_unknown_constructor_keyword(offline_connection, cursor_class, source): From 769b115c3d9b00bdd9764f8d8746fcde54458711 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 9 Oct 2026 23:03:34 +0900 Subject: [PATCH 09/15] Share reader query preparation and skip unnecessary module lookups --- pyathena/aio/pandas/cursor.py | 36 ++++------- pyathena/aio/polars/cursor.py | 37 ++++-------- pyathena/common.py | 63 ++++++++++++++++++++ pyathena/pandas/async_cursor.py | 32 ++++------ pyathena/pandas/cursor.py | 36 ++++------- pyathena/pandas/result_set.py | 22 ++++--- pyathena/polars/async_cursor.py | 33 ++++------ pyathena/polars/cursor.py | 37 ++++-------- pyathena/polars/result_set.py | 13 +++- tests/pyathena/pandas/test_result_set.py | 42 +++++++++++++ tests/pyathena/polars/test_result_set.py | 44 +++++++++++++- tests/pyathena/test_kwargs.py | 76 ++++++++++++++++++++++++ 12 files changed, 316 insertions(+), 155 deletions(-) diff --git a/pyathena/aio/pandas/cursor.py b/pyathena/aio/pandas/cursor.py index 9c92684df..2b920c9c9 100644 --- a/pyathena/aio/pandas/cursor.py +++ b/pyathena/aio/pandas/cursor.py @@ -180,33 +180,21 @@ async def execute( Returns: Self reference for method chaining. """ - try: - options = ExecuteOptions.resolve( - options, - work_group=work_group, - s3_staging_dir=s3_staging_dir, - cache_size=cache_size, - cache_expiration_time=cache_expiration_time, - result_reuse_enable=result_reuse_enable, - result_reuse_minutes=result_reuse_minutes, - paramstyle=paramstyle, - on_start_query_execution=on_start_query_execution, - result_set_type_hints=result_set_type_hints, - ) - operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) - except Exception: - self._validate_reader_kwargs( - operation, - kwargs, - validate_execute_kwargs, - prepared=False, - ) - self._reset_state() - raise - self._validate_reader_kwargs( + operation, unload_location, options = self._prepare_reader_query( operation, kwargs, validate_execute_kwargs, + options=options, + work_group=work_group, + s3_staging_dir=s3_staging_dir, + cache_size=cache_size, + cache_expiration_time=cache_expiration_time, + result_reuse_enable=result_reuse_enable, + result_reuse_minutes=result_reuse_minutes, + paramstyle=paramstyle, + on_start_query_execution=on_start_query_execution, + result_set_type_hints=result_set_type_hints, + on_prepare_error=self._reset_state, ) self._reset_state() self.query_id = await self._execute( diff --git a/pyathena/aio/polars/cursor.py b/pyathena/aio/polars/cursor.py index 4a1234b49..fa9a674db 100644 --- a/pyathena/aio/polars/cursor.py +++ b/pyathena/aio/polars/cursor.py @@ -160,35 +160,22 @@ async def execute( Returns: Self reference for method chaining. """ - try: - options = ExecuteOptions.resolve( - options, - work_group=work_group, - s3_staging_dir=s3_staging_dir, - cache_size=cache_size, - cache_expiration_time=cache_expiration_time, - result_reuse_enable=result_reuse_enable, - result_reuse_minutes=result_reuse_minutes, - paramstyle=paramstyle, - on_start_query_execution=on_start_query_execution, - result_set_type_hints=result_set_type_hints, - ) - operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) - except Exception: - self._validate_reader_kwargs( - operation, - kwargs, - validate_execute_kwargs, - kwargs.get("chunksize", self._chunksize), - prepared=False, - ) - self._reset_state() - raise - self._validate_reader_kwargs( + operation, unload_location, options = self._prepare_reader_query( operation, kwargs, validate_execute_kwargs, kwargs.get("chunksize", self._chunksize), + options=options, + work_group=work_group, + s3_staging_dir=s3_staging_dir, + cache_size=cache_size, + cache_expiration_time=cache_expiration_time, + result_reuse_enable=result_reuse_enable, + result_reuse_minutes=result_reuse_minutes, + paramstyle=paramstyle, + on_start_query_execution=on_start_query_execution, + result_set_type_hints=result_set_type_hints, + on_prepare_error=self._reset_state, ) self._reset_state() self.query_id = await self._execute( diff --git a/pyathena/common.py b/pyathena/common.py index 802df3f9e..e8ddb85d6 100644 --- a/pyathena/common.py +++ b/pyathena/common.py @@ -1239,6 +1239,69 @@ def _validate_reader_kwargs( *reader_args, ) + def _prepare_reader_query( + self, + operation: str, + kwargs: dict[str, Any], + validator: Callable[..., None], + *reader_args: Any, + options: ExecuteOptions | None = None, + work_group: str | None = None, + s3_staging_dir: str | None = None, + 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, + on_prepare_error: Callable[[], None] | None = None, + ) -> tuple[str, str | None, ExecuteOptions]: + """Resolve query options, prepare UNLOAD, and validate reader keywords. + + Args: + operation: The SQL operation to prepare. + kwargs: Extra execution arguments passed to the reader. + validator: The backend's reader keyword validator. + *reader_args: Additional settings needed by the validator. + options: Shared execution options. + work_group: Override for the Athena workgroup. + s3_staging_dir: Override for the query result location. + cache_size: Override for the number of cached queries to check. + cache_expiration_time: Override for cache expiration in seconds. + result_reuse_enable: Override for Athena result reuse. + result_reuse_minutes: Override for the result reuse duration. + paramstyle: Override for the parameter style. + on_start_query_execution: Override for the query-start callback. + result_set_type_hints: Override for complex-type signatures. + on_prepare_error: Reset cursor state after a preparation failure, + only if the reader keywords are valid. + + Returns: + The prepared operation, UNLOAD location, and resolved options. + """ + try: + options = ExecuteOptions.resolve( + options, + work_group=work_group, + s3_staging_dir=s3_staging_dir, + cache_size=cache_size, + cache_expiration_time=cache_expiration_time, + result_reuse_enable=result_reuse_enable, + result_reuse_minutes=result_reuse_minutes, + paramstyle=paramstyle, + on_start_query_execution=on_start_query_execution, + result_set_type_hints=result_set_type_hints, + ) + operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) + except Exception: + self._validate_reader_kwargs(operation, kwargs, validator, *reader_args, prepared=False) + if on_prepare_error is not None: + on_prepare_error() + raise + self._validate_reader_kwargs(operation, kwargs, validator, *reader_args) + return operation, unload_location, options + def _prepare_unload( self, operation: str, diff --git a/pyathena/pandas/async_cursor.py b/pyathena/pandas/async_cursor.py index aaa167de6..18b26ccb3 100644 --- a/pyathena/pandas/async_cursor.py +++ b/pyathena/pandas/async_cursor.py @@ -241,31 +241,19 @@ def execute( Returns: Tuple of (query_id, future) where future resolves to AthenaPandasResultSet. """ - try: - options = ExecuteOptions.resolve( - options, - work_group=work_group, - s3_staging_dir=s3_staging_dir, - cache_size=cache_size, - cache_expiration_time=cache_expiration_time, - result_reuse_enable=result_reuse_enable, - result_reuse_minutes=result_reuse_minutes, - paramstyle=paramstyle, - result_set_type_hints=result_set_type_hints, - ) - operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) - except Exception: - self._validate_reader_kwargs( - operation, - kwargs, - validate_execute_kwargs, - prepared=False, - ) - raise - self._validate_reader_kwargs( + operation, unload_location, options = self._prepare_reader_query( operation, kwargs, validate_execute_kwargs, + options=options, + work_group=work_group, + s3_staging_dir=s3_staging_dir, + cache_size=cache_size, + cache_expiration_time=cache_expiration_time, + result_reuse_enable=result_reuse_enable, + result_reuse_minutes=result_reuse_minutes, + paramstyle=paramstyle, + result_set_type_hints=result_set_type_hints, ) query_id = self._execute( operation, diff --git a/pyathena/pandas/cursor.py b/pyathena/pandas/cursor.py index 586eeefb0..a5b1b6e65 100644 --- a/pyathena/pandas/cursor.py +++ b/pyathena/pandas/cursor.py @@ -211,33 +211,21 @@ def execute( ... {"year": 2023}) >>> df = cursor.as_pandas() # Returns pandas DataFrame """ - try: - options = ExecuteOptions.resolve( - options, - work_group=work_group, - s3_staging_dir=s3_staging_dir, - cache_size=cache_size, - cache_expiration_time=cache_expiration_time, - result_reuse_enable=result_reuse_enable, - result_reuse_minutes=result_reuse_minutes, - paramstyle=paramstyle, - on_start_query_execution=on_start_query_execution, - result_set_type_hints=result_set_type_hints, - ) - operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) - except Exception: - self._validate_reader_kwargs( - operation, - kwargs, - validate_execute_kwargs, - prepared=False, - ) - self._reset_state() - raise - self._validate_reader_kwargs( + operation, unload_location, options = self._prepare_reader_query( operation, kwargs, validate_execute_kwargs, + options=options, + work_group=work_group, + s3_staging_dir=s3_staging_dir, + cache_size=cache_size, + cache_expiration_time=cache_expiration_time, + result_reuse_enable=result_reuse_enable, + result_reuse_minutes=result_reuse_minutes, + paramstyle=paramstyle, + on_start_query_execution=on_start_query_execution, + result_set_type_hints=result_set_type_hints, + on_prepare_error=self._reset_state, ) self._reset_state() self.query_id = self._execute( diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 881ec2a68..7cd90aec3 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -36,6 +36,17 @@ _logger = logging.getLogger(__name__) +_CURSOR_KWARGS = frozenset( + { + "engine", + "chunksize", + "block_size", + "cache_type", + "max_workers", + "auto_optimize_chunksize", + } +) + def validate_execute_kwargs(method: str, kwargs: dict[str, Any], unload: bool) -> None: """Validate extra execution keyword names before starting a query. @@ -48,7 +59,7 @@ def validate_execute_kwargs(method: str, kwargs: dict[str, Any], unload: bool) - Raises: TypeError: If an argument is not supported by the selected reader. """ - if not kwargs: + if not kwargs or kwargs.keys() <= _CURSOR_KWARGS: return pandas = import_module("pandas") if unload: @@ -57,14 +68,7 @@ def validate_execute_kwargs(method: str, kwargs: dict[str, Any], unload: bool) - else: allowed = keyword_parameters(pandas.read_csv) allowed -= {"filepath_or_buffer", "path", "source"} - allowed |= { - "engine", - "chunksize", - "block_size", - "cache_type", - "max_workers", - "auto_optimize_chunksize", - } + allowed |= _CURSOR_KWARGS validate_kwargs(method, kwargs, allowed) diff --git a/pyathena/polars/async_cursor.py b/pyathena/polars/async_cursor.py index d68ac6323..68d892c04 100644 --- a/pyathena/polars/async_cursor.py +++ b/pyathena/polars/async_cursor.py @@ -244,33 +244,20 @@ def execute( >>> result_set = future.result() >>> df = result_set.as_polars() # Returns Polars DataFrame """ - try: - options = ExecuteOptions.resolve( - options, - work_group=work_group, - s3_staging_dir=s3_staging_dir, - cache_size=cache_size, - cache_expiration_time=cache_expiration_time, - result_reuse_enable=result_reuse_enable, - result_reuse_minutes=result_reuse_minutes, - paramstyle=paramstyle, - result_set_type_hints=result_set_type_hints, - ) - operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) - except Exception: - self._validate_reader_kwargs( - operation, - kwargs, - validate_execute_kwargs, - kwargs.get("chunksize", self._chunksize), - prepared=False, - ) - raise - self._validate_reader_kwargs( + operation, unload_location, options = self._prepare_reader_query( operation, kwargs, validate_execute_kwargs, kwargs.get("chunksize", self._chunksize), + options=options, + work_group=work_group, + s3_staging_dir=s3_staging_dir, + cache_size=cache_size, + cache_expiration_time=cache_expiration_time, + result_reuse_enable=result_reuse_enable, + result_reuse_minutes=result_reuse_minutes, + paramstyle=paramstyle, + result_set_type_hints=result_set_type_hints, ) query_id = self._execute( operation, diff --git a/pyathena/polars/cursor.py b/pyathena/polars/cursor.py index 26ba8b89f..2ec006f27 100644 --- a/pyathena/polars/cursor.py +++ b/pyathena/polars/cursor.py @@ -202,35 +202,22 @@ def execute( >>> cursor.execute("SELECT * FROM sales WHERE year = 2023") >>> df = cursor.as_polars() # Returns Polars DataFrame """ - try: - options = ExecuteOptions.resolve( - options, - work_group=work_group, - s3_staging_dir=s3_staging_dir, - cache_size=cache_size, - cache_expiration_time=cache_expiration_time, - result_reuse_enable=result_reuse_enable, - result_reuse_minutes=result_reuse_minutes, - paramstyle=paramstyle, - on_start_query_execution=on_start_query_execution, - result_set_type_hints=result_set_type_hints, - ) - operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) - except Exception: - self._validate_reader_kwargs( - operation, - kwargs, - validate_execute_kwargs, - kwargs.get("chunksize", self._chunksize), - prepared=False, - ) - self._reset_state() - raise - self._validate_reader_kwargs( + operation, unload_location, options = self._prepare_reader_query( operation, kwargs, validate_execute_kwargs, kwargs.get("chunksize", self._chunksize), + options=options, + work_group=work_group, + s3_staging_dir=s3_staging_dir, + cache_size=cache_size, + cache_expiration_time=cache_expiration_time, + result_reuse_enable=result_reuse_enable, + result_reuse_minutes=result_reuse_minutes, + paramstyle=paramstyle, + on_start_query_execution=on_start_query_execution, + result_set_type_hints=result_set_type_hints, + on_prepare_error=self._reset_state, ) self._reset_state() self.query_id = self._execute( diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 97e7202a4..17a141a70 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -40,6 +40,15 @@ _logger = logging.getLogger(__name__) +_CURSOR_KWARGS = frozenset( + { + "block_size", + "cache_type", + "max_workers", + "chunksize", + } +) + def validate_execute_kwargs( method: str, @@ -61,7 +70,7 @@ def validate_execute_kwargs( Raises: TypeError: If an argument is not supported by the selected reader. """ - if not kwargs: + if not kwargs or kwargs.keys() <= _CURSOR_KWARGS: return polars = import_module("polars") if unload: @@ -81,7 +90,7 @@ def validate_execute_kwargs( }.items(): if new in allowed: allowed.add(old) - allowed |= {"block_size", "cache_type", "max_workers", "chunksize"} + allowed |= _CURSOR_KWARGS validate_kwargs(method, kwargs, allowed) diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index 93b2c9c3d..df3e7d75d 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -20,6 +20,7 @@ PandasDataFrameIterator, _no_trunc_date, _read_csv_with_pyarrow, + validate_execute_kwargs, ) @@ -699,3 +700,44 @@ def test_read_csv_with_pyarrow_multiline_values_across_blocks(): ) assert df["id"].tolist() == [i for i, _ in rows] assert df["v"].tolist() == [v for _, v in rows] + + +@pytest.mark.parametrize("unload", [False, True]) +@pytest.mark.parametrize( + "kwargs", + [ + {}, + {"chunksize": 10}, + { + "engine": "c", + "chunksize": 10, + "block_size": 64, + "cache_type": "bytes", + "max_workers": 2, + "auto_optimize_chunksize": False, + }, + ], +) +def test_cursor_settings_skip_reader_imports(kwargs, unload): + with ( + patch("pyathena.pandas.result_set.import_module") as load_module, + patch("pyathena.pandas.result_set.keyword_parameters") as reader_parameters, + ): + validate_execute_kwargs("Cursor.execute", kwargs, unload) + load_module.assert_not_called() + reader_parameters.assert_not_called() + + +def test_reader_replacement_updates_keyword_validation(): + def first_reader(source, *, old_option=None): + pass + + def second_reader(source, *, new_option=None): + pass + + with patch.object(pd, "read_csv", first_reader): + validate_execute_kwargs("Cursor.execute", {"old_option": True}, False) + with patch.object(pd, "read_csv", second_reader): + validate_execute_kwargs("Cursor.execute", {"chunksize": 10, "new_option": True}, False) + with pytest.raises(TypeError, match="unexpected keyword argument 'old_option'"): + validate_execute_kwargs("Cursor.execute", {"chunksize": 10, "old_option": True}, False) diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index bc5736e83..941bcae9d 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -13,7 +13,11 @@ from pyathena.error import OperationalError from pyathena.polars.converter import DefaultPolarsTypeConverter -from pyathena.polars.result_set import AthenaPolarsResultSet, PolarsDataFrameIterator +from pyathena.polars.result_set import ( + AthenaPolarsResultSet, + PolarsDataFrameIterator, + validate_execute_kwargs, +) _ROWS_BEFORE_FAILURE = 300_000 @@ -419,3 +423,41 @@ def test_close_stops_iteration(self, reader): df_iter = PolarsDataFrameIterator(reader, {}, ["a"]) df_iter.close() assert list(df_iter) == [] + + +@pytest.mark.parametrize("unload", [False, True]) +@pytest.mark.parametrize( + "kwargs", + [ + {}, + {"chunksize": 10}, + {"chunksize": 10, "block_size": 64, "cache_type": "bytes", "max_workers": 2}, + ], +) +def test_cursor_settings_skip_reader_imports(kwargs, unload): + with ( + patch("pyathena.polars.result_set.import_module") as load_module, + patch("pyathena.polars.result_set.keyword_parameters") as reader_parameters, + ): + validate_execute_kwargs("Cursor.execute", kwargs, unload, None) + load_module.assert_not_called() + reader_parameters.assert_not_called() + + +def test_reader_replacement_updates_keyword_validation(): + def first_reader(source, *, old_option=None): + pass + + def second_reader(source, *, new_option=None): + pass + + with patch.object(pl, "read_csv", first_reader): + validate_execute_kwargs("Cursor.execute", {"old_option": True}, False, None) + with patch.object(pl, "read_csv", second_reader): + validate_execute_kwargs( + "Cursor.execute", {"chunksize": 10, "new_option": True}, False, None + ) + with pytest.raises(TypeError, match="unexpected keyword argument 'old_option'"): + validate_execute_kwargs( + "Cursor.execute", {"chunksize": 10, "old_option": True}, False, None + ) diff --git a/tests/pyathena/test_kwargs.py b/tests/pyathena/test_kwargs.py index 921a04c6e..b93082e5f 100644 --- a/tests/pyathena/test_kwargs.py +++ b/tests/pyathena/test_kwargs.py @@ -552,3 +552,79 @@ def test_internal_cursor_does_not_hide_unknown_defaults(offline_connection): offline_connection.cursor_kwargs = {"unload": True, "work_gruop": "typo"} with pytest.raises(TypeError, match="unexpected keyword argument 'work_gruop'"): offline_connection._internal_cursor(Cursor) + + +@pytest.mark.parametrize( + "cursor_class", + [ + PandasCursor, + AsyncPandasCursor, + AioPandasCursor, + PolarsCursor, + AsyncPolarsCursor, + AioPolarsCursor, + ], +) +@pytest.mark.parametrize("valid_for_prepared_reader", [False, True]) +async def test_validation_uses_prepared_operation( + offline_connection, cursor_class, valid_for_prepared_reader +): + cursor = offline_connection.cursor(cursor_class, unload=True) + previous_result = MagicMock(is_closed=False) + cursor._query_id = "previous" + cursor._result_set = previous_result + before = vars(cursor).copy() + if "Pandas" in cursor_class.__name__: + kwargs = {"sep": "|"} if valid_for_prepared_reader else {"use_threads": False} + else: + kwargs = {"separator": "|"} if valid_for_prepared_reader else {"parallel": "none"} + try: + with ( + patch.object(cursor, "_prepare_unload", return_value=("DESCRIBE t", None)), + patch.object( + cursor, "_execute", side_effect=RuntimeError("execution reached") + ) as execute, + pytest.raises( + RuntimeError if valid_for_prepared_reader else TypeError, + match="execution reached" if valid_for_prepared_reader else "unexpected keyword", + ), + ): + await _execute(cursor, **kwargs) + if valid_for_prepared_reader: + assert execute.call_args.args[0] == "DESCRIBE t" + else: + execute.assert_not_called() + assert vars(cursor) == before + previous_result.close.assert_not_called() + finally: + await _close(cursor) + + +@pytest.mark.parametrize( + "cursor_class", + [ + PandasCursor, + AsyncPandasCursor, + AioPandasCursor, + PolarsCursor, + AsyncPolarsCursor, + AioPolarsCursor, + ], +) +@pytest.mark.parametrize("error", [KeyboardInterrupt, SystemExit]) +async def test_preparation_interrupt_preserves_state(offline_connection, cursor_class, error): + cursor = offline_connection.cursor(cursor_class, unload=True) + previous_result = MagicMock(is_closed=False) + cursor._query_id = "previous" + cursor._result_set = previous_result + before = vars(cursor).copy() + try: + with ( + patch.object(cursor, "_prepare_unload", side_effect=error), + pytest.raises(error), + ): + await _execute(cursor, work_gruop="typo") + assert vars(cursor) == before + previous_result.close.assert_not_called() + finally: + await _close(cursor) From bb993fa8c3985bee117006f07c6912517d9bdea6 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 9 Oct 2026 23:15:55 +0900 Subject: [PATCH 10/15] Remove the extra reader validation helper call --- pyathena/common.py | 41 ++++++++++++++--------------------------- 1 file changed, 14 insertions(+), 27 deletions(-) diff --git a/pyathena/common.py b/pyathena/common.py index af52474ca..6a8f02530 100644 --- a/pyathena/common.py +++ b/pyathena/common.py @@ -1215,31 +1215,6 @@ def _is_unload_query(self, operation: object, *, prepared: bool = True) -> bool: return False return operation[start.start() : start.start() + 6].upper().startswith(prefixes) - def _validate_reader_kwargs( - self, - operation: object, - kwargs: dict[str, Any], - validator: Callable[..., None], - *reader_args: Any, - prepared: bool = True, - ) -> None: - """Validate reader keyword names when extra arguments are present. - - Args: - operation: The SQL operation used to select the reader. - kwargs: Extra keyword arguments given to ``execute()``. - validator: The backend's reader keyword validator. - *reader_args: Additional reader settings needed by the validator. - prepared: Whether ``_prepare_unload()`` completed. - """ - if kwargs: - validator( - f"{type(self).__name__}.execute", - kwargs, - self._is_unload_query(operation, prepared=prepared), - *reader_args, - ) - def _prepare_reader_query( self, operation: str, @@ -1296,11 +1271,23 @@ def _prepare_reader_query( ) operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) except Exception: - self._validate_reader_kwargs(operation, kwargs, validator, *reader_args, prepared=False) + if kwargs: + validator( + f"{type(self).__name__}.execute", + kwargs, + self._is_unload_query(operation, prepared=False), + *reader_args, + ) if on_prepare_error is not None: on_prepare_error() raise - self._validate_reader_kwargs(operation, kwargs, validator, *reader_args) + if kwargs: + validator( + f"{type(self).__name__}.execute", + kwargs, + self._is_unload_query(operation), + *reader_args, + ) return operation, unload_location, options def _prepare_unload( From c42d65d04a54d3fe1b5cb90122635b495a64d09d Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Fri, 9 Oct 2026 23:48:33 +0900 Subject: [PATCH 11/15] Use native imports for optional reader validation --- pyathena/pandas/result_set.py | 7 ++++--- pyathena/polars/result_set.py | 5 +++-- tests/pyathena/pandas/test_result_set.py | 24 ++++++++++++++++++++++-- tests/pyathena/polars/test_result_set.py | 8 ++++++-- 4 files changed, 35 insertions(+), 9 deletions(-) diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 7cd90aec3..565b3f5fe 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -8,7 +8,6 @@ from collections.abc import Callable, Iterable, Iterator from contextlib import ExitStack from functools import partial -from importlib import import_module from io import BufferedReader, BytesIO, IOBase, StringIO, TextIOWrapper from multiprocessing import cpu_count from typing import ( @@ -61,9 +60,11 @@ def validate_execute_kwargs(method: str, kwargs: dict[str, Any], unload: bool) - """ if not kwargs or kwargs.keys() <= _CURSOR_KWARGS: return - pandas = import_module("pandas") + import pandas + if unload: - parquet = import_module("pyarrow.parquet") + import pyarrow.parquet as parquet + allowed = keyword_parameters(pandas.read_parquet) | keyword_parameters(parquet.read_table) else: allowed = keyword_parameters(pandas.read_csv) diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 17a141a70..2b9d5a046 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -13,7 +13,6 @@ import logging from collections import abc from collections.abc import Callable, Iterator -from importlib import import_module from io import BytesIO, StringIO from multiprocessing import cpu_count from typing import ( @@ -72,7 +71,9 @@ def validate_execute_kwargs( """ if not kwargs or kwargs.keys() <= _CURSOR_KWARGS: return - polars = import_module("polars") + import polars + + reader: Callable[..., Any] if unload: reader = polars.scan_parquet if chunksize is not None else polars.read_parquet else: diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index df3e7d75d..df34db9db 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -5,12 +5,14 @@ # # SPDX-License-Identifier: MIT +import builtins import csv import io from unittest.mock import MagicMock, PropertyMock, patch, sentinel import pandas as pd import pyarrow as pa +import pyarrow.parquet as pq import pytest from pandas.testing import assert_frame_equal @@ -720,11 +722,14 @@ def test_read_csv_with_pyarrow_multiline_values_across_blocks(): ) def test_cursor_settings_skip_reader_imports(kwargs, unload): with ( - patch("pyathena.pandas.result_set.import_module") as load_module, + patch("builtins.__import__", wraps=builtins.__import__) as load_module, patch("pyathena.pandas.result_set.keyword_parameters") as reader_parameters, ): validate_execute_kwargs("Cursor.execute", kwargs, unload) - load_module.assert_not_called() + assert not any( + call.args and call.args[0] in {"pandas", "pyarrow.parquet"} + for call in load_module.call_args_list + ) reader_parameters.assert_not_called() @@ -741,3 +746,18 @@ def second_reader(source, *, new_option=None): validate_execute_kwargs("Cursor.execute", {"chunksize": 10, "new_option": True}, False) with pytest.raises(TypeError, match="unexpected keyword argument 'old_option'"): validate_execute_kwargs("Cursor.execute", {"chunksize": 10, "old_option": True}, False) + + +def test_parquet_reader_replacement_updates_keyword_validation(): + def first_reader(source, *, old_option=None): + pass + + def second_reader(source, *, new_option=None): + pass + + with patch.object(pq, "read_table", first_reader): + validate_execute_kwargs("Cursor.execute", {"old_option": True}, True) + with patch.object(pq, "read_table", second_reader): + validate_execute_kwargs("Cursor.execute", {"new_option": True}, True) + with pytest.raises(TypeError, match="unexpected keyword argument 'old_option'"): + validate_execute_kwargs("Cursor.execute", {"old_option": True}, True) diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index 941bcae9d..b8a728edf 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -5,6 +5,7 @@ # # SPDX-License-Identifier: MIT +import builtins from datetime import datetime from unittest.mock import PropertyMock, patch @@ -436,11 +437,14 @@ def test_close_stops_iteration(self, reader): ) def test_cursor_settings_skip_reader_imports(kwargs, unload): with ( - patch("pyathena.polars.result_set.import_module") as load_module, + patch("builtins.__import__", wraps=builtins.__import__) as load_module, patch("pyathena.polars.result_set.keyword_parameters") as reader_parameters, ): validate_execute_kwargs("Cursor.execute", kwargs, unload, None) - load_module.assert_not_called() + assert not any( + call.args and call.args[0] in {"polars", "pyarrow.parquet"} + for call in load_module.call_args_list + ) reader_parameters.assert_not_called() From 45639cd638ea3f3d71a7fe105084619607fa0296 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 10 Oct 2026 00:45:47 +0900 Subject: [PATCH 12/15] Leave reader option names to the DataFrame readers Pandas and Polars readers already reject unknown keyword names. Checking them before the query required signature inspection, UNLOAD classification, Polars alias handling, and a signature cache, so keep the pandas and Polars cursors as on master. Move the execute() keyword check and the constructor keyword names onto BaseCursor, drop pyathena/_kwargs.py, and move the tests next to the keyword-only tests in the connection test modules. Co-Authored-By: Claude Opus 5.5 --- docs/cursor.md | 25 - docs/usage.md | 4 + pyathena/_kwargs.py | 78 --- pyathena/aio/arrow/cursor.py | 7 +- pyathena/aio/cursor.py | 3 +- pyathena/aio/pandas/cursor.py | 17 +- pyathena/aio/polars/cursor.py | 14 +- pyathena/aio/s3fs/cursor.py | 3 +- pyathena/aio/spark/cursor.py | 3 +- pyathena/aio/sqlalchemy/base.py | 11 +- pyathena/arrow/async_cursor.py | 7 +- pyathena/arrow/cursor.py | 7 +- pyathena/async_cursor.py | 3 +- pyathena/common.py | 136 ++--- pyathena/connection.py | 29 +- pyathena/cursor.py | 3 +- pyathena/pandas/async_cursor.py | 10 +- pyathena/pandas/cursor.py | 17 +- pyathena/pandas/result_set.py | 38 -- pyathena/polars/async_cursor.py | 11 +- pyathena/polars/cursor.py | 14 +- pyathena/polars/result_set.py | 87 +--- pyathena/s3fs/async_cursor.py | 7 +- pyathena/s3fs/cursor.py | 7 +- pyathena/spark/async_cursor.py | 3 +- pyathena/spark/cursor.py | 3 +- tests/pyathena/aio/test_connection.py | 111 +++- tests/pyathena/pandas/test_result_set.py | 62 --- tests/pyathena/polars/test_result_set.py | 208 +------- tests/pyathena/sqlalchemy/test_base.py | 37 +- tests/pyathena/test_connection.py | 128 ++++- tests/pyathena/test_kwargs.py | 630 ----------------------- 32 files changed, 405 insertions(+), 1318 deletions(-) delete mode 100644 pyathena/_kwargs.py delete mode 100644 tests/pyathena/test_kwargs.py diff --git a/docs/cursor.md b/docs/cursor.md index c72b71124..a3db13a2a 100644 --- a/docs/cursor.md +++ b/docs/cursor.md @@ -1,30 +1,5 @@ # Cursor -## Keyword arguments in 4.0.0 - -Starting in 4.0.0, cursor constructors and `execute()` raise `TypeError` for unknown keyword arguments. -Earlier versions silently ignored some extra arguments. -Correct misspelled names and pass backend options to the cursor that supports them: - -```python -cursor = connection.cursor(work_group="analytics") -cursor.execute("SELECT 1", work_group="analytics") -``` - -Unknown keyword names are rejected before starting a query or a Spark calculation. -For native asyncio cursors, `execute()` validates arguments when awaited. -Pandas and Polars cursors continue to accept their reader options, with names checked against the installed library and the prepared query's CSV or UNLOAD mode. -Polars also checks the chunking mode. -For chunked CSV queries, both eager and scan options are accepted before execution because managed query results use the eager reader. -Options specific to the other CSV reader raise `TypeError` when the result set is created. -For Future-based cursors, this error is reported by `future.result()`. -Invalid values and combinations remain subject to the reader's validation. -Arrow and S3FS cursors accept their supported result-set options; arbitrary reader options are rejected. - -Connection-level `on_start_query_execution` and `on_poll` remain accepted by every cursor constructor. -The Future-based Async cursors and Spark cursors do not invoke `on_start_query_execution`. -For reader options, such as pandas `parse_dates`, pass the option to `execute()` rather than the constructor. - (default_cursor)= ## DefaultCursor diff --git a/docs/usage.md b/docs/usage.md index 4a2febcb6..2324b1b1d 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -40,6 +40,10 @@ cursor.execute("SELECT %(value)s", {"value": 1}, work_group="YOUR_WORK_GROUP") Existing keyword arguments retain their defaults and behavior. The cursor class can still be selected positionally with `connection.cursor(DictCursor)`. +Unknown keyword arguments raise `TypeError` in cursor constructors and in `execute()` instead of being ignored. +The connection's `cursor_kwargs` are passed to every cursor that it creates, so the requested cursor class must accept them. +The pandas and Polars cursors pass the `execute()` keyword arguments that they do not use themselves to the reader, such as `pandas.read_csv()`, so an unknown name fails when the results are read. + ## Managed query result storage When using a workgroup with [managed query result storage](https://docs.aws.amazon.com/athena/latest/ug/managed-results.html) enabled, diff --git a/pyathena/_kwargs.py b/pyathena/_kwargs.py deleted file mode 100644 index 2c8ad531a..000000000 --- a/pyathena/_kwargs.py +++ /dev/null @@ -1,78 +0,0 @@ -# Copyright 2026 The PyAthena authors -# -# Licensed under the MIT License. -# See LICENSE or https://opensource.org/licenses/MIT. -# -# SPDX-License-Identifier: MIT - -"""Validation of keyword names forwarded by cursors.""" - -from collections.abc import Callable, Collection -from functools import lru_cache -from inspect import Parameter, signature -from typing import Any - - -@lru_cache(maxsize=64) -def _cached_keyword_parameters(func: Callable[..., Any]) -> frozenset[str]: - return frozenset( - name - for name, parameter in signature(func).parameters.items() - if parameter.kind in (Parameter.POSITIONAL_OR_KEYWORD, Parameter.KEYWORD_ONLY) - ) - - -def keyword_parameters(func: Callable[..., Any]) -> set[str]: - """Return the explicitly named parameters that accept keyword arguments. - - Args: - func: The callable whose signature to inspect. - - Returns: - Its positional-or-keyword and keyword-only parameter names. - - Notes: - Cache up to 64 callable signatures, returning a fresh set on each call. - Replacing a callable selects a new cache entry. In-place signature changes - require clearing ``_cached_keyword_parameters.cache_clear()``. - Unhashable callables are inspected without caching. - """ - try: - hash(func) - except TypeError: - names = _cached_keyword_parameters.__wrapped__(func) - else: - names = _cached_keyword_parameters(func) - return set(names) - - -def constructor_keyword_parameters(cls: type[Any]) -> set[str]: - """Return explicit constructor keyword names across a class's MRO. - - Args: - cls: The class whose forwarding constructors to inspect. - - Returns: - The named constructor parameters, excluding ``self``. - """ - names: set[str] = set() - for base in cls.__mro__: - if "__init__" in vars(base): - names.update(keyword_parameters(vars(base)["__init__"])) - return names - {"self"} - - -def validate_kwargs(method: str, kwargs: dict[str, Any], allowed: Collection[str] = ()) -> None: - """Reject the first keyword name that a cursor does not support. - - Args: - method: The method name included in the error. - kwargs: Extra keyword arguments given to the method. - allowed: The extra keyword names the method supports. - - Raises: - TypeError: If a keyword name is not allowed. - """ - for name in kwargs: - if name not in allowed: - raise TypeError(f"{method}() got an unexpected keyword argument '{name}'") diff --git a/pyathena/aio/arrow/cursor.py b/pyathena/aio/arrow/cursor.py index 6da39d592..a6d5fb17a 100644 --- a/pyathena/aio/arrow/cursor.py +++ b/pyathena/aio/arrow/cursor.py @@ -7,7 +7,6 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Any, cast -from pyathena._kwargs import validate_kwargs from pyathena.aio.common import WithAsyncFetch from pyathena.arrow.converter import ( DefaultArrowTypeConverter, @@ -153,11 +152,7 @@ async def execute( Returns: Self reference for method chaining. """ - validate_kwargs( - f"{type(self).__name__}.execute", - kwargs, - ("block_size", "connect_timeout", "request_timeout"), - ) + self._validate_execute_kwargs(kwargs, ("block_size", "connect_timeout", "request_timeout")) self._reset_state() options = ExecuteOptions.resolve( options, diff --git a/pyathena/aio/cursor.py b/pyathena/aio/cursor.py index 78f5e75be..71ed84d49 100644 --- a/pyathena/aio/cursor.py +++ b/pyathena/aio/cursor.py @@ -13,7 +13,6 @@ from collections.abc import Callable from typing import Any, cast -from pyathena._kwargs import validate_kwargs from pyathena.aio.common import WithAsyncFetch from pyathena.aio.result_set import AthenaAioDictResultSet, AthenaAioResultSet from pyathena.common import CursorIterator @@ -145,7 +144,7 @@ async def execute( Returns: Self reference for method chaining. """ - validate_kwargs(f"{type(self).__name__}.execute", kwargs) + self._validate_execute_kwargs(kwargs) self._reset_state() options = ExecuteOptions.resolve( options, diff --git a/pyathena/aio/pandas/cursor.py b/pyathena/aio/pandas/cursor.py index 238770216..4000d4baa 100644 --- a/pyathena/aio/pandas/cursor.py +++ b/pyathena/aio/pandas/cursor.py @@ -21,11 +21,7 @@ DefaultPandasTypeConverter, DefaultPandasUnloadTypeConverter, ) -from pyathena.pandas.result_set import ( - AthenaPandasResultSet, - PandasDataFrameIterator, - validate_execute_kwargs, -) +from pyathena.pandas.result_set import AthenaPandasResultSet, PandasDataFrameIterator from pyathena.util import override if TYPE_CHECKING: @@ -181,11 +177,9 @@ async def execute( Returns: Self reference for method chaining. """ - operation, unload_location, options = self._prepare_reader_query( - operation, - kwargs, - validate_execute_kwargs, - options=options, + self._reset_state() + options = ExecuteOptions.resolve( + options, work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -195,9 +189,8 @@ async def execute( paramstyle=paramstyle, on_start_query_execution=on_start_query_execution, result_set_type_hints=result_set_type_hints, - on_prepare_error=self._reset_state, ) - self._reset_state() + operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) self.query_id = await self._execute( operation, parameters=parameters, diff --git a/pyathena/aio/polars/cursor.py b/pyathena/aio/polars/cursor.py index 2d1dc82d0..ff9c7a60c 100644 --- a/pyathena/aio/polars/cursor.py +++ b/pyathena/aio/polars/cursor.py @@ -17,7 +17,7 @@ DefaultPolarsTypeConverter, DefaultPolarsUnloadTypeConverter, ) -from pyathena.polars.result_set import AthenaPolarsResultSet, validate_execute_kwargs +from pyathena.polars.result_set import AthenaPolarsResultSet from pyathena.util import override if TYPE_CHECKING: @@ -161,12 +161,9 @@ async def execute( Returns: Self reference for method chaining. """ - operation, unload_location, options = self._prepare_reader_query( - operation, - kwargs, - validate_execute_kwargs, - kwargs.get("chunksize", self._chunksize), - options=options, + self._reset_state() + options = ExecuteOptions.resolve( + options, work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -176,9 +173,8 @@ async def execute( paramstyle=paramstyle, on_start_query_execution=on_start_query_execution, result_set_type_hints=result_set_type_hints, - on_prepare_error=self._reset_state, ) - self._reset_state() + operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) self.query_id = await self._execute( operation, parameters=parameters, diff --git a/pyathena/aio/s3fs/cursor.py b/pyathena/aio/s3fs/cursor.py index 342d04b18..7fa6ff463 100644 --- a/pyathena/aio/s3fs/cursor.py +++ b/pyathena/aio/s3fs/cursor.py @@ -14,7 +14,6 @@ from collections.abc import Callable from typing import Any -from pyathena._kwargs import validate_kwargs from pyathena.aio.common import WithAsyncFetch from pyathena.common import CursorIterator from pyathena.error import OperationalError @@ -155,7 +154,7 @@ async def execute( Returns: Self reference for method chaining. """ - validate_kwargs(f"{type(self).__name__}.execute", kwargs, ("block_size", "csv_reader")) + self._validate_execute_kwargs(kwargs, ("block_size", "csv_reader")) self._reset_state() options = ExecuteOptions.resolve( options, diff --git a/pyathena/aio/spark/cursor.py b/pyathena/aio/spark/cursor.py index 9b13873c0..a50d4abf4 100644 --- a/pyathena/aio/spark/cursor.py +++ b/pyathena/aio/spark/cursor.py @@ -14,7 +14,6 @@ import uuid from typing import Any, cast -from pyathena._kwargs import validate_kwargs from pyathena.aio.util import async_retry_api_call from pyathena.error import DatabaseError, NotSupportedError, OperationalError, ProgrammingError from pyathena.model import ( @@ -359,7 +358,7 @@ async def execute( Returns: Self reference for method chaining. """ - validate_kwargs(f"{type(self).__name__}.execute", kwargs) + self._validate_execute_kwargs(kwargs) # A failure below must not leave the previous calculation on the cursor. self._calculation_id = None self._calculation_execution = None diff --git a/pyathena/aio/sqlalchemy/base.py b/pyathena/aio/sqlalchemy/base.py index 8c114b7c4..4fa5bd87b 100644 --- a/pyathena/aio/sqlalchemy/base.py +++ b/pyathena/aio/sqlalchemy/base.py @@ -277,7 +277,16 @@ def cursor(self, cursor: Any = None, **kwargs: Any) -> AsyncAdapt_pyathena_curso return AsyncAdapt_pyathena_cursor(raw_cursor) def _internal_cursor(self, cursor: Any) -> AsyncAdapt_pyathena_cursor: - """Adapt an API cursor with the default backend's result settings excluded.""" + """Create a cursor with the wrapped connection's ``_internal_cursor()`` and adapt it. + + A synchronous cursor class is replaced as in ``cursor()``. + + Args: + cursor: The cursor class to create. + + Returns: + The created cursor wrapped in ``AsyncAdapt_pyathena_cursor``. + """ raw_cursor = self._connection._internal_cursor(_ASYNC_CURSOR_CLASSES.get(cursor, cursor)) return AsyncAdapt_pyathena_cursor(raw_cursor) diff --git a/pyathena/arrow/async_cursor.py b/pyathena/arrow/async_cursor.py index 46a0bbff8..7d56e64e1 100644 --- a/pyathena/arrow/async_cursor.py +++ b/pyathena/arrow/async_cursor.py @@ -8,7 +8,6 @@ from typing import Any, cast from pyathena import ProgrammingError -from pyathena._kwargs import validate_kwargs from pyathena.arrow.converter import ( DefaultArrowTypeConverter, DefaultArrowUnloadTypeConverter, @@ -221,11 +220,7 @@ def execute( Returns: Tuple of (query_id, future) where future resolves to AthenaArrowResultSet. """ - validate_kwargs( - f"{type(self).__name__}.execute", - kwargs, - ("block_size", "connect_timeout", "request_timeout"), - ) + self._validate_execute_kwargs(kwargs, ("block_size", "connect_timeout", "request_timeout")) options = ExecuteOptions.resolve( options, work_group=work_group, diff --git a/pyathena/arrow/cursor.py b/pyathena/arrow/cursor.py index 23a3c1e75..e9e7774c3 100644 --- a/pyathena/arrow/cursor.py +++ b/pyathena/arrow/cursor.py @@ -6,7 +6,6 @@ from collections.abc import Callable from typing import TYPE_CHECKING, Any, cast -from pyathena._kwargs import validate_kwargs from pyathena.arrow.converter import ( DefaultArrowTypeConverter, DefaultArrowUnloadTypeConverter, @@ -183,11 +182,7 @@ def execute( >>> cursor.execute("SELECT * FROM sales WHERE year = 2023") >>> table = cursor.as_arrow() # Returns Apache Arrow Table """ - validate_kwargs( - f"{type(self).__name__}.execute", - kwargs, - ("block_size", "connect_timeout", "request_timeout"), - ) + self._validate_execute_kwargs(kwargs, ("block_size", "connect_timeout", "request_timeout")) self._reset_state() options = ExecuteOptions.resolve( options, diff --git a/pyathena/async_cursor.py b/pyathena/async_cursor.py index ea35f330f..5f047933c 100644 --- a/pyathena/async_cursor.py +++ b/pyathena/async_cursor.py @@ -8,7 +8,6 @@ from multiprocessing import cpu_count from typing import Any, cast -from pyathena._kwargs import validate_kwargs from pyathena.common import BaseCursor, CursorIterator from pyathena.error import NotSupportedError, ProgrammingError from pyathena.model import AthenaQueryExecution @@ -255,7 +254,7 @@ def execute( >>> # Do other work while query runs... >>> result_set = future.result() # Wait for completion """ - validate_kwargs(f"{type(self).__name__}.execute", kwargs) + self._validate_execute_kwargs(kwargs) options = ExecuteOptions.resolve( options, work_group=work_group, diff --git a/pyathena/common.py b/pyathena/common.py index 6a8f02530..ae09d313e 100644 --- a/pyathena/common.py +++ b/pyathena/common.py @@ -3,14 +3,14 @@ from __future__ import annotations import logging -import re import sys import threading import time from abc import ABCMeta, abstractmethod -from collections.abc import Callable +from collections.abc import Callable, Collection from concurrent.futures import Future, wait from datetime import UTC, datetime, timedelta +from inspect import Parameter, signature from typing import TYPE_CHECKING, Any, TypeVar, cast from botocore.exceptions import BotoCoreError, ClientError @@ -48,7 +48,6 @@ # a KeyboardInterrupt is raised promptly where an untimed lock wait cannot be # interrupted by signals (Windows before Python 3.14). _INTERRUPT_CHECK_INTERVAL = 0.1 -_OPERATION_START = re.compile(r"\S") OnPollCallback = Callable[[AthenaQueryExecution | AthenaCalculationExecutionStatus], None] """Type of the optional ``on_poll`` callback. @@ -272,11 +271,48 @@ def get_default_converter(unload: bool = False) -> DefaultTypeConverter | Any: """ return DefaultTypeConverter() + @classmethod + def _constructor_keyword_names(cls) -> set[str]: + """Get the keyword names accepted by the constructors in this class's MRO. + + Returns: + The named parameters of each ``__init__`` defined in the MRO, + excluding ``self``. + """ + names: set[str] = set() + for base in cls.__mro__: + if "__init__" in vars(base): + names.update( + name + for name, parameter in signature(vars(base)["__init__"]).parameters.items() + if parameter.kind in (Parameter.POSITIONAL_OR_KEYWORD, Parameter.KEYWORD_ONLY) + ) + names.discard("self") + return names + @property def connection(self) -> Connection[Any]: """The connection that created this cursor.""" return self._connection + def _validate_execute_kwargs( + self, kwargs: dict[str, Any], allowed: Collection[str] = () + ) -> None: + """Reject the keyword arguments that ``execute()`` does not support. + + Args: + kwargs: The extra keyword arguments given to ``execute()``. + allowed: The extra keyword names that this cursor supports. + + Raises: + TypeError: If a keyword name is not in ``allowed``. + """ + for name in kwargs: + if name not in allowed: + raise TypeError( + f"{type(self).__name__}.execute() got an unexpected keyword argument '{name}'" + ) + def _build_start_query_execution_request( self, query: str, @@ -1196,100 +1232,6 @@ def _prepare_query( _logger.debug(query) return query, execution_parameters - def _is_unload_query(self, operation: object, *, prepared: bool = True) -> bool: - """Check whether an operation selects the UNLOAD Parquet reader. - - Args: - operation: The SQL operation, or an invalid value to be checked later. - prepared: Whether ``_prepare_unload()`` completed. Before preparation, - SELECT and WITH operations also select the UNLOAD reader when enabled. - - Returns: - True if UNLOAD is enabled and the operation selects its reader. - """ - if not getattr(self, "_unload", False) or not isinstance(operation, str): - return False - prefixes = ("UNLOAD",) if prepared else ("SELECT", "WITH", "UNLOAD") - start = _OPERATION_START.search(operation) - if start is None: - return False - return operation[start.start() : start.start() + 6].upper().startswith(prefixes) - - def _prepare_reader_query( - self, - operation: str, - kwargs: dict[str, Any], - validator: Callable[..., None], - *reader_args: Any, - options: ExecuteOptions | None = None, - work_group: str | None = None, - s3_staging_dir: str | None = None, - 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, - on_prepare_error: Callable[[], None] | None = None, - ) -> tuple[str, str | None, ExecuteOptions]: - """Resolve query options, prepare UNLOAD, and validate reader keywords. - - Args: - operation: The SQL operation to prepare. - kwargs: Extra execution arguments passed to the reader. - validator: The backend's reader keyword validator. - *reader_args: Additional settings needed by the validator. - options: Shared execution options. - work_group: Override for the Athena workgroup. - s3_staging_dir: Override for the query result location. - cache_size: Override for the number of cached queries to check. - cache_expiration_time: Override for cache expiration in seconds. - result_reuse_enable: Override for Athena result reuse. - result_reuse_minutes: Override for the result reuse duration. - paramstyle: Override for the parameter style. - on_start_query_execution: Override for the query-start callback. - result_set_type_hints: Override for complex-type signatures. - on_prepare_error: Reset cursor state after a preparation failure, - only if the reader keywords are valid. - - Returns: - The prepared operation, UNLOAD location, and resolved options. - """ - try: - options = ExecuteOptions.resolve( - options, - work_group=work_group, - s3_staging_dir=s3_staging_dir, - cache_size=cache_size, - cache_expiration_time=cache_expiration_time, - result_reuse_enable=result_reuse_enable, - result_reuse_minutes=result_reuse_minutes, - paramstyle=paramstyle, - on_start_query_execution=on_start_query_execution, - result_set_type_hints=result_set_type_hints, - ) - operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) - except Exception: - if kwargs: - validator( - f"{type(self).__name__}.execute", - kwargs, - self._is_unload_query(operation, prepared=False), - *reader_args, - ) - if on_prepare_error is not None: - on_prepare_error() - raise - if kwargs: - validator( - f"{type(self).__name__}.execute", - kwargs, - self._is_unload_query(operation), - *reader_args, - ) - return operation, unload_location, options - def _prepare_unload( self, operation: str, diff --git a/pyathena/connection.py b/pyathena/connection.py index db87168a5..b1d356902 100644 --- a/pyathena/connection.py +++ b/pyathena/connection.py @@ -21,7 +21,6 @@ from botocore.config import Config import pyathena -from pyathena._kwargs import constructor_keyword_parameters from pyathena.common import BaseCursor, CursorIterator, OnPollCallback from pyathena.converter import Converter from pyathena.cursor import Cursor @@ -618,7 +617,16 @@ def cursor( def _cursor( self, cursor: type[FunctionalCursor] | None = None, **kwargs: Any ) -> FunctionalCursor | ConnectionCursor: - """Create a cursor from already resolved cursor defaults.""" + """Create a cursor without applying the connection's ``cursor_kwargs``. + + Args: + cursor: The cursor class. If None, the connection's default cursor class. + **kwargs: The cursor constructor arguments. Omitted connection settings + default to the connection's values. + + Returns: + The created cursor. + """ _cursor = cursor or self.cursor_class converter = kwargs.pop("converter", self._converter) if not converter: @@ -646,9 +654,20 @@ def _cursor( ) def _internal_cursor(self, cursor: type[FunctionalCursor]) -> FunctionalCursor: - """Create an API cursor without the default backend's result settings.""" - backend_options = constructor_keyword_parameters(self.cursor_class) - ( - constructor_keyword_parameters(cursor) + """Create a cursor for queries that PyAthena runs on its own behalf. + + The connection's ``cursor_kwargs`` are applied except for the constructor + arguments that only the default cursor class accepts, such as ``unload`` + of a DataFrame cursor. The cursor class's default converter is used. + + Args: + cursor: The cursor class. + + Returns: + The created cursor. + """ + backend_options = ( + self.cursor_class._constructor_keyword_names() - cursor._constructor_keyword_names() ) kwargs = {k: v for k, v in self.cursor_kwargs.items() if k not in backend_options} kwargs["converter"] = cursor.get_default_converter() diff --git a/pyathena/cursor.py b/pyathena/cursor.py index 2b5222fef..95f3a8933 100644 --- a/pyathena/cursor.py +++ b/pyathena/cursor.py @@ -6,7 +6,6 @@ from collections.abc import Callable from typing import Any, cast -from pyathena._kwargs import validate_kwargs from pyathena.common import CursorIterator from pyathena.error import OperationalError, ProgrammingError from pyathena.model import AthenaQueryExecution @@ -164,7 +163,7 @@ def execute( ... } ... ) """ - validate_kwargs(f"{type(self).__name__}.execute", kwargs) + self._validate_execute_kwargs(kwargs) self._reset_state() options = ExecuteOptions.resolve( options, diff --git a/pyathena/pandas/async_cursor.py b/pyathena/pandas/async_cursor.py index 04a8b7dfe..01c968052 100644 --- a/pyathena/pandas/async_cursor.py +++ b/pyathena/pandas/async_cursor.py @@ -17,7 +17,7 @@ DefaultPandasTypeConverter, DefaultPandasUnloadTypeConverter, ) -from pyathena.pandas.result_set import AthenaPandasResultSet, validate_execute_kwargs +from pyathena.pandas.result_set import AthenaPandasResultSet from pyathena.util import override _logger = logging.getLogger(__name__) @@ -242,11 +242,8 @@ def execute( Returns: Tuple of (query_id, future) where future resolves to AthenaPandasResultSet. """ - operation, unload_location, options = self._prepare_reader_query( - operation, - kwargs, - validate_execute_kwargs, - options=options, + options = ExecuteOptions.resolve( + options, work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -256,6 +253,7 @@ def execute( 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, diff --git a/pyathena/pandas/cursor.py b/pyathena/pandas/cursor.py index b31c894c5..1ed437935 100644 --- a/pyathena/pandas/cursor.py +++ b/pyathena/pandas/cursor.py @@ -19,11 +19,7 @@ DefaultPandasTypeConverter, DefaultPandasUnloadTypeConverter, ) -from pyathena.pandas.result_set import ( - AthenaPandasResultSet, - PandasDataFrameIterator, - validate_execute_kwargs, -) +from pyathena.pandas.result_set import AthenaPandasResultSet, PandasDataFrameIterator from pyathena.result_set import WithFetch from pyathena.util import override @@ -212,11 +208,9 @@ def execute( ... {"year": 2023}) >>> df = cursor.as_pandas() # Returns pandas DataFrame """ - operation, unload_location, options = self._prepare_reader_query( - operation, - kwargs, - validate_execute_kwargs, - options=options, + self._reset_state() + options = ExecuteOptions.resolve( + options, work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -226,9 +220,8 @@ def execute( paramstyle=paramstyle, on_start_query_execution=on_start_query_execution, result_set_type_hints=result_set_type_hints, - on_prepare_error=self._reset_state, ) - self._reset_state() + operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) self.query_id = self._execute( operation, parameters=parameters, diff --git a/pyathena/pandas/result_set.py b/pyathena/pandas/result_set.py index 565b3f5fe..669683d06 100644 --- a/pyathena/pandas/result_set.py +++ b/pyathena/pandas/result_set.py @@ -19,7 +19,6 @@ from fsspec import open as filesystem_open from pyathena import OperationalError -from pyathena._kwargs import keyword_parameters, validate_kwargs from pyathena.converter import Converter from pyathena.error import ProgrammingError from pyathena.model import AthenaQueryExecution @@ -35,43 +34,6 @@ _logger = logging.getLogger(__name__) -_CURSOR_KWARGS = frozenset( - { - "engine", - "chunksize", - "block_size", - "cache_type", - "max_workers", - "auto_optimize_chunksize", - } -) - - -def validate_execute_kwargs(method: str, kwargs: dict[str, Any], unload: bool) -> None: - """Validate extra execution keyword names before starting a query. - - Args: - method: The cursor method name included in an error. - kwargs: Extra execution arguments, including result-set overrides. - unload: Whether the cursor reads UNLOAD Parquet results. - - Raises: - TypeError: If an argument is not supported by the selected reader. - """ - if not kwargs or kwargs.keys() <= _CURSOR_KWARGS: - return - import pandas - - if unload: - import pyarrow.parquet as parquet - - allowed = keyword_parameters(pandas.read_parquet) | keyword_parameters(parquet.read_table) - else: - allowed = keyword_parameters(pandas.read_csv) - allowed -= {"filepath_or_buffer", "path", "source"} - allowed |= _CURSOR_KWARGS - validate_kwargs(method, kwargs, allowed) - def _convert_binary_csv(converter: Callable[[str | None], Any], value: str) -> Any: return converter(None if value == _BINARY_NULL else value) diff --git a/pyathena/polars/async_cursor.py b/pyathena/polars/async_cursor.py index 6e0aadf80..f836b2179 100644 --- a/pyathena/polars/async_cursor.py +++ b/pyathena/polars/async_cursor.py @@ -16,7 +16,7 @@ DefaultPolarsTypeConverter, DefaultPolarsUnloadTypeConverter, ) -from pyathena.polars.result_set import AthenaPolarsResultSet, validate_execute_kwargs +from pyathena.polars.result_set import AthenaPolarsResultSet from pyathena.util import override _logger = logging.getLogger(__name__) @@ -245,12 +245,8 @@ def execute( >>> result_set = future.result() >>> df = result_set.as_polars() # Returns Polars DataFrame """ - operation, unload_location, options = self._prepare_reader_query( - operation, - kwargs, - validate_execute_kwargs, - kwargs.get("chunksize", self._chunksize), - options=options, + options = ExecuteOptions.resolve( + options, work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -260,6 +256,7 @@ def execute( 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, diff --git a/pyathena/polars/cursor.py b/pyathena/polars/cursor.py index 06178f3cb..a7e31a9f2 100644 --- a/pyathena/polars/cursor.py +++ b/pyathena/polars/cursor.py @@ -19,7 +19,7 @@ DefaultPolarsTypeConverter, DefaultPolarsUnloadTypeConverter, ) -from pyathena.polars.result_set import AthenaPolarsResultSet, validate_execute_kwargs +from pyathena.polars.result_set import AthenaPolarsResultSet from pyathena.result_set import WithFetch from pyathena.util import override @@ -203,12 +203,9 @@ def execute( >>> cursor.execute("SELECT * FROM sales WHERE year = 2023") >>> df = cursor.as_polars() # Returns Polars DataFrame """ - operation, unload_location, options = self._prepare_reader_query( - operation, - kwargs, - validate_execute_kwargs, - kwargs.get("chunksize", self._chunksize), - options=options, + self._reset_state() + options = ExecuteOptions.resolve( + options, work_group=work_group, s3_staging_dir=s3_staging_dir, cache_size=cache_size, @@ -218,9 +215,8 @@ def execute( paramstyle=paramstyle, on_start_query_execution=on_start_query_execution, result_set_type_hints=result_set_type_hints, - on_prepare_error=self._reset_state, ) - self._reset_state() + operation, unload_location = self._prepare_unload(operation, options.s3_staging_dir) self.query_id = self._execute( operation, parameters=parameters, diff --git a/pyathena/polars/result_set.py b/pyathena/polars/result_set.py index 2b9d5a046..1ee5bc0bc 100644 --- a/pyathena/polars/result_set.py +++ b/pyathena/polars/result_set.py @@ -22,7 +22,6 @@ ) from pyathena import OperationalError -from pyathena._kwargs import keyword_parameters, validate_kwargs from pyathena.converter import Converter from pyathena.error import ProgrammingError from pyathena.model import AthenaQueryExecution @@ -39,61 +38,6 @@ _logger = logging.getLogger(__name__) -_CURSOR_KWARGS = frozenset( - { - "block_size", - "cache_type", - "max_workers", - "chunksize", - } -) - - -def validate_execute_kwargs( - method: str, - kwargs: dict[str, Any], - unload: bool, - chunksize: int | None, - *, - allow_eager_fallback: bool = True, -) -> None: - """Validate extra execution keyword names before starting a query. - - Args: - method: The cursor method name included in an error. - kwargs: Extra execution arguments, including result-set overrides. - unload: Whether the cursor reads UNLOAD Parquet results. - chunksize: The effective chunk size, or None for eager reading. - allow_eager_fallback: Include eager CSV options until the result location is known. - - Raises: - TypeError: If an argument is not supported by the selected reader. - """ - if not kwargs or kwargs.keys() <= _CURSOR_KWARGS: - return - import polars - - reader: Callable[..., Any] - if unload: - reader = polars.scan_parquet if chunksize is not None else polars.read_parquet - else: - reader = polars.scan_csv if chunksize is not None else polars.read_csv - allowed = keyword_parameters(reader) - {"source"} - if not unload and chunksize is not None and allow_eager_fallback: - # Managed query results use read_csv even when chunking was requested. - allowed |= keyword_parameters(polars.read_csv) - {"source"} - # Polars still accepts these renamed parameters through its reader decorators. - for old, new in { - "dtypes": "schema_overrides", - "row_count_name": "row_index_name", - "row_count_offset": "row_index_offset", - "missing_utf8_is_empty_string": "empty_string_is_null", - }.items(): - if new in allowed: - allowed.add(old) - allowed |= _CURSOR_KWARGS - validate_kwargs(method, kwargs, allowed) - def _identity(x: Any) -> Any: """Identity function for use as default converter.""" @@ -384,8 +328,6 @@ def _read_kwargs( The arguments for the read function. A value given to ``execute()`` replaces the one the result set chose, including the whole ``storage_options``. """ - if "dtypes" in self._kwargs and "schema_overrides" not in self._kwargs: - defaults.pop("schema_overrides", None) return { **defaults, **self._kwargs, @@ -492,20 +434,17 @@ def _get_timestamp_dtypes(self, has_header: bool) -> dict[str, Any]: """Get the Datetime dtypes of the timestamp columns, which are read as text. Args: - has_header: Whether the CSV data has a header. The fallback is disabled - without a header or with ``schema_overrides``, its ``dtypes`` alias, - or ``with_column_names`` given to ``execute()``. + has_header: Whether the CSV data has a header. Without one, or with + ``schema_overrides`` or ``with_column_names`` given to ``execute()``, + Polars does not name the columns by the header, and they are not read + as text. Returns: The Datetime dtypes keyed by the header of a CSV file. """ import polars as pl - if not has_header or self._kwargs.keys() & { - "schema_overrides", - "dtypes", - "with_column_names", - }: + if not has_header or self._kwargs.keys() & {"schema_overrides", "with_column_names"}: return {} dtypes = self._csv_dtypes return { @@ -560,13 +499,6 @@ def _create_dataframe_iterator(self) -> PolarsDataFrameIterator: """ reader: Iterator[pl.DataFrame] if not self.is_unload: - validate_execute_kwargs( - "Polars CSV reader", - self._kwargs, - False, - self._chunksize, - allow_eager_fallback=False, - ) reader = self._iter_csv_chunks() elif self._prepare_parquet_location(): self._metadata = self._read_parquet_schema() @@ -637,7 +569,6 @@ def _read_csv(self) -> pl.DataFrame: """ import polars as pl - validate_execute_kwargs("Polars CSV reader", self._kwargs, False, None) if not self._is_csv_readable(): return pl.DataFrame() @@ -648,7 +579,7 @@ def _read_csv(self) -> pl.DataFrame: has_header=has_header, schema_overrides=self._csv_dtypes, ) - if not {"schema_overrides", "dtypes"}.intersection(self._kwargs): + if "schema_overrides" not in self._kwargs: # Renamed after reading, so that Polars matches the types to the header. kwargs.pop("new_columns", None) source: str | bytes @@ -819,9 +750,7 @@ def _get_csv_params(self) -> tuple[str, bool, list[str] | None]: separator = "," has_header = True new_columns = ( - None - if {"schema_overrides", "dtypes"}.intersection(self._kwargs) - else self._kwargs.get("new_columns") + None if "schema_overrides" in self._kwargs else self._kwargs.get("new_columns") ) return separator, has_header, new_columns @@ -852,7 +781,7 @@ def _iter_csv_chunks(self) -> Iterator[pl.DataFrame]: has_header=has_header, schema_overrides=self._csv_dtypes, ) - if not {"schema_overrides", "dtypes"}.intersection(self._kwargs): + if "schema_overrides" not in self._kwargs: # Renamed after reading, so that Polars matches the types to the header. read_kwargs.pop("new_columns", None) diff --git a/pyathena/s3fs/async_cursor.py b/pyathena/s3fs/async_cursor.py index 878232c24..999a37b63 100644 --- a/pyathena/s3fs/async_cursor.py +++ b/pyathena/s3fs/async_cursor.py @@ -7,7 +7,6 @@ from multiprocessing import cpu_count from typing import Any, cast -from pyathena._kwargs import validate_kwargs from pyathena.async_cursor import AsyncCursor from pyathena.common import CursorIterator from pyathena.error import ProgrammingError @@ -229,11 +228,7 @@ def execute( >>> result_set = future.result() >>> rows = result_set.fetchall() """ - validate_kwargs( - f"{type(self).__name__}.execute", - kwargs, - ("block_size", "csv_reader", "filesystem_class"), - ) + self._validate_execute_kwargs(kwargs, ("block_size", "csv_reader", "filesystem_class")) options = ExecuteOptions.resolve( options, work_group=work_group, diff --git a/pyathena/s3fs/cursor.py b/pyathena/s3fs/cursor.py index abc1fba12..44f243a90 100644 --- a/pyathena/s3fs/cursor.py +++ b/pyathena/s3fs/cursor.py @@ -6,7 +6,6 @@ from collections.abc import Callable from typing import Any, cast -from pyathena._kwargs import validate_kwargs from pyathena.common import CursorIterator from pyathena.error import OperationalError from pyathena.model import AthenaQueryExecution @@ -178,11 +177,7 @@ def execute( >>> cursor.execute("SELECT * FROM my_table WHERE id = %(id)s", {"id": 123}) >>> rows = cursor.fetchall() """ - validate_kwargs( - f"{type(self).__name__}.execute", - kwargs, - ("block_size", "csv_reader", "filesystem_class"), - ) + self._validate_execute_kwargs(kwargs, ("block_size", "csv_reader", "filesystem_class")) self._reset_state() options = ExecuteOptions.resolve( options, diff --git a/pyathena/spark/async_cursor.py b/pyathena/spark/async_cursor.py index 85d6fdaf8..68c50bae1 100644 --- a/pyathena/spark/async_cursor.py +++ b/pyathena/spark/async_cursor.py @@ -12,7 +12,6 @@ from multiprocessing import cpu_count from typing import TYPE_CHECKING, Any, cast -from pyathena._kwargs import validate_kwargs from pyathena.model import AthenaCalculationExecution from pyathena.spark.common import SparkBaseCursor from pyathena.util import override @@ -213,7 +212,7 @@ def execute( work_group: str | None = None, **kwargs, ) -> tuple[str, "Future[AthenaQueryExecution | AthenaCalculationExecution]"]: - validate_kwargs(f"{type(self).__name__}.execute", kwargs) + self._validate_execute_kwargs(kwargs) calculation_id = self._calculate( session_id=session_id if session_id else self._session_id, code_block=operation, diff --git a/pyathena/spark/cursor.py b/pyathena/spark/cursor.py index c402e40d9..374fc5eac 100644 --- a/pyathena/spark/cursor.py +++ b/pyathena/spark/cursor.py @@ -13,7 +13,6 @@ from typing import Any, cast from pyathena import OperationalError, ProgrammingError -from pyathena._kwargs import validate_kwargs from pyathena.model import AthenaCalculationExecution, AthenaCalculationExecutionStatus from pyathena.spark.common import SparkBaseCursor, WithCalculationExecution from pyathena.util import override @@ -112,7 +111,7 @@ def execute( work_group: str | None = None, **kwargs, ) -> SparkCursor: - validate_kwargs(f"{type(self).__name__}.execute", kwargs) + self._validate_execute_kwargs(kwargs) # A failure below must not leave the previous calculation on the cursor. self._calculation_id = None self._calculation_execution = None diff --git a/tests/pyathena/aio/test_connection.py b/tests/pyathena/aio/test_connection.py index c1b2f5d70..e36d7e73f 100644 --- a/tests/pyathena/aio/test_connection.py +++ b/tests/pyathena/aio/test_connection.py @@ -7,7 +7,7 @@ from inspect import Parameter, signature from typing import Any -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest @@ -19,13 +19,17 @@ from pyathena.aio.polars.cursor import AioPolarsCursor from pyathena.aio.s3fs.cursor import AioS3FSCursor from pyathena.aio.spark.cursor import AioSparkCursor +from pyathena.converter import DefaultTypeConverter from pyathena.error import ProgrammingError +from pyathena.spark.common import SparkBaseCursor # Cursors whose arraysize is the GetQueryResults page size, capped at 1000. PAGED_CURSORS = [AioCursor, AioDictCursor] # Cursors whose arraysize only sets the fetchmany() batch, with no upper limit. UNCAPPED_CURSORS = [AioArrowCursor, AioPandasCursor, AioPolarsCursor, AioS3FSCursor] CURSOR_CLASSES = PAGED_CURSORS + UNCAPPED_CURSORS + [AioSparkCursor] +# Cursors that pass execute() keyword arguments to a DataFrame reader. +READER_CURSOR_CLASSES = [AioPandasCursor, AioPolarsCursor] @pytest.mark.parametrize("factory", [pyathena.aio_connect, AioConnection.create]) @@ -93,6 +97,93 @@ async def _connection(**kwargs: Any) -> AioConnection: ) +@pytest.fixture +async def offline_connection(): + """Yield a connection whose cursors fail the test if they call AWS. + + A Spark cursor reports a started session without calling AWS. + """ + conn = await _connection() + with ( + patch.object( + conn.client, "_make_api_call", side_effect=AssertionError("unexpected AWS request") + ) as request, + patch.object(SparkBaseCursor, "_start_session", return_value="session"), + ): + yield conn + request.assert_not_called() + conn.close() + + +def _open_cursor(conn: AioConnection, cursor_class: type[Any], **kwargs: Any) -> Any: + if issubclass(cursor_class, SparkBaseCursor): + # The session was not started, so closing the cursor must not terminate it. + kwargs["terminate_session_on_close"] = False + return conn.cursor(cursor_class, **kwargs) + + +@pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) +async def test_cursor_unknown_constructor_keyword(offline_connection, cursor_class): + with pytest.raises(TypeError, match="unexpected keyword argument 'work_gruop'"): + offline_connection.cursor(cursor_class, work_gruop="typo") + offline_connection.cursor_kwargs = {"work_gruop": "typo"} + with pytest.raises(TypeError, match="unexpected keyword argument 'work_gruop'"): + offline_connection.cursor(cursor_class) + + +@pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) +async def test_cursor_accepts_connection_callbacks(offline_connection, cursor_class): + on_start_query_execution, on_poll = MagicMock(), MagicMock() + offline_connection.on_start_query_execution = on_start_query_execution + offline_connection.on_poll = on_poll + + async with _open_cursor(offline_connection, cursor_class) as cursor: + assert cursor._on_start_query_execution is on_start_query_execution + assert cursor._on_poll is on_poll + + +@pytest.mark.parametrize( + "cursor_class", [c for c in CURSOR_CLASSES if c not in READER_CURSOR_CLASSES] +) +async def test_cursor_execute_unknown_keyword(offline_connection, cursor_class): + async with _open_cursor(offline_connection, cursor_class) as cursor: + previous = MagicMock(is_closed=False) + if isinstance(cursor, SparkBaseCursor): + cursor._calculation_id = "previous" + else: + cursor._query_id = "previous" + cursor._result_set = previous + state = vars(cursor).copy() + + with pytest.raises( + TypeError, + match=rf"{cursor_class.__name__}\.execute\(\) got an unexpected keyword " + "argument 'work_gruop'", + ): + await cursor.execute("SELECT 1", work_gruop="typo") + + assert vars(cursor) == state + previous.close.assert_not_called() + + +@pytest.mark.parametrize( + ("cursor_class", "kwargs"), + [ + (AioArrowCursor, {"block_size": 64, "connect_timeout": 1, "request_timeout": 2}), + (AioS3FSCursor, {"block_size": 64, "csv_reader": None}), + (AioPandasCursor, {"chunksize": 10, "parse_dates": []}), + (AioPolarsCursor, {"chunksize": 10, "separator": "|"}), + ], +) +async def test_cursor_execute_supported_keywords(offline_connection, cursor_class, kwargs): + async with _open_cursor(offline_connection, cursor_class) as cursor: + with ( + patch.object(cursor, "_execute", side_effect=RuntimeError("query started")), + pytest.raises(RuntimeError, match="query started"), + ): + await cursor.execute("SELECT 1", **kwargs) + + class TestAioConnection: async def test_cursor_arguments_override_cursor_kwargs(self): conn = await _connection( @@ -132,3 +223,21 @@ async def test_cursor_arraysize_not_positive(self, cursor_class): with pytest.raises(ProgrammingError): conn.cursor(cursor_class, arraysize=0) + + async def test_internal_cursor_leaves_out_default_cursor_options(self): + cursor_kwargs = {"unload": True, "chunksize": 10, "schema_name": "configured"} + conn = await _connection( + cursor_class=AioPandasCursor, + cursor_kwargs=cursor_kwargs, + converter=AioPandasCursor.get_default_converter(True), + ) + + cursor = conn._internal_cursor(AioCursor) + + assert type(cursor) is AioCursor + assert cursor._schema_name == "configured" + assert type(cursor._converter) is DefaultTypeConverter + assert conn.cursor_kwargs == cursor_kwargs + # Creating the cursor directly still passes the options it does not accept. + with pytest.raises(TypeError, match="unexpected keyword argument 'unload'"): + conn.cursor(AioCursor) diff --git a/tests/pyathena/pandas/test_result_set.py b/tests/pyathena/pandas/test_result_set.py index df34db9db..93b2c9c3d 100644 --- a/tests/pyathena/pandas/test_result_set.py +++ b/tests/pyathena/pandas/test_result_set.py @@ -5,14 +5,12 @@ # # SPDX-License-Identifier: MIT -import builtins import csv import io from unittest.mock import MagicMock, PropertyMock, patch, sentinel import pandas as pd import pyarrow as pa -import pyarrow.parquet as pq import pytest from pandas.testing import assert_frame_equal @@ -22,7 +20,6 @@ PandasDataFrameIterator, _no_trunc_date, _read_csv_with_pyarrow, - validate_execute_kwargs, ) @@ -702,62 +699,3 @@ def test_read_csv_with_pyarrow_multiline_values_across_blocks(): ) assert df["id"].tolist() == [i for i, _ in rows] assert df["v"].tolist() == [v for _, v in rows] - - -@pytest.mark.parametrize("unload", [False, True]) -@pytest.mark.parametrize( - "kwargs", - [ - {}, - {"chunksize": 10}, - { - "engine": "c", - "chunksize": 10, - "block_size": 64, - "cache_type": "bytes", - "max_workers": 2, - "auto_optimize_chunksize": False, - }, - ], -) -def test_cursor_settings_skip_reader_imports(kwargs, unload): - with ( - patch("builtins.__import__", wraps=builtins.__import__) as load_module, - patch("pyathena.pandas.result_set.keyword_parameters") as reader_parameters, - ): - validate_execute_kwargs("Cursor.execute", kwargs, unload) - assert not any( - call.args and call.args[0] in {"pandas", "pyarrow.parquet"} - for call in load_module.call_args_list - ) - reader_parameters.assert_not_called() - - -def test_reader_replacement_updates_keyword_validation(): - def first_reader(source, *, old_option=None): - pass - - def second_reader(source, *, new_option=None): - pass - - with patch.object(pd, "read_csv", first_reader): - validate_execute_kwargs("Cursor.execute", {"old_option": True}, False) - with patch.object(pd, "read_csv", second_reader): - validate_execute_kwargs("Cursor.execute", {"chunksize": 10, "new_option": True}, False) - with pytest.raises(TypeError, match="unexpected keyword argument 'old_option'"): - validate_execute_kwargs("Cursor.execute", {"chunksize": 10, "old_option": True}, False) - - -def test_parquet_reader_replacement_updates_keyword_validation(): - def first_reader(source, *, old_option=None): - pass - - def second_reader(source, *, new_option=None): - pass - - with patch.object(pq, "read_table", first_reader): - validate_execute_kwargs("Cursor.execute", {"old_option": True}, True) - with patch.object(pq, "read_table", second_reader): - validate_execute_kwargs("Cursor.execute", {"new_option": True}, True) - with pytest.raises(TypeError, match="unexpected keyword argument 'old_option'"): - validate_execute_kwargs("Cursor.execute", {"old_option": True}, True) diff --git a/tests/pyathena/polars/test_result_set.py b/tests/pyathena/polars/test_result_set.py index b8a728edf..7cec99ea4 100644 --- a/tests/pyathena/polars/test_result_set.py +++ b/tests/pyathena/polars/test_result_set.py @@ -5,7 +5,6 @@ # # SPDX-License-Identifier: MIT -import builtins from datetime import datetime from unittest.mock import PropertyMock, patch @@ -14,11 +13,7 @@ from pyathena.error import OperationalError from pyathena.polars.converter import DefaultPolarsTypeConverter -from pyathena.polars.result_set import ( - AthenaPolarsResultSet, - PolarsDataFrameIterator, - validate_execute_kwargs, -) +from pyathena.polars.result_set import AthenaPolarsResultSet, PolarsDataFrameIterator _ROWS_BEFORE_FAILURE = 300_000 @@ -37,135 +32,6 @@ def _chunked_result_set() -> AthenaPolarsResultSet: class TestAthenaPolarsResultSet: - def test_managed_csv_accepts_eager_options_with_chunksize(self): - result_set = _chunked_result_set() - result_set._query_execution = None - result_set._converter = DefaultPolarsTypeConverter() - result_set._metadata = tuple( - {"Name": name, "Type": dtype, "Precision": 0, "Scale": 0, "Nullable": "UNKNOWN"} - for name, dtype in (("a", "integer"), ("b", "varchar")) - ) - result_set._kwargs = { - "columns": ["a"], - "n_threads": 1, - "use_pyarrow": False, - "batch_size": 1024, - } - with ( - patch.object( - AthenaPolarsResultSet, "_fetch_all_rows_as_csv", return_value=b"a,b\n1,x\n2,y\n" - ), - patch.object( - AthenaPolarsResultSet, - "_csv_storage_options", - new_callable=PropertyMock, - return_value={}, - ), - ): - frame = result_set._read_csv() - assert frame.to_dict(as_series=False) == {"a": [1, 2]} - - @pytest.mark.parametrize("reader", ["_read_csv", "_create_dataframe_iterator"]) - def test_csv_rejects_options_for_other_actual_reader(self, reader): - result_set = _chunked_result_set() - result_set._query_execution = None - result_set._unload = False - key = "include_file_paths" if reader == "_read_csv" else "columns" - result_set._kwargs = {key: "a"} - for _ in range(2): - with pytest.raises(TypeError, match=f"unexpected keyword argument '{key}'"): - ( - result_set._read_csv() - if reader == "_read_csv" - else result_set._create_dataframe_iterator() - ) - - @pytest.mark.parametrize("reader", ["_read_csv", "_iter_csv_chunks"]) - @pytest.mark.parametrize("rename", [False, True]) - def test_csv_dtypes_alias_replaces_inferred_schema(self, tmp_path, reader, rename): - path = tmp_path / "result.csv" - path.write_text("a\n001\n002\n") - result_set = _chunked_result_set() - name = "z" if rename else "a" - result_set._kwargs = {"dtypes": {name: pl.String}} - if rename: - result_set._kwargs["new_columns"] = [name] - with ( - patch.object( - AthenaPolarsResultSet, - "output_location", - new_callable=PropertyMock, - return_value=str(path), - ), - patch.object( - AthenaPolarsResultSet, - "_csv_dtypes", - new_callable=PropertyMock, - return_value={"a": pl.Int64}, - ), - patch.object( - AthenaPolarsResultSet, - "_csv_storage_options", - new_callable=PropertyMock, - return_value={}, - ), - patch.object( - AthenaPolarsResultSet, - "_parquet_storage_options", - new_callable=PropertyMock, - return_value={}, - ), - patch.object(AthenaPolarsResultSet, "_is_csv_readable", return_value=True), - pytest.warns(DeprecationWarning, match="dtypes"), - ): - frame = ( - result_set._read_csv() - if reader == "_read_csv" - else pl.concat(list(result_set._iter_csv_chunks())) - ) - assert frame.to_dict(as_series=False) == {name: ["001", "002"]} - - @pytest.mark.parametrize("reader", ["_read_csv", "_iter_csv_chunks"]) - def test_csv_keeps_deprecated_row_index_options(self, tmp_path, reader): - path = tmp_path / "result.csv" - path.write_text("a\n1\n2\n") - result_set = _chunked_result_set() - result_set._kwargs = {"row_count_name": "row_number", "row_count_offset": 2} - with ( - patch.object( - AthenaPolarsResultSet, - "output_location", - new_callable=PropertyMock, - return_value=str(path), - ), - patch.object( - AthenaPolarsResultSet, - "_csv_dtypes", - new_callable=PropertyMock, - return_value={"a": pl.Int64}, - ), - patch.object( - AthenaPolarsResultSet, - "_csv_storage_options", - new_callable=PropertyMock, - return_value={}, - ), - patch.object( - AthenaPolarsResultSet, - "_parquet_storage_options", - new_callable=PropertyMock, - return_value={}, - ), - patch.object(AthenaPolarsResultSet, "_is_csv_readable", return_value=True), - pytest.warns(DeprecationWarning, match="row_count_(name|offset)"), - ): - frame = ( - result_set._read_csv() - if reader == "_read_csv" - else pl.concat(list(result_set._iter_csv_chunks())) - ) - assert frame.to_dict(as_series=False) == {"row_number": [2, 3], "a": [1, 2]} - def test_iter_csv_chunks_raises_when_read_fails_partway(self, tmp_path): """A CSV read that fails partway through the data raises instead of ending early.""" path = tmp_path / "result.csv" @@ -335,7 +201,7 @@ def test_storage_options_replace_defaults(self, reader, function): ), patch.object(AthenaPolarsResultSet, "_is_csv_readable", return_value=True), patch.object(AthenaPolarsResultSet, "_prepare_parquet_location", return_value=True), - patch(f"polars.{function}", autospec=True) as read, + patch(f"polars.{function}") as read, patch("pyathena.polars.result_set.to_column_info"), ): result = getattr(result_set, reader)() @@ -347,30 +213,17 @@ def test_storage_options_replace_defaults(self, reader, function): ("kwargs", "expected"), [ ({}, {"t": [datetime(2020, 1, 2, 3, 4, 5, 123456), None]}), - ({"dtypes": {"t": pl.Datetime("ms")}}, OperationalError), - ({"schema_overrides": {"t": pl.Datetime("ms")}}, OperationalError), - ({"dtypes": {"v": pl.Int64}}, OperationalError), - ({"schema_overrides": {"v": pl.Int64}}, OperationalError), - ( - {"dtypes": {"t2": pl.Datetime("us")}, "new_columns": ["t2", "v2"]}, - OperationalError, - ), - ( - {"schema_overrides": {"t2": pl.Datetime("us")}, "new_columns": ["t2", "v2"]}, - OperationalError, - ), ({"columns": ["v"]}, {}), ({"new_columns": ["t2", "v2"]}, {"t2": [datetime(2020, 1, 2, 3, 4, 5, 123456), None]}), - ({"with_column_names": lambda names: [n.upper() for n in names]}, TypeError), + ({"with_column_names": lambda names: [n.upper() for n in names]}, OperationalError), ], ) def test_read_csv_truncates_timestamps(self, kwargs, expected): """Timestamps that fail to parse are read again as text and truncated. - ``with_column_names`` is a scan-only option and is rejected by the eager reader. - The ``schema_overrides`` option and its ``dtypes`` alias disable timestamp - retries and preserve reader errors. - No AWS calls; the GetQueryResults rows are mocked. + With ``with_column_names`` given to execute(), Polars renames the columns as it + reads them, so they are not, and the read fails as before. No AWS calls; the + GetQueryResults rows are mocked. """ result_set = AthenaPolarsResultSet.__new__(AthenaPolarsResultSet) # bypass __init__ result_set._query_execution = None @@ -391,13 +244,7 @@ def test_read_csv_truncates_timestamps(self, kwargs, expected): ), ): if expected is OperationalError: - with pytest.raises(OperationalError, match="could not parse"): - result_set._read_csv() - return - if expected is TypeError: - with pytest.raises( - TypeError, match="unexpected keyword argument 'with_column_names'" - ): + with pytest.raises(OperationalError): result_set._read_csv() return df = result_set._read_csv() @@ -424,44 +271,3 @@ def test_close_stops_iteration(self, reader): df_iter = PolarsDataFrameIterator(reader, {}, ["a"]) df_iter.close() assert list(df_iter) == [] - - -@pytest.mark.parametrize("unload", [False, True]) -@pytest.mark.parametrize( - "kwargs", - [ - {}, - {"chunksize": 10}, - {"chunksize": 10, "block_size": 64, "cache_type": "bytes", "max_workers": 2}, - ], -) -def test_cursor_settings_skip_reader_imports(kwargs, unload): - with ( - patch("builtins.__import__", wraps=builtins.__import__) as load_module, - patch("pyathena.polars.result_set.keyword_parameters") as reader_parameters, - ): - validate_execute_kwargs("Cursor.execute", kwargs, unload, None) - assert not any( - call.args and call.args[0] in {"polars", "pyarrow.parquet"} - for call in load_module.call_args_list - ) - reader_parameters.assert_not_called() - - -def test_reader_replacement_updates_keyword_validation(): - def first_reader(source, *, old_option=None): - pass - - def second_reader(source, *, new_option=None): - pass - - with patch.object(pl, "read_csv", first_reader): - validate_execute_kwargs("Cursor.execute", {"old_option": True}, False, None) - with patch.object(pl, "read_csv", second_reader): - validate_execute_kwargs( - "Cursor.execute", {"chunksize": 10, "new_option": True}, False, None - ) - with pytest.raises(TypeError, match="unexpected keyword argument 'old_option'"): - validate_execute_kwargs( - "Cursor.execute", {"chunksize": 10, "old_option": True}, False, None - ) diff --git a/tests/pyathena/sqlalchemy/test_base.py b/tests/pyathena/sqlalchemy/test_base.py index 9b928f520..0bcb5f661 100644 --- a/tests/pyathena/sqlalchemy/test_base.py +++ b/tests/pyathena/sqlalchemy/test_base.py @@ -8,6 +8,7 @@ from decimal import Decimal from pathlib import Path from types import SimpleNamespace +from unittest.mock import MagicMock from urllib.parse import quote_plus import numpy as np @@ -24,11 +25,15 @@ from sqlalchemy.sql.selectable import TextualSelect from sqlalchemy.util import PluginLoader -from pyathena.aio.sqlalchemy.base import AthenaAioDialect +from pyathena.aio.connection import AioConnection +from pyathena.aio.cursor import AioCursor +from pyathena.aio.pandas.cursor import AioPandasCursor +from pyathena.aio.sqlalchemy.base import AsyncAdapt_pyathena_connection, AthenaAioDialect from pyathena.connection import Connection from pyathena.cursor import Cursor from pyathena.error import DatabaseError, OperationalError from pyathena.formatter import DefaultParameterFormatter +from pyathena.pandas.cursor import PandasCursor from pyathena.sqlalchemy.base import AthenaDialect from pyathena.sqlalchemy.compiler import AthenaTypeCompiler from pyathena.sqlalchemy.rest import AthenaRestDialect @@ -187,6 +192,36 @@ def test_conn_str_unload(self, driver, value, expected): assert opts["cursor_kwargs"]["unload"] is expected assert "unload" not in opts + @pytest.mark.parametrize( + ("connection_class", "cursor_class", "expected"), + [(Connection, PandasCursor, Cursor), (AioConnection, AioPandasCursor, AioCursor)], + ) + def test_internal_cursor_leaves_out_default_cursor_options( + self, connection_class, cursor_class, expected + ): + conn = connection_class( + region_name="us-west-2", + s3_staging_dir="s3://bucket/path/", + aws_access_key_id="access_key", + aws_secret_access_key="secret_key", + cursor_class=cursor_class, + cursor_kwargs={"unload": True, "chunksize": 10}, + ) + driver_connection = ( + AsyncAdapt_pyathena_connection(MagicMock(), conn) + if connection_class is AioConnection + else conn + ) + try: + cursor = AthenaDialect._internal_cursor( + SimpleNamespace(driver_connection=driver_connection) + ) + if connection_class is AioConnection: + cursor = cursor._cursor + assert type(cursor) is expected + finally: + conn.close() + @pytest.mark.parametrize("dialect_class", [AthenaDialect, AthenaAioDialect]) def test_type_compiler(self, dialect_class): # SQLAlchemy 2.0 builds the type compiler from type_compiler_cls. A legacy diff --git a/tests/pyathena/test_connection.py b/tests/pyathena/test_connection.py index 6c5acf5c9..4335d8877 100644 --- a/tests/pyathena/test_connection.py +++ b/tests/pyathena/test_connection.py @@ -10,7 +10,7 @@ from concurrent.futures import ThreadPoolExecutor from inspect import Parameter, signature from typing import Any -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest from botocore.config import Config @@ -51,6 +51,8 @@ AsyncS3FSCursor, ] CURSOR_CLASSES = PAGED_CURSORS + UNCAPPED_CURSORS + [SparkCursor, AsyncSparkCursor] +# Cursors that pass execute() keyword arguments to a DataFrame reader. +READER_CURSOR_CLASSES = [PandasCursor, PolarsCursor, AsyncPandasCursor, AsyncPolarsCursor] @pytest.mark.parametrize("factory", [pyathena.connect, Connection]) @@ -128,6 +130,104 @@ def isolated_aws_config(monkeypatch, tmp_path): monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(tmp_path / "credentials")) +@pytest.fixture +def offline_connection(): + """Yield a connection whose cursors fail the test if they call AWS. + + A Spark cursor reports a started session without calling AWS. + """ + conn = _connection() + with ( + patch.object( + conn.client, "_make_api_call", side_effect=AssertionError("unexpected AWS request") + ) as request, + patch.object(SparkBaseCursor, "_start_session", return_value="session"), + ): + yield conn + request.assert_not_called() + conn.close() + + +def _open_cursor(conn: Connection[Any], cursor_class: type[Any], **kwargs: Any) -> Any: + if issubclass(cursor_class, SparkBaseCursor): + # The session was not started, so closing the cursor must not terminate it. + kwargs["terminate_session_on_close"] = False + return conn.cursor(cursor_class, **kwargs) + + +@pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) +def test_cursor_unknown_constructor_keyword(offline_connection, cursor_class): + with pytest.raises(TypeError, match="unexpected keyword argument 'work_gruop'"): + offline_connection.cursor(cursor_class, work_gruop="typo") + offline_connection.cursor_kwargs = {"work_gruop": "typo"} + with pytest.raises(TypeError, match="unexpected keyword argument 'work_gruop'"): + offline_connection.cursor(cursor_class) + + +@pytest.mark.parametrize("cursor_class", CURSOR_CLASSES) +def test_cursor_accepts_connection_callbacks(offline_connection, cursor_class): + on_start_query_execution, on_poll = MagicMock(), MagicMock() + offline_connection.on_start_query_execution = on_start_query_execution + offline_connection.on_poll = on_poll + + with _open_cursor(offline_connection, cursor_class) as cursor: + assert cursor._on_start_query_execution is on_start_query_execution + assert cursor._on_poll is on_poll + + +@pytest.mark.parametrize( + "cursor_class", [c for c in CURSOR_CLASSES if c not in READER_CURSOR_CLASSES] +) +def test_cursor_execute_unknown_keyword(offline_connection, cursor_class): + with _open_cursor(offline_connection, cursor_class) as cursor: + previous = MagicMock(is_closed=False) + if isinstance(cursor, SparkBaseCursor): + cursor._calculation_id = "previous" + elif isinstance(cursor, WithResultSet): + cursor._query_id = "previous" + cursor._result_set = previous + state = vars(cursor).copy() + + with pytest.raises( + TypeError, + match=rf"{cursor_class.__name__}\.execute\(\) got an unexpected keyword " + "argument 'work_gruop'", + ): + cursor.execute("SELECT 1", work_gruop="typo") + + assert vars(cursor) == state + previous.close.assert_not_called() + + +@pytest.mark.parametrize( + ("cursor_class", "kwargs"), + [ + (cursor_class, {"block_size": 64, "connect_timeout": 1, "request_timeout": 2}) + for cursor_class in (ArrowCursor, AsyncArrowCursor) + ] + + [ + (cursor_class, {"block_size": 64, "csv_reader": None, "filesystem_class": None}) + for cursor_class in (S3FSCursor, AsyncS3FSCursor) + ] + + [ + (cursor_class, {"chunksize": 10, "parse_dates": []}) + for cursor_class in (PandasCursor, AsyncPandasCursor) + ] + + [ + (cursor_class, {"chunksize": 10, "separator": "|"}) + for cursor_class in (PolarsCursor, AsyncPolarsCursor) + ], +) +def test_cursor_execute_supported_keywords(offline_connection, cursor_class, kwargs): + cursor = _open_cursor(offline_connection, cursor_class) + with ( + cursor, + patch.object(cursor, "_execute", side_effect=RuntimeError("query started")), + pytest.raises(RuntimeError, match="query started"), + ): + cursor.execute("SELECT 1", **kwargs) + + class TestConnection: @pytest.mark.parametrize( ("key", "configured", "explicit"), @@ -278,3 +378,29 @@ def test_close_closes_built_clients(self): athena_close.assert_called_once_with() glue_close.assert_called_once_with() s3_close.assert_called_once_with() + + def test_internal_cursor_leaves_out_default_cursor_options(self): + cursor_kwargs = {"unload": True, "chunksize": 10, "schema_name": "configured"} + conn = _connection( + cursor_class=PandasCursor, + cursor_kwargs=cursor_kwargs, + converter=PandasCursor.get_default_converter(True), + ) + + cursor = conn._internal_cursor(Cursor) + + assert type(cursor) is Cursor + assert cursor._schema_name == "configured" + assert type(cursor._converter) is DefaultTypeConverter + assert conn.cursor_kwargs == cursor_kwargs + # Creating the cursor directly still passes the options it does not accept. + with pytest.raises(TypeError, match="unexpected keyword argument 'unload'"): + conn.cursor(Cursor) + + def test_internal_cursor_unknown_cursor_kwargs(self): + conn = _connection( + cursor_class=PandasCursor, cursor_kwargs={"unload": True, "work_gruop": "typo"} + ) + + with pytest.raises(TypeError, match="unexpected keyword argument 'work_gruop'"): + conn._internal_cursor(Cursor) diff --git a/tests/pyathena/test_kwargs.py b/tests/pyathena/test_kwargs.py deleted file mode 100644 index b93082e5f..000000000 --- a/tests/pyathena/test_kwargs.py +++ /dev/null @@ -1,630 +0,0 @@ -# Copyright 2026 The PyAthena authors -# -# Licensed under the MIT License. -# See LICENSE or https://opensource.org/licenses/MIT. -# -# SPDX-License-Identifier: MIT - -from inspect import iscoroutinefunction, signature -from types import SimpleNamespace -from unittest.mock import MagicMock, patch - -import pytest - -from pyathena._kwargs import constructor_keyword_parameters, keyword_parameters -from pyathena.aio.arrow.cursor import AioArrowCursor -from pyathena.aio.connection import AioConnection -from pyathena.aio.cursor import AioCursor, AioDictCursor -from pyathena.aio.pandas.cursor import AioPandasCursor -from pyathena.aio.polars.cursor import AioPolarsCursor -from pyathena.aio.s3fs.cursor import AioS3FSCursor -from pyathena.aio.spark.cursor import AioSparkCursor -from pyathena.aio.sqlalchemy.base import AsyncAdapt_pyathena_connection -from pyathena.arrow.async_cursor import AsyncArrowCursor -from pyathena.arrow.cursor import ArrowCursor -from pyathena.async_cursor import AsyncCursor, AsyncDictCursor -from pyathena.common import BaseCursor -from pyathena.connection import Connection -from pyathena.cursor import Cursor, DictCursor -from pyathena.error import ProgrammingError -from pyathena.formatter import DefaultParameterFormatter -from pyathena.pandas.async_cursor import AsyncPandasCursor -from pyathena.pandas.cursor import PandasCursor -from pyathena.polars.async_cursor import AsyncPolarsCursor -from pyathena.polars.cursor import PolarsCursor -from pyathena.s3fs.async_cursor import AsyncS3FSCursor -from pyathena.s3fs.cursor import S3FSCursor -from pyathena.spark.async_cursor import AsyncSparkCursor -from pyathena.spark.common import SparkBaseCursor -from pyathena.spark.cursor import SparkCursor -from pyathena.sqlalchemy.base import AthenaDialect -from pyathena.util import RetryConfig - -CURSORS = [ - Cursor, - DictCursor, - AsyncCursor, - AsyncDictCursor, - AioCursor, - AioDictCursor, - PandasCursor, - AsyncPandasCursor, - AioPandasCursor, - ArrowCursor, - AsyncArrowCursor, - AioArrowCursor, - PolarsCursor, - AsyncPolarsCursor, - AioPolarsCursor, - S3FSCursor, - AsyncS3FSCursor, - AioS3FSCursor, - SparkCursor, - AsyncSparkCursor, - AioSparkCursor, -] - - -def test_keyword_parameters_reuses_inspection_without_sharing_mutable_names(): - def reader(source, /, value=None, *args, option=None, **kwargs): - pass - - with patch("pyathena._kwargs.signature", wraps=signature) as inspect_signature: - names = keyword_parameters(reader) - assert names == {"value", "option"} - names.remove("value") - names.add("unsupported") - assert keyword_parameters(reader) == {"value", "option"} - inspect_signature.assert_called_once_with(reader) - - -def test_keyword_parameters_accepts_unhashable_callable(): - class Reader: - __hash__ = None - - def __call__(self, *, setting=None): - pass - - reader = Reader() - assert keyword_parameters(reader) == {"setting"} - assert keyword_parameters(reader) == {"setting"} - - -def test_constructor_keyword_parameters_observes_replaced_constructor(): - class Parent: - def __init__(self, *, shared=None): - pass - - class Child(Parent): - def __init__(self, *, original=None): - pass - - def replacement(self, *, updated=None): - pass - - assert constructor_keyword_parameters(Child) == {"shared", "original"} - Child.__init__ = replacement - assert constructor_keyword_parameters(Child) == {"shared", "updated"} - - -@pytest.mark.parametrize( - ("operation", "before_preparation", "after_preparation"), - [ - (" \n sElEcT 1", True, False), - ("\u2003WITH t AS (SELECT 1) SELECT * FROM t", True, False), - ("\u001cUNLOAD (SELECT 1)", True, True), - ("\u017felect 1", True, False), - ("w\u0131th t AS (SELECT 1) SELECT * FROM t", True, False), - ("w\u0130th t AS (SELECT 1) SELECT * FROM t", False, False), - ("\ufeffSELECT 1", False, False), - ("SELECTED", True, False), - ("DESCRIBE t", False, False), - ("", False, False), - (" \n\t", False, False), - (None, False, False), - ("UNLOAD " + "x" * 100_000, True, True), - ], -) -def test_unload_reader_classification(operation, before_preparation, after_preparation): - cursor = SimpleNamespace(_unload=True) - assert BaseCursor._is_unload_query(cursor, operation, prepared=False) is before_preparation - assert BaseCursor._is_unload_query(cursor, operation) is after_preparation - cursor._unload = False - assert BaseCursor._is_unload_query(cursor, operation, prepared=False) is False - assert BaseCursor._is_unload_query(cursor, operation) is False - - -@pytest.fixture -def offline_connection(monkeypatch): - conn = Connection( - region_name="us-east-1", - s3_staging_dir="s3://bucket/path/", - aws_access_key_id="access_key", - aws_secret_access_key="secret_key", - ) - request = MagicMock(side_effect=AssertionError("An offline test made an AWS request")) - monkeypatch.setattr(conn.client, "_make_api_call", request) - monkeypatch.setattr(SparkBaseCursor, "_start_session", MagicMock(return_value="session")) - yield conn - request.assert_not_called() - conn.close() - - -def _constructor_kwargs(conn, cursor_class): - kwargs = { - "connection": conn, - "converter": cursor_class.get_default_converter(), - "formatter": DefaultParameterFormatter(), - "retry_config": RetryConfig(), - "s3_staging_dir": conn.s3_staging_dir, - "schema_name": conn.schema_name, - "catalog_name": conn.catalog_name, - "work_group": conn.work_group, - "poll_interval": 0, - "encryption_option": None, - "kms_key": None, - "kill_on_interrupt": True, - "result_reuse_enable": False, - "result_reuse_minutes": 60, - } - if issubclass(cursor_class, SparkBaseCursor): - kwargs["terminate_session_on_close"] = False - return kwargs - - -async def _close(cursor): - if iscoroutinefunction(cursor.close): - await cursor.close() - else: - cursor.close() - - -async def _execute(cursor, operation="SELECT 1", **kwargs): - if iscoroutinefunction(cursor.execute): - return await cursor.execute(operation, **kwargs) - return cursor.execute(operation, **kwargs) - - -@pytest.mark.parametrize( - "cursor_class", - [ - PandasCursor, - AsyncPandasCursor, - AioPandasCursor, - PolarsCursor, - AsyncPolarsCursor, - AioPolarsCursor, - ], -) -@pytest.mark.parametrize("preparation_failure", [False, True]) -async def test_execute_without_reader_keywords_skips_classification( - offline_connection, cursor_class, preparation_failure -): - cursor = offline_connection.cursor(cursor_class, unload=True) - try: - target = "_prepare_unload" if preparation_failure else "_execute" - with ( - patch.object(cursor, "_is_unload_query", side_effect=AssertionError("SQL was scanned")), - patch.object(cursor, target, side_effect=ProgrammingError("execution stopped")), - pytest.raises(ProgrammingError, match="execution stopped"), - ): - await _execute(cursor) - finally: - await _close(cursor) - - -@pytest.mark.parametrize("cursor_class", CURSORS) -@pytest.mark.parametrize("source", ["constructor", "cursor", "defaults"]) -def test_unknown_constructor_keyword(offline_connection, cursor_class, source): - if source == "constructor": - factory = cursor_class - kwargs = _constructor_kwargs(offline_connection, cursor_class) - else: - factory = offline_connection.cursor - kwargs = {"cursor": cursor_class} - if source == "defaults": - offline_connection.cursor_kwargs = {"work_gruop": "typo"} - if source != "defaults": - kwargs["work_gruop"] = "typo" - with pytest.raises(TypeError, match="unexpected keyword argument 'work_gruop'"): - factory(**kwargs) - - -@pytest.mark.parametrize("cursor_class", CURSORS) -async def test_shared_constructor_callbacks(offline_connection, cursor_class): - on_start, on_poll = MagicMock(), MagicMock() - offline_connection.on_start_query_execution = on_start - offline_connection.on_poll = on_poll - kwargs = ( - {"terminate_session_on_close": False} if issubclass(cursor_class, SparkBaseCursor) else {} - ) - cursor = offline_connection.cursor(cursor_class, **kwargs) - try: - assert cursor._on_start_query_execution is on_start - assert cursor._on_poll is on_poll - finally: - await _close(cursor) - - -@pytest.mark.parametrize("cursor_class", CURSORS) -async def test_unknown_execute_keyword_preserves_state(offline_connection, cursor_class): - cursor = cursor_class(**_constructor_kwargs(offline_connection, cursor_class)) - cursor._query_id = "previous" - cursor._result_set = MagicMock(is_closed=False) - if isinstance(cursor, SparkBaseCursor): - cursor._calculation_id = "previous-calculation" - cursor._calculation_execution = MagicMock() - before = vars(cursor).copy() - try: - with pytest.raises(TypeError, match="unexpected keyword argument 'work_gruop'"): - await _execute(cursor, work_gruop="typo") - assert vars(cursor) == before - cursor._result_set.close.assert_not_called() - finally: - await _close(cursor) - - -@pytest.mark.parametrize( - "cursor_class", - [ - PandasCursor, - AsyncPandasCursor, - AioPandasCursor, - PolarsCursor, - AsyncPolarsCursor, - AioPolarsCursor, - ], -) -@pytest.mark.parametrize( - "failure", ["empty", "missing_staging", "invalid_options", "invalid_operation"] -) -@pytest.mark.parametrize("unknown_keyword", [False, True]) -async def test_preparation_failure_resets_state_except_unknown_keyword( - offline_connection, cursor_class, failure, unknown_keyword -): - cursor = offline_connection.cursor(cursor_class, unload=True) - previous_result = MagicMock(is_closed=False) - cursor._query_id = "previous" - cursor._result_set = previous_result - operation, options = "SELECT 1", {} - if failure == "empty": - operation = " " - elif failure == "missing_staging": - cursor._s3_staging_dir = None - elif failure == "invalid_options": - options["options"] = object() - else: - operation = object() - if unknown_keyword: - options["work_gruop"] = "typo" - before = vars(cursor).copy() - error = ( - TypeError - if unknown_keyword - else ( - AttributeError - if failure in ("invalid_options", "invalid_operation") - else ProgrammingError - ) - ) - message = ( - "unexpected keyword argument 'work_gruop'" - if unknown_keyword - else ( - "merge" - if failure == "invalid_options" - else ( - "strip" - if failure == "invalid_operation" - else "Query is none|s3_staging_dir is required" - ) - ) - ) - try: - with pytest.raises(error, match=message): - await _execute(cursor, operation=operation, **options) - if unknown_keyword or cursor_class.__name__.startswith("Async"): - assert vars(cursor) == before - previous_result.close.assert_not_called() - else: - assert cursor.query_id is None - assert cursor.result_set is None - previous_result.close.assert_called_once() - finally: - await _close(cursor) - - -@pytest.mark.parametrize( - "cursor_class", - [ - PandasCursor, - AsyncPandasCursor, - AioPandasCursor, - PolarsCursor, - AsyncPolarsCursor, - AioPolarsCursor, - ], -) -@pytest.mark.parametrize( - ("operation", "unload_reader"), - [ - (" \n sElEcT 1", True), - (" \n wItH t AS (SELECT 1) SELECT * FROM t", True), - ( - " \n uNlOaD (SELECT 1) TO 's3://bucket/path/' WITH (format='PARQUET')", - True, - ), - (" \n dEsCrIbE t", False), - ], -) -async def test_preparation_failure_preserves_supported_reader_keywords( - offline_connection, cursor_class, operation, unload_reader -): - cursor = offline_connection.cursor(cursor_class, unload=True) - cursor._s3_staging_dir = None - if "Polars" in cursor_class.__name__: - options = {"parallel": "none"} if unload_reader else {"separator": "|"} - else: - options = {"use_threads": False} if unload_reader else {"sep": "|"} - try: - with pytest.raises(ProgrammingError, match="s3_staging_dir is required"): - await _execute(cursor, operation=operation, **options) - finally: - await _close(cursor) - - -@pytest.mark.parametrize( - ("cursor_class", "settings", "options"), - [(cls, {}, {"parse_dates": []}) for cls in (PandasCursor, AsyncPandasCursor, AioPandasCursor)] - + [ - (cls, {"unload": True}, {"use_threads": False}) - for cls in (PandasCursor, AsyncPandasCursor, AioPandasCursor) - ] - + [(cls, {}, {"separator": "|"}) for cls in (PolarsCursor, AsyncPolarsCursor, AioPolarsCursor)] - + [ - (cls, {"unload": True}, {"parallel": "none"}) - for cls in (PolarsCursor, AsyncPolarsCursor, AioPolarsCursor) - ] - + [ - (cls, {}, {"block_size": 64, "connect_timeout": None, "request_timeout": 2}) - for cls in (ArrowCursor, AsyncArrowCursor, AioArrowCursor) - ] - + [(cls, {}, {"block_size": 64}) for cls in (S3FSCursor, AsyncS3FSCursor, AioS3FSCursor)], -) -async def test_supported_execute_keywords(offline_connection, cursor_class, settings, options): - cursor = offline_connection.cursor(cursor_class, **settings) - try: - with ( - patch.object(cursor, "_execute", side_effect=RuntimeError("execution reached")), - pytest.raises(RuntimeError, match="execution reached"), - ): - await _execute(cursor, **options) - finally: - await _close(cursor) - - -@pytest.mark.parametrize( - "cursor_class", - [ - PandasCursor, - AsyncPandasCursor, - AioPandasCursor, - PolarsCursor, - AsyncPolarsCursor, - AioPolarsCursor, - ], -) -@pytest.mark.parametrize("unload", [False, True]) -async def test_keyword_rejected_for_other_reader_mode(offline_connection, cursor_class, unload): - cursor = offline_connection.cursor(cursor_class, unload=unload) - if "Pandas" in cursor_class.__name__: - key = "sep" if unload else "use_threads" - else: - key = "separator" if unload else "parallel" - try: - with pytest.raises(TypeError, match=f"unexpected keyword argument '{key}'"): - await _execute(cursor, **{key: "|"}) - finally: - await _close(cursor) - - -@pytest.mark.parametrize( - "cursor_class", - [ - PandasCursor, - AsyncPandasCursor, - AioPandasCursor, - PolarsCursor, - AsyncPolarsCursor, - AioPolarsCursor, - ], -) -async def test_reader_input_cannot_replace_query_results(offline_connection, cursor_class): - cursor = offline_connection.cursor(cursor_class) - key = "filepath_or_buffer" if "Pandas" in cursor_class.__name__ else "source" - try: - with pytest.raises(TypeError, match=f"unexpected keyword argument '{key}'"): - await _execute(cursor, **{key: "other.csv"}) - finally: - await _close(cursor) - - -@pytest.mark.parametrize("cursor_class", [PolarsCursor, AsyncPolarsCursor, AioPolarsCursor]) -@pytest.mark.parametrize("chunksize", [None, 10]) -async def test_polars_selects_effective_chunk_reader(offline_connection, cursor_class, chunksize): - cursor = offline_connection.cursor(cursor_class, chunksize=10) - try: - with patch.object(cursor, "_execute", side_effect=RuntimeError("execution reached")): - error = TypeError if chunksize is None else RuntimeError - message = "include_file_paths" if chunksize is None else "execution reached" - with pytest.raises(error, match=message): - await _execute(cursor, chunksize=chunksize, include_file_paths="path") - finally: - await _close(cursor) - - -@pytest.mark.parametrize( - ("cursor_class", "options"), - [(cls, {"nrows": 10}) for cls in (PandasCursor, AsyncPandasCursor, AioPandasCursor)] - + [ - (cls, {"separator": "\t", "null_values": "NULL"}) - for cls in (PolarsCursor, AsyncPolarsCursor, AioPolarsCursor) - ], -) -async def test_unload_non_select_uses_csv_keywords(offline_connection, cursor_class, options): - cursor = offline_connection.cursor(cursor_class, unload=True) - try: - with ( - patch.object(cursor, "_execute", side_effect=RuntimeError("execution reached")), - pytest.raises(RuntimeError, match="execution reached"), - ): - await _execute(cursor, operation="DESCRIBE t", **options) - finally: - await _close(cursor) - - -@pytest.mark.parametrize("cursor_class", [PolarsCursor, AsyncPolarsCursor, AioPolarsCursor]) -async def test_polars_chunked_csv_allows_managed_reader_options(offline_connection, cursor_class): - cursor = offline_connection.cursor(cursor_class, chunksize=10) - try: - with ( - patch.object(cursor, "_execute", side_effect=RuntimeError("execution reached")), - pytest.raises(RuntimeError, match="execution reached"), - ): - await _execute(cursor, columns=["a"], n_threads=1, use_pyarrow=False, batch_size=1024) - finally: - await _close(cursor) - - -@pytest.mark.parametrize("cursor_class", [PolarsCursor, AsyncPolarsCursor, AioPolarsCursor]) -@pytest.mark.parametrize("unload", [False, True]) -async def test_polars_accepts_reader_aliases(offline_connection, cursor_class, unload): - cursor = offline_connection.cursor(cursor_class, unload=unload) - try: - with ( - patch.object(cursor, "_execute", side_effect=RuntimeError("execution reached")), - pytest.raises(RuntimeError, match="execution reached"), - ): - await _execute(cursor, row_count_name="row_number", row_count_offset=2) - finally: - await _close(cursor) - - -@pytest.mark.parametrize( - "cursor_class", - [PandasCursor, ArrowCursor, PolarsCursor, AioPandasCursor, AioArrowCursor, AioPolarsCursor], -) -async def test_internal_cursor_isolates_backend_defaults(cursor_class): - is_async = cursor_class.__name__.startswith("Aio") - cls = AioConnection if is_async else Connection - defaults = {"unload": True, "schema_name": "configured", "arraysize": 25} - conn = cls( - region_name="us-east-1", - s3_staging_dir="s3://bucket/path/", - aws_access_key_id="access_key", - aws_secret_access_key="secret_key", - cursor_class=cursor_class, - cursor_kwargs=defaults, - converter=cursor_class.get_default_converter(True), - ) - adapted = AsyncAdapt_pyathena_connection(MagicMock(), conn) if is_async else conn - raw_connection = SimpleNamespace(driver_connection=adapted) - cursor = None - try: - internal = AthenaDialect._internal_cursor(raw_connection) - cursor = internal._cursor if is_async else internal - assert type(cursor) is (AioCursor if is_async else Cursor) - assert cursor._schema_name == "configured" - assert cursor.arraysize == 25 - assert type(cursor._converter) is type(Cursor.get_default_converter()) - assert conn.cursor_kwargs == defaults - # Public creation still rejects settings unsupported by the requested class. - with pytest.raises(TypeError, match="unexpected keyword argument 'unload'"): - conn.cursor(AioCursor if is_async else Cursor) - finally: - if cursor is not None: - await _close(cursor) - conn.close() - - -def test_internal_cursor_does_not_hide_unknown_defaults(offline_connection): - offline_connection.cursor_class = PandasCursor - offline_connection.cursor_kwargs = {"unload": True, "work_gruop": "typo"} - with pytest.raises(TypeError, match="unexpected keyword argument 'work_gruop'"): - offline_connection._internal_cursor(Cursor) - - -@pytest.mark.parametrize( - "cursor_class", - [ - PandasCursor, - AsyncPandasCursor, - AioPandasCursor, - PolarsCursor, - AsyncPolarsCursor, - AioPolarsCursor, - ], -) -@pytest.mark.parametrize("valid_for_prepared_reader", [False, True]) -async def test_validation_uses_prepared_operation( - offline_connection, cursor_class, valid_for_prepared_reader -): - cursor = offline_connection.cursor(cursor_class, unload=True) - previous_result = MagicMock(is_closed=False) - cursor._query_id = "previous" - cursor._result_set = previous_result - before = vars(cursor).copy() - if "Pandas" in cursor_class.__name__: - kwargs = {"sep": "|"} if valid_for_prepared_reader else {"use_threads": False} - else: - kwargs = {"separator": "|"} if valid_for_prepared_reader else {"parallel": "none"} - try: - with ( - patch.object(cursor, "_prepare_unload", return_value=("DESCRIBE t", None)), - patch.object( - cursor, "_execute", side_effect=RuntimeError("execution reached") - ) as execute, - pytest.raises( - RuntimeError if valid_for_prepared_reader else TypeError, - match="execution reached" if valid_for_prepared_reader else "unexpected keyword", - ), - ): - await _execute(cursor, **kwargs) - if valid_for_prepared_reader: - assert execute.call_args.args[0] == "DESCRIBE t" - else: - execute.assert_not_called() - assert vars(cursor) == before - previous_result.close.assert_not_called() - finally: - await _close(cursor) - - -@pytest.mark.parametrize( - "cursor_class", - [ - PandasCursor, - AsyncPandasCursor, - AioPandasCursor, - PolarsCursor, - AsyncPolarsCursor, - AioPolarsCursor, - ], -) -@pytest.mark.parametrize("error", [KeyboardInterrupt, SystemExit]) -async def test_preparation_interrupt_preserves_state(offline_connection, cursor_class, error): - cursor = offline_connection.cursor(cursor_class, unload=True) - previous_result = MagicMock(is_closed=False) - cursor._query_id = "previous" - cursor._result_set = previous_result - before = vars(cursor).copy() - try: - with ( - patch.object(cursor, "_prepare_unload", side_effect=error), - pytest.raises(error), - ): - await _execute(cursor, work_gruop="typo") - assert vars(cursor) == before - previous_result.close.assert_not_called() - finally: - await _close(cursor) From a8d79da40e176d553555233c0e5fa5cbb110569d Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 10 Oct 2026 00:53:06 +0900 Subject: [PATCH 13/15] Apply only accepted cursor_kwargs to the internal cursor A custom default cursor class can consume its own options through **kwargs, which constructor introspection cannot see. Apply only the cursor_kwargs that the internal cursor class accepts. Document that cursor subclasses must not pass their own settings to the base constructor, and how pandas and Polars report unknown reader option names. Co-Authored-By: Claude Opus 5.5 --- docs/usage.md | 4 +++- pyathena/connection.py | 10 ++++------ tests/pyathena/test_connection.py | 15 +++++++++------ 3 files changed, 16 insertions(+), 13 deletions(-) diff --git a/docs/usage.md b/docs/usage.md index 2324b1b1d..f1a6a063d 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -42,7 +42,9 @@ The cursor class can still be selected positionally with `connection.cursor(Dict Unknown keyword arguments raise `TypeError` in cursor constructors and in `execute()` instead of being ignored. The connection's `cursor_kwargs` are passed to every cursor that it creates, so the requested cursor class must accept them. -The pandas and Polars cursors pass the `execute()` keyword arguments that they do not use themselves to the reader, such as `pandas.read_csv()`, so an unknown name fails when the results are read. +A cursor subclass that receives settings of its own must not pass them on to the base class constructor. +The pandas and Polars cursors pass the `execute()` keyword arguments that they do not use themselves to the reader, such as `pandas.read_csv()`. +The reader rejects an unknown name when it reads a query result file, which PyAthena raises as `OperationalError`; the name is not checked when no result file is read. ## Managed query result storage diff --git a/pyathena/connection.py b/pyathena/connection.py index b1d356902..e1218ea3d 100644 --- a/pyathena/connection.py +++ b/pyathena/connection.py @@ -656,8 +656,8 @@ def _cursor( def _internal_cursor(self, cursor: type[FunctionalCursor]) -> FunctionalCursor: """Create a cursor for queries that PyAthena runs on its own behalf. - The connection's ``cursor_kwargs`` are applied except for the constructor - arguments that only the default cursor class accepts, such as ``unload`` + Only the connection's ``cursor_kwargs`` that the cursor class accepts are + applied, leaving out options of the default cursor class such as ``unload`` of a DataFrame cursor. The cursor class's default converter is used. Args: @@ -666,10 +666,8 @@ def _internal_cursor(self, cursor: type[FunctionalCursor]) -> FunctionalCursor: Returns: The created cursor. """ - backend_options = ( - self.cursor_class._constructor_keyword_names() - cursor._constructor_keyword_names() - ) - kwargs = {k: v for k, v in self.cursor_kwargs.items() if k not in backend_options} + accepted = cursor._constructor_keyword_names() + kwargs = {k: v for k, v in self.cursor_kwargs.items() if k in accepted} kwargs["converter"] = cursor.get_default_converter() return cast(FunctionalCursor, self._cursor(cursor, **kwargs)) diff --git a/tests/pyathena/test_connection.py b/tests/pyathena/test_connection.py index 4335d8877..2a091b192 100644 --- a/tests/pyathena/test_connection.py +++ b/tests/pyathena/test_connection.py @@ -397,10 +397,13 @@ def test_internal_cursor_leaves_out_default_cursor_options(self): with pytest.raises(TypeError, match="unexpected keyword argument 'unload'"): conn.cursor(Cursor) - def test_internal_cursor_unknown_cursor_kwargs(self): - conn = _connection( - cursor_class=PandasCursor, cursor_kwargs={"unload": True, "work_gruop": "typo"} - ) + def test_internal_cursor_leaves_out_custom_cursor_options(self): + class CustomCursor(Cursor): + def __init__(self, **kwargs: Any) -> None: + self.trace_id = kwargs.pop("trace_id") + super().__init__(**kwargs) + + conn = _connection(cursor_class=CustomCursor, cursor_kwargs={"trace_id": "x"}) - with pytest.raises(TypeError, match="unexpected keyword argument 'work_gruop'"): - conn._internal_cursor(Cursor) + assert conn.cursor().trace_id == "x" + assert type(conn._internal_cursor(Cursor)) is Cursor From 2bc1ef0e09ce17fb2a81adbf1dfcccaaa61ec26b Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 10 Oct 2026 01:01:41 +0900 Subject: [PATCH 14/15] Keep a cursor_kwargs option named cursor for the cursor class Make the cursor class parameter of Connection._cursor() positional-only so that cursor_kwargs={"cursor": ...} still reaches a custom cursor. Managed query results are read without a result file, so describe the unchecked case as having no results to read. Co-Authored-By: Claude Opus 5.5 --- docs/usage.md | 2 +- pyathena/connection.py | 2 +- tests/pyathena/test_connection.py | 10 ++++++++++ 3 files changed, 12 insertions(+), 2 deletions(-) diff --git a/docs/usage.md b/docs/usage.md index f1a6a063d..3f6be5a25 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -44,7 +44,7 @@ Unknown keyword arguments raise `TypeError` in cursor constructors and in `execu The connection's `cursor_kwargs` are passed to every cursor that it creates, so the requested cursor class must accept them. A cursor subclass that receives settings of its own must not pass them on to the base class constructor. The pandas and Polars cursors pass the `execute()` keyword arguments that they do not use themselves to the reader, such as `pandas.read_csv()`. -The reader rejects an unknown name when it reads a query result file, which PyAthena raises as `OperationalError`; the name is not checked when no result file is read. +The reader rejects an unknown name when it reads the query results, which PyAthena raises as `OperationalError`; the name is not checked when there are no results to read. ## Managed query result storage diff --git a/pyathena/connection.py b/pyathena/connection.py index e1218ea3d..6df8561a3 100644 --- a/pyathena/connection.py +++ b/pyathena/connection.py @@ -615,7 +615,7 @@ def cursor( return self._cursor(cursor, **{**self.cursor_kwargs, **kwargs}) def _cursor( - self, cursor: type[FunctionalCursor] | None = None, **kwargs: Any + self, cursor: type[FunctionalCursor] | None, /, **kwargs: Any ) -> FunctionalCursor | ConnectionCursor: """Create a cursor without applying the connection's ``cursor_kwargs``. diff --git a/tests/pyathena/test_connection.py b/tests/pyathena/test_connection.py index 2a091b192..923259ab4 100644 --- a/tests/pyathena/test_connection.py +++ b/tests/pyathena/test_connection.py @@ -407,3 +407,13 @@ def __init__(self, **kwargs: Any) -> None: assert conn.cursor().trace_id == "x" assert type(conn._internal_cursor(Cursor)) is Cursor + + def test_cursor_kwargs_named_cursor(self): + class CustomCursor(Cursor): + def __init__(self, *, cursor: Any = None, **kwargs: Any) -> None: + self.custom_cursor = cursor + super().__init__(**kwargs) + + conn = _connection(cursor_class=CustomCursor, cursor_kwargs={"cursor": "x"}) + + assert conn.cursor().custom_cursor == "x" From dd296bf5d231c5025f2f428384eec44a35ae876b Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 10 Oct 2026 01:30:38 +0900 Subject: [PATCH 15/15] Leave out arraysize from the internal cursor's cursor_kwargs A DataFrame cursor allows an arraysize above Cursor's page size limit of 1000, so applying it to the internal API cursor made SQLAlchemy reflection raise ProgrammingError. Co-Authored-By: Claude Opus 5.5 --- pyathena/connection.py | 6 ++++-- tests/pyathena/aio/test_connection.py | 8 +++++++- tests/pyathena/sqlalchemy/test_base.py | 2 +- tests/pyathena/test_connection.py | 8 +++++++- 4 files changed, 19 insertions(+), 5 deletions(-) diff --git a/pyathena/connection.py b/pyathena/connection.py index 6df8561a3..78c39e023 100644 --- a/pyathena/connection.py +++ b/pyathena/connection.py @@ -658,7 +658,9 @@ def _internal_cursor(self, cursor: type[FunctionalCursor]) -> FunctionalCursor: Only the connection's ``cursor_kwargs`` that the cursor class accepts are applied, leaving out options of the default cursor class such as ``unload`` - of a DataFrame cursor. The cursor class's default converter is used. + of a DataFrame cursor. ``arraysize`` is left out too, since a DataFrame + cursor allows values above this cursor's page size limit. The cursor + class's default converter is used. Args: cursor: The cursor class. @@ -666,7 +668,7 @@ def _internal_cursor(self, cursor: type[FunctionalCursor]) -> FunctionalCursor: Returns: The created cursor. """ - accepted = cursor._constructor_keyword_names() + accepted = cursor._constructor_keyword_names() - {"arraysize"} kwargs = {k: v for k, v in self.cursor_kwargs.items() if k in accepted} kwargs["converter"] = cursor.get_default_converter() return cast(FunctionalCursor, self._cursor(cursor, **kwargs)) diff --git a/tests/pyathena/aio/test_connection.py b/tests/pyathena/aio/test_connection.py index e36d7e73f..6fb2cb59a 100644 --- a/tests/pyathena/aio/test_connection.py +++ b/tests/pyathena/aio/test_connection.py @@ -225,7 +225,12 @@ async def test_cursor_arraysize_not_positive(self, cursor_class): conn.cursor(cursor_class, arraysize=0) async def test_internal_cursor_leaves_out_default_cursor_options(self): - cursor_kwargs = {"unload": True, "chunksize": 10, "schema_name": "configured"} + cursor_kwargs = { + "unload": True, + "chunksize": 10, + "arraysize": 5000, + "schema_name": "configured", + } conn = await _connection( cursor_class=AioPandasCursor, cursor_kwargs=cursor_kwargs, @@ -236,6 +241,7 @@ async def test_internal_cursor_leaves_out_default_cursor_options(self): assert type(cursor) is AioCursor assert cursor._schema_name == "configured" + assert cursor.arraysize == AioCursor.DEFAULT_FETCH_SIZE assert type(cursor._converter) is DefaultTypeConverter assert conn.cursor_kwargs == cursor_kwargs # Creating the cursor directly still passes the options it does not accept. diff --git a/tests/pyathena/sqlalchemy/test_base.py b/tests/pyathena/sqlalchemy/test_base.py index 0bcb5f661..ff7e39724 100644 --- a/tests/pyathena/sqlalchemy/test_base.py +++ b/tests/pyathena/sqlalchemy/test_base.py @@ -205,7 +205,7 @@ def test_internal_cursor_leaves_out_default_cursor_options( aws_access_key_id="access_key", aws_secret_access_key="secret_key", cursor_class=cursor_class, - cursor_kwargs={"unload": True, "chunksize": 10}, + cursor_kwargs={"unload": True, "chunksize": 10, "arraysize": 5000}, ) driver_connection = ( AsyncAdapt_pyathena_connection(MagicMock(), conn) diff --git a/tests/pyathena/test_connection.py b/tests/pyathena/test_connection.py index 923259ab4..8b9171a36 100644 --- a/tests/pyathena/test_connection.py +++ b/tests/pyathena/test_connection.py @@ -380,7 +380,12 @@ def test_close_closes_built_clients(self): s3_close.assert_called_once_with() def test_internal_cursor_leaves_out_default_cursor_options(self): - cursor_kwargs = {"unload": True, "chunksize": 10, "schema_name": "configured"} + cursor_kwargs = { + "unload": True, + "chunksize": 10, + "arraysize": 5000, + "schema_name": "configured", + } conn = _connection( cursor_class=PandasCursor, cursor_kwargs=cursor_kwargs, @@ -391,6 +396,7 @@ def test_internal_cursor_leaves_out_default_cursor_options(self): assert type(cursor) is Cursor assert cursor._schema_name == "configured" + assert cursor.arraysize == Cursor.DEFAULT_FETCH_SIZE assert type(cursor._converter) is DefaultTypeConverter assert conn.cursor_kwargs == cursor_kwargs # Creating the cursor directly still passes the options it does not accept.