Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
21 changes: 20 additions & 1 deletion pyathena/aio/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
21 changes: 20 additions & 1 deletion pyathena/common.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
51 changes: 51 additions & 0 deletions tests/pyathena/aio/test_cursor.py
Original file line number Diff line number Diff line change
@@ -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

Expand Down Expand Up @@ -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()
Expand Down
65 changes: 65 additions & 0 deletions tests/pyathena/test_cursor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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 = []
Expand Down