From 8d299e4c0aa15bfb0ead58804c034729e81e3f20 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 11 Jul 2026 00:44:37 +0900 Subject: [PATCH 1/2] Restore legacy keyword arguments on the private _execute() methods PR #733 (v3.35.0) consolidated the shared execute() keyword arguments into an ExecuteOptions dataclass and narrowed the private _execute() signature to (operation, parameters, options). External callers that predate ExecuteOptions - most notably dbt-athena <= 1.10.x - call _execute() directly with individual keywords (work_group, s3_staging_dir, cache_size, cache_expiration_time), which now raises: TypeError: BaseCursor._execute() got an unexpected keyword argument 'work_group' Restore the pre-3.35 individual keyword arguments on BaseCursor._execute() and AioBaseCursor._execute(), merging them into ExecuteOptions via the existing ExecuteOptions.resolve() (keywords win, None values are ignored). Closes #734 Co-Authored-By: Claude Fable 5 --- pyathena/aio/common.py | 21 +++++++++- pyathena/common.py | 21 +++++++++- tests/pyathena/aio/test_cursor.py | 17 ++++++++ tests/pyathena/test_cursor.py | 65 +++++++++++++++++++++++++++++++ 4 files changed, 122 insertions(+), 2 deletions(-) diff --git a/pyathena/aio/common.py b/pyathena/aio/common.py index 7cfd7c4b..19914164 100644 --- a/pyathena/aio/common.py +++ b/pyathena/aio/common.py @@ -28,9 +28,28 @@ async def _execute( # type: ignore[override] self, operation: str, parameters: dict[str, Any] | list[str] | None = None, + work_group: str | None = None, + s3_staging_dir: str | None = None, + cache_size: int | None = None, + cache_expiration_time: int | None = None, + result_reuse_enable: bool | None = None, + result_reuse_minutes: int | None = None, + paramstyle: str | None = None, options: ExecuteOptions | None = None, ) -> str: - options = ExecuteOptions.resolve(options) + # The individual keyword arguments are retained for backward compatibility + # with external callers that predate ExecuteOptions, mirroring + # BaseCursor._execute(). + 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, + ) query, execution_parameters = self._prepare_query(operation, parameters, options.paramstyle) request = self._build_start_query_execution_request( diff --git a/pyathena/common.py b/pyathena/common.py index 4ae3836f..9364e394 100644 --- a/pyathena/common.py +++ b/pyathena/common.py @@ -714,9 +714,28 @@ def _execute( self, operation: str, parameters: dict[str, Any] | list[str] | None = None, + work_group: str | None = None, + s3_staging_dir: str | None = None, + cache_size: int | None = None, + cache_expiration_time: int | None = None, + result_reuse_enable: bool | None = None, + result_reuse_minutes: int | None = None, + paramstyle: str | None = None, options: ExecuteOptions | None = None, ) -> str: - options = ExecuteOptions.resolve(options) + # The individual keyword arguments are retained for backward compatibility + # with external callers that predate ExecuteOptions (e.g. dbt-athena <= 1.10.x + # calls _execute() with work_group/s3_staging_dir/cache_* keywords). + 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, + ) query, execution_parameters = self._prepare_query(operation, parameters, options.paramstyle) request = self._build_start_query_execution_request( diff --git a/tests/pyathena/aio/test_cursor.py b/tests/pyathena/aio/test_cursor.py index 87bd7cf9..3ccaae8b 100644 --- a/tests/pyathena/aio/test_cursor.py +++ b/tests/pyathena/aio/test_cursor.py @@ -79,6 +79,23 @@ async def test_execute_with_options(self, aio_cursor): assert callback_results == [aio_cursor.query_id] assert await aio_cursor.fetchone() == (1,) + async def test_execute_internal_legacy_kwargs(self, aio_cursor): + """The private _execute() accepts the pre-3.35 individual keyword arguments. + + Regression test for #734, mirroring the synchronous cursor test for + external callers that invoke _execute() directly with individual keywords. + """ + query_id = await aio_cursor._execute( + "SELECT 1", + parameters=None, + work_group=ENV.default_work_group, + s3_staging_dir=None, + cache_size=0, + cache_expiration_time=0, + ) + query_execution = await aio_cursor._poll(query_id) + assert query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED + async def test_no_result_set_raises(self, aio_cursor): with pytest.raises(ProgrammingError): await aio_cursor.fetchone() diff --git a/tests/pyathena/test_cursor.py b/tests/pyathena/test_cursor.py index 30c5bb7b..bfe9e82a 100644 --- a/tests/pyathena/test_cursor.py +++ b/tests/pyathena/test_cursor.py @@ -20,6 +20,7 @@ from pyathena.cursor import Cursor from pyathena.error import DatabaseError, NotSupportedError, ProgrammingError from pyathena.model import AthenaQueryExecution +from pyathena.util import RetryConfig from tests import ENV from tests.pyathena.conftest import connect @@ -879,6 +880,70 @@ def test_execute_kwargs_override_options(self, cursor): assert from_kwargs == [cursor.query_id] assert cursor.fetchone() == (1,) + def test_execute_internal_legacy_kwargs(self, cursor): + """The private _execute() accepts the pre-3.35 individual keyword arguments. + + Regression test for #734: external callers such as dbt-athena <= 1.10.x + invoke _execute() directly with individual keywords instead of options. + """ + query_id = cursor._execute( + "SELECT 1", + parameters=None, + work_group=ENV.default_work_group, + s3_staging_dir=None, + cache_size=0, + cache_expiration_time=0, + ) + query_execution = cursor._poll(query_id) + assert query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED + + def test_execute_internal_legacy_kwargs_passthrough(self): + """The pre-3.35 _execute() keywords are forwarded to the request (no AWS). + + Regression test for #734: on 3.35.0 this call raised + ``TypeError: _execute() got an unexpected keyword argument 'work_group'``. + """ + cursor = Cursor.__new__(Cursor) # bypass __init__ to avoid AWS calls + cursor._connection = MagicMock() + cursor._connection.client.start_query_execution.return_value = { + "QueryExecutionId": "test_query_id" + } + cursor._retry_config = RetryConfig() + + with ( + patch.object( + Cursor, "_build_start_query_execution_request", return_value={} + ) as request_mock, + patch.object(Cursor, "_find_previous_query_id", return_value=None) as cache_mock, + ): + query_id = cursor._execute( + "SELECT 1", + parameters=None, + work_group="test_work_group", + s3_staging_dir="s3://test-bucket/path/", + cache_size=10, + cache_expiration_time=100, + result_reuse_enable=True, + result_reuse_minutes=5, + paramstyle="qmark", + ) + + assert query_id == "test_query_id" + request_mock.assert_called_once_with( + query="SELECT 1", + work_group="test_work_group", + s3_staging_dir="s3://test-bucket/path/", + result_reuse_enable=True, + result_reuse_minutes=5, + execution_parameters=None, + ) + cache_mock.assert_called_once_with( + "SELECT 1", + "test_work_group", + cache_size=10, + cache_expiration_time=100, + ) + def test_connection_level_callback(self): """Test connection-level default callback.""" callback_results = [] From 322ce8afa81a5795c561c6eb5bb5b7a632a24fa2 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 11 Jul 2026 01:02:09 +0900 Subject: [PATCH 2/2] Add no-AWS passthrough regression tests from self-review Add a mocked passthrough test for AioBaseCursor._execute() so the aio compat surface is guarded in CI without AWS credentials, mirroring the synchronous test. Drop the aio live-Athena legacy-kwargs test flagged as redundant in self-review: the mocked test verifies forwarding of all legacy keywords, and the synchronous integration test already proves the dbt-athena call pattern end-to-end. Co-Authored-By: Claude Fable 5 --- tests/pyathena/aio/test_cursor.py | 58 ++++++++++++++++++++++++------- 1 file changed, 46 insertions(+), 12 deletions(-) diff --git a/tests/pyathena/aio/test_cursor.py b/tests/pyathena/aio/test_cursor.py index 3ccaae8b..520b8795 100644 --- a/tests/pyathena/aio/test_cursor.py +++ b/tests/pyathena/aio/test_cursor.py @@ -1,12 +1,15 @@ import re from datetime import datetime +from unittest.mock import AsyncMock, MagicMock, patch import pytest from pyathena import ExecuteOptions +from pyathena.aio.cursor import AioCursor from pyathena.error import DatabaseError, ProgrammingError from pyathena.model import AthenaQueryExecution from pyathena.result_set import AthenaResultSet +from pyathena.util import RetryConfig from tests import ENV from tests.pyathena.aio.conftest import _aio_connect @@ -79,22 +82,53 @@ async def test_execute_with_options(self, aio_cursor): assert callback_results == [aio_cursor.query_id] assert await aio_cursor.fetchone() == (1,) - async def test_execute_internal_legacy_kwargs(self, aio_cursor): - """The private _execute() accepts the pre-3.35 individual keyword arguments. + async def test_execute_internal_legacy_kwargs_passthrough(self): + """The pre-3.35 _execute() keywords are forwarded to the request (no AWS). - Regression test for #734, mirroring the synchronous cursor test for - external callers that invoke _execute() directly with individual keywords. + Mirrors the synchronous cursor test (regression test for #734). """ - query_id = await aio_cursor._execute( + cursor = AioCursor.__new__(AioCursor) # bypass __init__ to avoid AWS calls + cursor._connection = MagicMock() + cursor._connection.client.start_query_execution.return_value = { + "QueryExecutionId": "test_query_id" + } + cursor._retry_config = RetryConfig() + + with ( + patch.object( + AioCursor, "_build_start_query_execution_request", return_value={} + ) as request_mock, + patch.object( + AioCursor, "_find_previous_query_id", new_callable=AsyncMock, return_value=None + ) as cache_mock, + ): + query_id = await cursor._execute( + "SELECT 1", + parameters=None, + work_group="test_work_group", + s3_staging_dir="s3://test-bucket/path/", + cache_size=10, + cache_expiration_time=100, + result_reuse_enable=True, + result_reuse_minutes=5, + paramstyle="qmark", + ) + + assert query_id == "test_query_id" + request_mock.assert_called_once_with( + query="SELECT 1", + work_group="test_work_group", + s3_staging_dir="s3://test-bucket/path/", + result_reuse_enable=True, + result_reuse_minutes=5, + execution_parameters=None, + ) + cache_mock.assert_awaited_once_with( "SELECT 1", - parameters=None, - work_group=ENV.default_work_group, - s3_staging_dir=None, - cache_size=0, - cache_expiration_time=0, + "test_work_group", + cache_size=10, + cache_expiration_time=100, ) - query_execution = await aio_cursor._poll(query_id) - assert query_execution.state == AthenaQueryExecution.STATE_SUCCEEDED async def test_no_result_set_raises(self, aio_cursor): with pytest.raises(ProgrammingError):