diff --git a/docs/usage.md b/docs/usage.md index 4a2febcb6..3f6be5a25 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -40,6 +40,12 @@ 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. +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 the query results, which PyAthena raises as `OperationalError`; the name is not checked when there are no results to 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/aio/arrow/cursor.py b/pyathena/aio/arrow/cursor.py index c297dd0e1..a6d5fb17a 100644 --- a/pyathena/aio/arrow/cursor.py +++ b/pyathena/aio/arrow/cursor.py @@ -144,13 +144,15 @@ 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. """ + 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 d98f3f56c..71ed84d49 100644 --- a/pyathena/aio/cursor.py +++ b/pyathena/aio/cursor.py @@ -139,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. """ + self._validate_execute_kwargs(kwargs) self._reset_state() options = ExecuteOptions.resolve( options, diff --git a/pyathena/aio/s3fs/cursor.py b/pyathena/aio/s3fs/cursor.py index 37d2a81de..7fa6ff463 100644 --- a/pyathena/aio/s3fs/cursor.py +++ b/pyathena/aio/s3fs/cursor.py @@ -147,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. """ + 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 12c19221b..a50d4abf4 100644 --- a/pyathena/aio/spark/cursor.py +++ b/pyathena/aio/spark/cursor.py @@ -353,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. """ + 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 abb600c53..4fa5bd87b 100644 --- a/pyathena/aio/sqlalchemy/base.py +++ b/pyathena/aio/sqlalchemy/base.py @@ -276,6 +276,20 @@ 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: + """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) + 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 4be5972fd..7d56e64e1 100644 --- a/pyathena/arrow/async_cursor.py +++ b/pyathena/arrow/async_cursor.py @@ -212,13 +212,15 @@ 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. """ + 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 e39c58237..e9e7774c3 100644 --- a/pyathena/arrow/cursor.py +++ b/pyathena/arrow/cursor.py @@ -170,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. @@ -181,6 +182,7 @@ def execute( >>> cursor.execute("SELECT * FROM sales WHERE year = 2023") >>> table = cursor.as_arrow() # Returns Apache Arrow Table """ + 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 825662d45..5f047933c 100644 --- a/pyathena/async_cursor.py +++ b/pyathena/async_cursor.py @@ -241,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: @@ -254,6 +254,7 @@ def execute( >>> # Do other work while query runs... >>> result_set = future.result() # Wait for completion """ + self._validate_execute_kwargs(kwargs) options = ExecuteOptions.resolve( options, work_group=work_group, diff --git a/pyathena/common.py b/pyathena/common.py index dca0ba3ac..ae09d313e 100644 --- a/pyathena/common.py +++ b/pyathena/common.py @@ -7,9 +7,10 @@ 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 @@ -212,7 +213,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. @@ -236,7 +236,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 @@ -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, diff --git a/pyathena/connection.py b/pyathena/connection.py index 3a3d543ee..78c39e023 100644 --- a/pyathena/connection.py +++ b/pyathena/connection.py @@ -612,7 +612,21 @@ 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, /, **kwargs: Any + ) -> FunctionalCursor | ConnectionCursor: + """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: @@ -639,6 +653,26 @@ def cursor( **kwargs, ) + def _internal_cursor(self, cursor: type[FunctionalCursor]) -> FunctionalCursor: + """Create a cursor for queries that PyAthena runs on its own behalf. + + 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. ``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. + + Returns: + The created cursor. + """ + 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)) + def close(self) -> None: """Close the connection. diff --git a/pyathena/cursor.py b/pyathena/cursor.py index cab989858..95f3a8933 100644 --- a/pyathena/cursor.py +++ b/pyathena/cursor.py @@ -149,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. @@ -163,6 +163,7 @@ def execute( ... } ... ) """ + self._validate_execute_kwargs(kwargs) self._reset_state() options = ExecuteOptions.resolve( options, diff --git a/pyathena/s3fs/async_cursor.py b/pyathena/s3fs/async_cursor.py index 5b2f18272..999a37b63 100644 --- a/pyathena/s3fs/async_cursor.py +++ b/pyathena/s3fs/async_cursor.py @@ -215,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]). @@ -227,6 +228,7 @@ def execute( >>> result_set = future.result() >>> rows = result_set.fetchall() """ + 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 8ec1f7204..44f243a90 100644 --- a/pyathena/s3fs/cursor.py +++ b/pyathena/s3fs/cursor.py @@ -165,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. @@ -176,6 +177,7 @@ def execute( >>> cursor.execute("SELECT * FROM my_table WHERE id = %(id)s", {"id": 123}) >>> rows = cursor.fetchall() """ + 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 4548d1fe4..68c50bae1 100644 --- a/pyathena/spark/async_cursor.py +++ b/pyathena/spark/async_cursor.py @@ -212,6 +212,7 @@ def execute( work_group: str | None = None, **kwargs, ) -> tuple[str, "Future[AthenaQueryExecution | AthenaCalculationExecution]"]: + 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 a6f67d052..374fc5eac 100644 --- a/pyathena/spark/cursor.py +++ b/pyathena/spark/cursor.py @@ -111,6 +111,7 @@ def execute( work_group: str | None = None, **kwargs, ) -> SparkCursor: + 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/sqlalchemy/base.py b/pyathena/sqlalchemy/base.py index 392952b87..2a9b184d0 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/aio/test_connection.py b/tests/pyathena/aio/test_connection.py index c1b2f5d70..6fb2cb59a 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,27 @@ 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, + "arraysize": 5000, + "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 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. + with pytest.raises(TypeError, match="unexpected keyword argument 'unload'"): + conn.cursor(AioCursor) diff --git a/tests/pyathena/sqlalchemy/test_base.py b/tests/pyathena/sqlalchemy/test_base.py index 93885c40d..ff7e39724 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.converter import DefaultTypeConverter +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, "arraysize": 5000}, + ) + 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 @@ -213,11 +248,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" @@ -225,7 +262,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) @@ -386,7 +423,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) @@ -423,7 +461,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) @@ -476,7 +515,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) @@ -516,7 +556,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 @@ -596,8 +637,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), @@ -608,14 +649,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): @@ -640,7 +681,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) ) ), @@ -1449,7 +1490,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_connection.py b/tests/pyathena/test_connection.py index bdc68178d..8b9171a36 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]) @@ -98,8 +100,8 @@ def test_cursor_execute_keyword_only(cursor_class): 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) @@ -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"), @@ -136,7 +236,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): @@ -278,3 +378,48 @@ 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, + "arraysize": 5000, + "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 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. + with pytest.raises(TypeError, match="unexpected keyword argument 'unload'"): + conn.cursor(Cursor) + + 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"}) + + 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"