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..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,6 +82,54 @@ 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_passthrough(self): + """The pre-3.35 _execute() keywords are forwarded to the request (no AWS). + + Mirrors the synchronous cursor test (regression test for #734). + """ + 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", + "test_work_group", + cache_size=10, + cache_expiration_time=100, + ) + 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 = []