Merge pull request #735 from pyathena-dev/fix/734-execute-legacy-kwargs

Restore legacy keyword arguments on the private _execute() methods

Author: laughingman7743

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(

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(

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()

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 = []