From 6d0f51536f9826ea8dd99b8123638418cb751c4b Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Wed, 7 Oct 2026 09:30:24 +0900 Subject: [PATCH 1/5] BREAKING: Reject ineffective Spark arguments and unify environment fallbacks --- docs/spark.md | 12 +++ docs/usage.md | 21 ++++- pyathena/aio/spark/cursor.py | 11 ++- pyathena/connection.py | 11 ++- pyathena/spark/async_cursor.py | 21 ++++- pyathena/spark/cursor.py | 23 ++++- tests/pyathena/aio/test_connection.py | 51 ++++++++++- tests/pyathena/spark/test_common.py | 117 +++++++++++++++++++++++++- tests/pyathena/test_connection.py | 68 ++++++++++++++- 9 files changed, 318 insertions(+), 17 deletions(-) diff --git a/docs/spark.md b/docs/spark.md index 295954fcf..0bb8d4d2c 100644 --- a/docs/spark.md +++ b/docs/spark.md @@ -65,6 +65,18 @@ with Connection(region_name="us-west-2", This cursor allows you to send PySpark code blocks and use Spark DataFrame and SQL. +### Breaking changes in PyAthena 4.0 + +`SparkCursor.execute()`, `AsyncSparkCursor.execute()`, and `AioSparkCursor.execute()` raise +`NotSupportedError` when `parameters` is not `None`, including an empty dictionary or list. +Earlier versions silently ignored these parameters and executed the code block without them. +Omit `parameters` or pass `None` when executing Spark code. + +These methods no longer accept `work_group`; passing it raises `TypeError`, even if its value is `None`. +Earlier versions silently ignored this argument. +Set the Spark workgroup when creating the connection or cursor instead. +The calculation runs in the Spark session selected by the cursor's `session_id` or the execution's `session_id` override. + (spark-session-lifecycle)= ### Session lifecycle diff --git a/docs/usage.md b/docs/usage.md index 61a6fa0c8..118bc3901 100644 --- a/docs/usage.md +++ b/docs/usage.md @@ -37,7 +37,6 @@ cursor = connection.cursor() 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. @@ -69,6 +68,24 @@ Passing `max_workers` to those four constructors raises `TypeError`. Passing it to `execute()` raises `TypeError` after the query has run; with `AsyncPandasCursor` and `AsyncPolarsCursor`, the returned future raises it. `AsyncPolarsCursor` used its `max_workers` for both the thread pool and the S3 read workers; pass both arguments to keep the former values. +## Environment fallbacks in PyAthena 4.0 + +For connection settings, `s3_staging_dir` and `work_group` use the same fallback rules: + +| Argument | Resolved setting | +|----------|------------------| +| Omitted or `None` | Read the corresponding environment variable. | +| `""` | Use `None` and ignore the corresponding environment variable. | +| Nonempty string | Use the supplied value. | + +The environment variables are `AWS_ATHENA_S3_STAGING_DIR` and `AWS_ATHENA_WORK_GROUP`. +These rules apply to `connect()`, `aio_connect()`, `Connection`, and `AioConnection`. +At least one resolved setting must be nonempty; otherwise connection creation raises `ProgrammingError`. + +This is a breaking change for `work_group=""`, which previously read `AWS_ATHENA_WORK_GROUP`. +Omit `work_group` or pass `None` to retain that fallback. +The behavior of `s3_staging_dir=""` is unchanged, including its use with managed query result storage. + ## 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, @@ -844,7 +861,7 @@ AWS_ATHENA_S3_STAGING_DIR : The S3 location where Athena automatically stores the query results and metadata information. Required if you have not set up workgroups. Not required if a workgroup has been set up. When connecting to a workgroup with [managed query result storage](https://docs.aws.amazon.com/athena/latest/ug/managed-results.html), pass ``s3_staging_dir=""`` to explicitly disable this environment variable fallback (see [Managed query result storage](#managed-query-result-storage)). AWS_ATHENA_WORK_GROUP -: The setting of the workgroup to execute the query. +: The workgroup used to execute queries when the connection's `work_group` argument is omitted or `None`. Pass `work_group=""` to disable this fallback. ## Credentials diff --git a/pyathena/aio/spark/cursor.py b/pyathena/aio/spark/cursor.py index a50d4abf4..cc9e6502c 100644 --- a/pyathena/aio/spark/cursor.py +++ b/pyathena/aio/spark/cursor.py @@ -341,24 +341,29 @@ async def execute( session_id: str | None = None, description: str | None = None, client_request_token: str | None = None, - work_group: str | None = None, **kwargs, ) -> AioSparkCursor: """Execute PySpark code asynchronously. Args: operation: PySpark code to execute. - parameters: Unused, kept for API compatibility. + parameters: Must be None; Spark calculations do not support parameters. session_id: Spark session ID override. description: Calculation description. client_request_token: Idempotency token. - work_group: Unused, kept for API compatibility. **kwargs: Unknown keyword arguments raise TypeError. Returns: Self reference for method chaining. + + Raises: + TypeError: If an unknown keyword argument is given. + NotSupportedError: If ``parameters`` is not None. + OperationalError: If the calculation does not complete successfully. """ self._validate_execute_kwargs(kwargs) + if parameters is not None: + raise NotSupportedError("Spark cursors do not support parameters.") # A failure below must not leave the previous calculation on the cursor. self._calculation_id = None self._calculation_execution = None diff --git a/pyathena/connection.py b/pyathena/connection.py index 78c39e023..0e416f43e 100644 --- a/pyathena/connection.py +++ b/pyathena/connection.py @@ -211,6 +211,7 @@ def __init__( Args: s3_staging_dir: S3 location to store query results. Required if not using workgroups or if workgroup doesn't have result location. + None uses the ``AWS_ATHENA_S3_STAGING_DIR`` environment variable. Pass an empty string to explicitly disable S3 staging and skip the ``AWS_ATHENA_S3_STAGING_DIR`` environment variable fallback. This is required when connecting to a workgroup with managed @@ -220,6 +221,9 @@ def __init__( catalog_name: Data catalog name. Defaults to "awsdatacatalog". work_group: Athena workgroup name. Can substitute for s3_staging_dir if workgroup has result location configured. + None uses the ``AWS_ATHENA_WORK_GROUP`` environment variable. + Pass an empty string to skip that environment variable fallback + and leave the workgroup unset. poll_interval: Seconds between query status polls. Defaults to 1.0. encryption_option: S3 encryption for results ("SSE_S3", "SSE_KMS", "CSE_KMS"). kms_key: KMS key ID when using SSE_KMS or CSE_KMS encryption. @@ -260,7 +264,8 @@ def __init__( Note: Either s3_staging_dir or work_group must be specified. Environment variables AWS_ATHENA_S3_STAGING_DIR and AWS_ATHENA_WORK_GROUP are - checked if parameters are not provided. + checked only when the corresponding argument is omitted or None. + An explicit empty string disables that fallback and becomes None. When using a workgroup with managed query result storage, pass ``s3_staging_dir=""`` to prevent the environment variable fallback @@ -282,8 +287,8 @@ def __init__( self.region_name = region_name self.schema_name = schema_name self.catalog_name = catalog_name - if work_group: - self.work_group: str | None = work_group + if work_group is not None: + self.work_group: str | None = work_group or None else: self.work_group = os.getenv(self._ENV_WORK_GROUP) self.poll_interval = poll_interval diff --git a/pyathena/spark/async_cursor.py b/pyathena/spark/async_cursor.py index 68c50bae1..1be413982 100644 --- a/pyathena/spark/async_cursor.py +++ b/pyathena/spark/async_cursor.py @@ -12,6 +12,7 @@ from multiprocessing import cpu_count from typing import TYPE_CHECKING, Any, cast +from pyathena import NotSupportedError from pyathena.model import AthenaCalculationExecution from pyathena.spark.common import SparkBaseCursor from pyathena.util import override @@ -209,10 +210,28 @@ def execute( session_id: str | None = None, description: str | None = None, client_request_token: str | None = None, - work_group: str | None = None, **kwargs, ) -> tuple[str, "Future[AthenaQueryExecution | AthenaCalculationExecution]"]: + """Start a PySpark calculation and return its ID and polling future. + + Args: + operation: PySpark code to execute. + parameters: Must be None; Spark calculations do not support parameters. + session_id: Spark session ID override. + description: Calculation description. + client_request_token: Idempotency token. + **kwargs: Unknown keyword arguments raise TypeError. + + Returns: + The calculation ID and a future containing its terminal execution. + + Raises: + TypeError: If an unknown keyword argument is given. + NotSupportedError: If ``parameters`` is not None. + """ self._validate_execute_kwargs(kwargs) + if parameters is not None: + raise NotSupportedError("Spark cursors do not support parameters.") 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 374fc5eac..c2e44977c 100644 --- a/pyathena/spark/cursor.py +++ b/pyathena/spark/cursor.py @@ -12,7 +12,7 @@ import logging from typing import Any, cast -from pyathena import OperationalError, ProgrammingError +from pyathena import NotSupportedError, OperationalError, ProgrammingError from pyathena.model import AthenaCalculationExecution, AthenaCalculationExecutionStatus from pyathena.spark.common import SparkBaseCursor, WithCalculationExecution from pyathena.util import override @@ -108,10 +108,29 @@ def execute( session_id: str | None = None, description: str | None = None, client_request_token: str | None = None, - work_group: str | None = None, **kwargs, ) -> SparkCursor: + """Execute PySpark code and wait for the calculation to complete. + + Args: + operation: PySpark code to execute. + parameters: Must be None; Spark calculations do not support parameters. + session_id: Spark session ID override. + description: Calculation description. + client_request_token: Idempotency token. + **kwargs: Unknown keyword arguments raise TypeError. + + Returns: + Self reference for method chaining. + + Raises: + TypeError: If an unknown keyword argument is given. + NotSupportedError: If ``parameters`` is not None. + OperationalError: If the calculation does not complete successfully. + """ self._validate_execute_kwargs(kwargs) + if parameters is not None: + raise NotSupportedError("Spark cursors do not support parameters.") # A failure below must not leave the previous calculation on the cursor. self._calculation_id = None self._calculation_execution = None diff --git a/tests/pyathena/aio/test_connection.py b/tests/pyathena/aio/test_connection.py index 6fb2cb59a..cd22bd6f6 100644 --- a/tests/pyathena/aio/test_connection.py +++ b/tests/pyathena/aio/test_connection.py @@ -32,6 +32,49 @@ READER_CURSOR_CLASSES = [AioPandasCursor, AioPolarsCursor] +@pytest.mark.parametrize("factory", [pyathena.aio_connect, AioConnection.create]) +@pytest.mark.parametrize("setting", ["s3_staging_dir", "work_group"]) +@pytest.mark.parametrize( + ("value", "expected"), + [(None, "from-environment"), ("", None), ("explicit", "explicit")], +) +async def test_environment_fallback(factory, setting, value, expected, monkeypatch): + environment_name = ( + "AWS_ATHENA_S3_STAGING_DIR" if setting == "s3_staging_dir" else "AWS_ATHENA_WORK_GROUP" + ) + monkeypatch.setenv(environment_name, "from-environment") + settings = ( + {"work_group": "other"} + if setting == "s3_staging_dir" + else {"s3_staging_dir": "s3://bucket/path/"} + ) + settings[setting] = value + connection = await factory( + region_name="us-east-1", + aws_access_key_id="access_key", + aws_secret_access_key="secret_key", + **settings, + ) + try: + if setting == "s3_staging_dir" and expected: + expected += "/" + assert getattr(connection, setting) == expected + finally: + connection.close() + + +@pytest.mark.parametrize("factory", [pyathena.aio_connect, AioConnection.create]) +async def test_empty_settings_disable_both_environment_fallbacks(factory, monkeypatch): + monkeypatch.setenv("AWS_ATHENA_S3_STAGING_DIR", "s3://environment/path/") + monkeypatch.setenv("AWS_ATHENA_WORK_GROUP", "environment-workgroup") + with ( + patch("pyathena.connection.Session") as session, + pytest.raises(ProgrammingError, match="Required argument"), + ): + await factory(s3_staging_dir="", work_group="") + session.assert_not_called() + + @pytest.mark.parametrize("factory", [pyathena.aio_connect, AioConnection.create]) async def test_connect_keyword_only(factory): assert all( @@ -73,9 +116,13 @@ def test_cursor_execute_keyword_only(cursor_class): ] == ["self", "operation", "parameters"] parameters = {"value": 1} cursor = object.__new__(cursor_class) - bound = execute_signature.bind(cursor, "SELECT %(value)s", parameters, work_group="group") + bound = execute_signature.bind(cursor, "SELECT %(value)s", parameters) assert bound.arguments["parameters"] is parameters - assert bound.arguments["work_group"] == "group" + if cursor_class is AioSparkCursor: + assert "work_group" not in execute_signature.parameters + else: + bound = execute_signature.bind(cursor, "SELECT %(value)s", parameters, work_group="group") + assert bound.arguments["work_group"] == "group" execute_signature.bind(cursor, operation="SELECT 1") with pytest.raises(TypeError, match="positional"): cursor.execute("SELECT %(value)s", parameters, "group") diff --git a/tests/pyathena/spark/test_common.py b/tests/pyathena/spark/test_common.py index ad83598b3..e84dbf41d 100644 --- a/tests/pyathena/spark/test_common.py +++ b/tests/pyathena/spark/test_common.py @@ -9,12 +9,13 @@ import logging import threading import uuid -from unittest.mock import MagicMock, PropertyMock, patch +from concurrent.futures import Future +from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch import pytest from botocore.exceptions import ClientError -from pyathena import DatabaseError, OperationalError +from pyathena import DatabaseError, NotSupportedError, OperationalError from pyathena.aio.spark.cursor import AioSparkCursor from pyathena.model import AthenaCalculationExecutionStatus, AthenaSessionStatus from pyathena.spark.async_cursor import AsyncSparkCursor @@ -142,6 +143,118 @@ def _close(cursor) -> None: cursor.close() +@pytest.fixture(params=SPARK_CURSOR_CLASSES, ids=lambda cls: cls.__name__) +def execute_cursor(request): + """A Spark cursor whose calculation and polling methods are mocked. + + Args: + request: The pytest fixture request, selecting the cursor class. + + Returns: + A cursor with a previous calculation and a completed next calculation. + """ + cursor = request.param.__new__(request.param) + cursor._session_id = "session_id" + cursor._calculation_id = "previous_calculation_id" + cursor._calculation_execution = MagicMock( + state=AthenaCalculationExecutionStatus.STATE_COMPLETED + ) + execution = MagicMock(state=AthenaCalculationExecutionStatus.STATE_COMPLETED) + mock = AsyncMock if isinstance(cursor, AioSparkCursor) else MagicMock + cursor._calculate = mock(return_value="calculation_id") + cursor._poll = mock(return_value=execution) + cursor._executor = MagicMock() + future = Future() + future.set_result(execution) + cursor._executor.submit.return_value = future + return cursor + + +async def _execute(cursor, *args, **kwargs): + """Execute code using the synchronous or native asyncio Spark interface. + + Args: + cursor: The Spark cursor to call. + *args: Positional execution arguments. + **kwargs: Keyword execution arguments. + + Returns: + The result of the cursor's ``execute()`` method. + """ + if isinstance(cursor, AioSparkCursor): + return await cursor.execute(*args, **kwargs) + return cursor.execute(*args, **kwargs) + + +class TestSparkExecute: + @pytest.mark.parametrize("parameters", [{}, [], {"value": 1}, ["value"]]) + @pytest.mark.parametrize("positional", [True, False]) + async def test_parameters_rejected(self, execute_cursor, parameters, positional): + previous_execution = execute_cursor._calculation_execution + arguments = (parameters,) if positional else () + keywords = {} if positional else {"parameters": parameters} + with pytest.raises(NotSupportedError, match="parameters"): + await _execute(execute_cursor, "code", *arguments, **keywords) + + execute_cursor._calculate.assert_not_called() + execute_cursor._poll.assert_not_called() + execute_cursor._executor.submit.assert_not_called() + assert execute_cursor.calculation_id == "previous_calculation_id" + assert execute_cursor._calculation_execution is previous_execution + + @pytest.mark.parametrize("work_group", [None, "", "other-workgroup"]) + async def test_work_group_rejected(self, execute_cursor, work_group): + previous_execution = execute_cursor._calculation_execution + keywords = {"work_group": work_group} + with pytest.raises(TypeError, match="work_group"): + await _execute(execute_cursor, "code", **keywords) + + execute_cursor._calculate.assert_not_called() + execute_cursor._poll.assert_not_called() + execute_cursor._executor.submit.assert_not_called() + assert execute_cursor.calculation_id == "previous_calculation_id" + assert execute_cursor._calculation_execution is previous_execution + + @pytest.mark.parametrize( + ("arguments", "keywords"), + [((), {}), ((None,), {}), ((), {"parameters": None})], + ) + @pytest.mark.parametrize("session_id", [None, "override-session"]) + async def test_none_parameters_accepted(self, execute_cursor, arguments, keywords, session_id): + result = await _execute( + execute_cursor, + "code", + *arguments, + **keywords, + session_id=session_id, + description="description", + client_request_token="token", + ) + + execute_cursor._calculate.assert_called_once_with( + session_id=session_id or "session_id", + code_block="code", + description="description", + client_request_token="token", + ) + if isinstance(execute_cursor, AioSparkCursor): + execute_cursor._calculate.assert_awaited_once() + execute_cursor._poll.assert_awaited_once_with("calculation_id") + elif isinstance(execute_cursor, AsyncSparkCursor): + calculation_id, future = result + assert calculation_id == "calculation_id" + assert future is execute_cursor._executor.submit.return_value + assert future.result().state == AthenaCalculationExecutionStatus.STATE_COMPLETED + execute_cursor._executor.submit.assert_called_once_with( + execute_cursor._poll, "calculation_id" + ) + else: + execute_cursor._poll.assert_called_once_with("calculation_id") + if not isinstance(execute_cursor, AsyncSparkCursor): + assert result is execute_cursor + assert execute_cursor.calculation_id == "calculation_id" + + class TestSparkBaseCursor: @pytest.mark.parametrize( "state", diff --git a/tests/pyathena/test_connection.py b/tests/pyathena/test_connection.py index 8b9171a36..ec72960a9 100644 --- a/tests/pyathena/test_connection.py +++ b/tests/pyathena/test_connection.py @@ -55,6 +55,66 @@ READER_CURSOR_CLASSES = [PandasCursor, PolarsCursor, AsyncPandasCursor, AsyncPolarsCursor] +@pytest.mark.parametrize("factory", [Connection, pyathena.connect]) +@pytest.mark.parametrize("setting", ["s3_staging_dir", "work_group"]) +@pytest.mark.parametrize( + ("arguments", "environment", "expected"), + [ + ({}, "from-environment", "from-environment"), + ({"value": None}, "from-environment", "from-environment"), + ({"value": ""}, "from-environment", None), + ({"value": "explicit"}, "from-environment", "explicit"), + ({}, None, None), + ({"value": None}, None, None), + ({"value": ""}, None, None), + ({"value": "explicit"}, None, "explicit"), + ({}, "", ""), + ({"value": None}, "", ""), + ({"value": ""}, "", None), + ({"value": "explicit"}, "", "explicit"), + ], +) +def test_environment_fallback(factory, setting, arguments, environment, expected, monkeypatch): + environment_name = ( + "AWS_ATHENA_S3_STAGING_DIR" if setting == "s3_staging_dir" else "AWS_ATHENA_WORK_GROUP" + ) + if environment is None: + monkeypatch.delenv(environment_name, raising=False) + else: + monkeypatch.setenv(environment_name, environment) + settings = ( + {"work_group": "other"} + if setting == "s3_staging_dir" + else {"s3_staging_dir": "s3://bucket/path/"} + ) + if arguments: + settings[setting] = arguments["value"] + connection = factory( + region_name="us-east-1", + aws_access_key_id="access_key", + aws_secret_access_key="secret_key", + **settings, + ) + try: + if setting == "s3_staging_dir" and expected: + expected += "/" + assert getattr(connection, setting) == expected + finally: + connection.close() + + +@pytest.mark.parametrize("factory", [Connection, pyathena.connect]) +def test_empty_settings_disable_both_environment_fallbacks(factory, monkeypatch): + monkeypatch.setenv("AWS_ATHENA_S3_STAGING_DIR", "s3://environment/path/") + monkeypatch.setenv("AWS_ATHENA_WORK_GROUP", "environment-workgroup") + with ( + patch("pyathena.connection.Session") as session, + pytest.raises(ProgrammingError, match="Required argument"), + ): + factory(s3_staging_dir="", work_group="") + session.assert_not_called() + + @pytest.mark.parametrize("factory", [pyathena.connect, Connection]) def test_connect_keyword_only(factory): assert all( @@ -91,9 +151,13 @@ def test_cursor_execute_keyword_only(cursor_class): ] == ["self", "operation", "parameters"] parameters = {"value": 1} cursor = object.__new__(cursor_class) - bound = execute_signature.bind(cursor, "SELECT %(value)s", parameters, work_group="group") + bound = execute_signature.bind(cursor, "SELECT %(value)s", parameters) assert bound.arguments["parameters"] is parameters - assert bound.arguments["work_group"] == "group" + if cursor_class in (SparkCursor, AsyncSparkCursor): + assert "work_group" not in execute_signature.parameters + else: + bound = execute_signature.bind(cursor, "SELECT %(value)s", parameters, work_group="group") + assert bound.arguments["work_group"] == "group" execute_signature.bind(cursor, operation="SELECT 1") with pytest.raises(TypeError, match="positional"): cursor.execute("SELECT %(value)s", parameters, "group") From 7e48f1d73bb6630684e812f1be11b852585db4a2 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 10 Oct 2026 12:12:58 +0900 Subject: [PATCH 2/5] Narrow the Spark execute and environment fallback tests Spark work_group rejection now comes from the generic unknown-keyword check added in #1102, so one case covers it. Fallback cases use explicit expected values instead of normalizing them in the test body. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/aio/test_connection.py | 41 ++-------- tests/pyathena/spark/test_common.py | 103 ++++++++------------------ tests/pyathena/test_connection.py | 59 +++++---------- 3 files changed, 54 insertions(+), 149 deletions(-) diff --git a/tests/pyathena/aio/test_connection.py b/tests/pyathena/aio/test_connection.py index cd22bd6f6..624cecaf1 100644 --- a/tests/pyathena/aio/test_connection.py +++ b/tests/pyathena/aio/test_connection.py @@ -33,46 +33,19 @@ @pytest.mark.parametrize("factory", [pyathena.aio_connect, AioConnection.create]) -@pytest.mark.parametrize("setting", ["s3_staging_dir", "work_group"]) -@pytest.mark.parametrize( - ("value", "expected"), - [(None, "from-environment"), ("", None), ("explicit", "explicit")], -) -async def test_environment_fallback(factory, setting, value, expected, monkeypatch): - environment_name = ( - "AWS_ATHENA_S3_STAGING_DIR" if setting == "s3_staging_dir" else "AWS_ATHENA_WORK_GROUP" - ) - monkeypatch.setenv(environment_name, "from-environment") - settings = ( - {"work_group": "other"} - if setting == "s3_staging_dir" - else {"s3_staging_dir": "s3://bucket/path/"} - ) - settings[setting] = value - connection = await factory( +async def test_empty_work_group_disables_environment_fallback(factory, monkeypatch): + monkeypatch.setenv("AWS_ATHENA_WORK_GROUP", "environment-workgroup") + conn = await factory( region_name="us-east-1", + s3_staging_dir="s3://bucket/path/", + work_group="", aws_access_key_id="access_key", aws_secret_access_key="secret_key", - **settings, ) try: - if setting == "s3_staging_dir" and expected: - expected += "/" - assert getattr(connection, setting) == expected + assert conn.work_group is None finally: - connection.close() - - -@pytest.mark.parametrize("factory", [pyathena.aio_connect, AioConnection.create]) -async def test_empty_settings_disable_both_environment_fallbacks(factory, monkeypatch): - monkeypatch.setenv("AWS_ATHENA_S3_STAGING_DIR", "s3://environment/path/") - monkeypatch.setenv("AWS_ATHENA_WORK_GROUP", "environment-workgroup") - with ( - patch("pyathena.connection.Session") as session, - pytest.raises(ProgrammingError, match="Required argument"), - ): - await factory(s3_staging_dir="", work_group="") - session.assert_not_called() + conn.close() @pytest.mark.parametrize("factory", [pyathena.aio_connect, AioConnection.create]) diff --git a/tests/pyathena/spark/test_common.py b/tests/pyathena/spark/test_common.py index e84dbf41d..8e6de24a3 100644 --- a/tests/pyathena/spark/test_common.py +++ b/tests/pyathena/spark/test_common.py @@ -9,7 +9,6 @@ import logging import threading import uuid -from concurrent.futures import Future from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch import pytest @@ -145,41 +144,37 @@ def _close(cursor) -> None: @pytest.fixture(params=SPARK_CURSOR_CLASSES, ids=lambda cls: cls.__name__) def execute_cursor(request): - """A Spark cursor whose calculation and polling methods are mocked. + """A Spark cursor whose calculation requests are mocked. Args: request: The pytest fixture request, selecting the cursor class. Returns: - A cursor with a previous calculation and a completed next calculation. + The cursor, with ``previous`` as its last calculation ID. """ - cursor = request.param.__new__(request.param) + cursor = request.param.__new__(request.param) # bypass __init__ to avoid AWS calls cursor._session_id = "session_id" - cursor._calculation_id = "previous_calculation_id" - cursor._calculation_execution = MagicMock( - state=AthenaCalculationExecutionStatus.STATE_COMPLETED - ) - execution = MagicMock(state=AthenaCalculationExecutionStatus.STATE_COMPLETED) - mock = AsyncMock if isinstance(cursor, AioSparkCursor) else MagicMock + cursor._calculation_id = "previous" + cursor._calculation_execution = None + mock = AsyncMock if request.param is AioSparkCursor else MagicMock cursor._calculate = mock(return_value="calculation_id") - cursor._poll = mock(return_value=execution) + cursor._poll = mock( + return_value=MagicMock(state=AthenaCalculationExecutionStatus.STATE_COMPLETED) + ) cursor._executor = MagicMock() - future = Future() - future.set_result(execution) - cursor._executor.submit.return_value = future return cursor async def _execute(cursor, *args, **kwargs): - """Execute code using the synchronous or native asyncio Spark interface. + """Call ``execute()``, awaiting it on ``AioSparkCursor``. Args: - cursor: The Spark cursor to call. - *args: Positional execution arguments. - **kwargs: Keyword execution arguments. + cursor: The Spark cursor. + *args: Positional arguments for ``execute()``. + **kwargs: Keyword arguments for ``execute()``. Returns: - The result of the cursor's ``execute()`` method. + The result of ``execute()``. """ if isinstance(cursor, AioSparkCursor): return await cursor.execute(*args, **kwargs) @@ -187,72 +182,32 @@ async def _execute(cursor, *args, **kwargs): class TestSparkExecute: - @pytest.mark.parametrize("parameters", [{}, [], {"value": 1}, ["value"]]) - @pytest.mark.parametrize("positional", [True, False]) - async def test_parameters_rejected(self, execute_cursor, parameters, positional): - previous_execution = execute_cursor._calculation_execution - arguments = (parameters,) if positional else () - keywords = {} if positional else {"parameters": parameters} - with pytest.raises(NotSupportedError, match="parameters"): - await _execute(execute_cursor, "code", *arguments, **keywords) + @pytest.mark.parametrize("parameters", [{}, [], {"value": 1}]) + async def test_parameters_rejected(self, execute_cursor, parameters): + with pytest.raises(NotSupportedError, match="do not support parameters"): + await _execute(execute_cursor, "code", parameters) execute_cursor._calculate.assert_not_called() - execute_cursor._poll.assert_not_called() - execute_cursor._executor.submit.assert_not_called() - assert execute_cursor.calculation_id == "previous_calculation_id" - assert execute_cursor._calculation_execution is previous_execution - - @pytest.mark.parametrize("work_group", [None, "", "other-workgroup"]) - async def test_work_group_rejected(self, execute_cursor, work_group): - previous_execution = execute_cursor._calculation_execution - keywords = {"work_group": work_group} - with pytest.raises(TypeError, match="work_group"): - await _execute(execute_cursor, "code", **keywords) + assert execute_cursor.calculation_id == "previous" + + async def test_work_group_rejected(self, execute_cursor): + with pytest.raises(TypeError, match="unexpected keyword argument 'work_group'"): + await _execute(execute_cursor, "code", work_group=None) execute_cursor._calculate.assert_not_called() - execute_cursor._poll.assert_not_called() - execute_cursor._executor.submit.assert_not_called() - assert execute_cursor.calculation_id == "previous_calculation_id" - assert execute_cursor._calculation_execution is previous_execution @pytest.mark.parametrize( - ("arguments", "keywords"), - [((), {}), ((None,), {}), ((), {"parameters": None})], + ("args", "kwargs"), [((), {}), ((None,), {}), ((), {"parameters": None})] ) - @pytest.mark.parametrize("session_id", [None, "override-session"]) - async def test_none_parameters_accepted(self, execute_cursor, arguments, keywords, session_id): - result = await _execute( - execute_cursor, - "code", - *arguments, - **keywords, - session_id=session_id, - description="description", - client_request_token="token", - ) + async def test_none_parameters_accepted(self, execute_cursor, args, kwargs): + await _execute(execute_cursor, "code", *args, **kwargs) execute_cursor._calculate.assert_called_once_with( - session_id=session_id or "session_id", + session_id="session_id", code_block="code", - description="description", - client_request_token="token", + description=None, + client_request_token=None, ) - if isinstance(execute_cursor, AioSparkCursor): - execute_cursor._calculate.assert_awaited_once() - execute_cursor._poll.assert_awaited_once_with("calculation_id") - elif isinstance(execute_cursor, AsyncSparkCursor): - calculation_id, future = result - assert calculation_id == "calculation_id" - assert future is execute_cursor._executor.submit.return_value - assert future.result().state == AthenaCalculationExecutionStatus.STATE_COMPLETED - execute_cursor._executor.submit.assert_called_once_with( - execute_cursor._poll, "calculation_id" - ) - else: - execute_cursor._poll.assert_called_once_with("calculation_id") - if not isinstance(execute_cursor, AsyncSparkCursor): - assert result is execute_cursor - assert execute_cursor.calculation_id == "calculation_id" class TestSparkBaseCursor: diff --git a/tests/pyathena/test_connection.py b/tests/pyathena/test_connection.py index ec72960a9..6df0b1686 100644 --- a/tests/pyathena/test_connection.py +++ b/tests/pyathena/test_connection.py @@ -56,63 +56,40 @@ @pytest.mark.parametrize("factory", [Connection, pyathena.connect]) -@pytest.mark.parametrize("setting", ["s3_staging_dir", "work_group"]) @pytest.mark.parametrize( - ("arguments", "environment", "expected"), + ("setting", "kwargs", "expected"), [ - ({}, "from-environment", "from-environment"), - ({"value": None}, "from-environment", "from-environment"), - ({"value": ""}, "from-environment", None), - ({"value": "explicit"}, "from-environment", "explicit"), - ({}, None, None), - ({"value": None}, None, None), - ({"value": ""}, None, None), - ({"value": "explicit"}, None, "explicit"), - ({}, "", ""), - ({"value": None}, "", ""), - ({"value": ""}, "", None), - ({"value": "explicit"}, "", "explicit"), + ("s3_staging_dir", {}, "s3://environment/path/"), + ("s3_staging_dir", {"s3_staging_dir": None}, "s3://environment/path/"), + ("s3_staging_dir", {"s3_staging_dir": ""}, None), + ("s3_staging_dir", {"s3_staging_dir": "s3://explicit/path/"}, "s3://explicit/path/"), + ("work_group", {}, "environment-workgroup"), + ("work_group", {"work_group": None}, "environment-workgroup"), + ("work_group", {"work_group": ""}, None), + ("work_group", {"work_group": "explicit-workgroup"}, "explicit-workgroup"), ], ) -def test_environment_fallback(factory, setting, arguments, environment, expected, monkeypatch): - environment_name = ( - "AWS_ATHENA_S3_STAGING_DIR" if setting == "s3_staging_dir" else "AWS_ATHENA_WORK_GROUP" - ) - if environment is None: - monkeypatch.delenv(environment_name, raising=False) - else: - monkeypatch.setenv(environment_name, environment) - settings = ( - {"work_group": "other"} - if setting == "s3_staging_dir" - else {"s3_staging_dir": "s3://bucket/path/"} - ) - if arguments: - settings[setting] = arguments["value"] - connection = factory( +def test_environment_fallback(factory, setting, kwargs, expected, monkeypatch): + monkeypatch.setenv("AWS_ATHENA_S3_STAGING_DIR", "s3://environment/path/") + monkeypatch.setenv("AWS_ATHENA_WORK_GROUP", "environment-workgroup") + conn = factory( region_name="us-east-1", aws_access_key_id="access_key", aws_secret_access_key="secret_key", - **settings, + **kwargs, ) try: - if setting == "s3_staging_dir" and expected: - expected += "/" - assert getattr(connection, setting) == expected + assert getattr(conn, setting) == expected finally: - connection.close() + conn.close() @pytest.mark.parametrize("factory", [Connection, pyathena.connect]) def test_empty_settings_disable_both_environment_fallbacks(factory, monkeypatch): monkeypatch.setenv("AWS_ATHENA_S3_STAGING_DIR", "s3://environment/path/") monkeypatch.setenv("AWS_ATHENA_WORK_GROUP", "environment-workgroup") - with ( - patch("pyathena.connection.Session") as session, - pytest.raises(ProgrammingError, match="Required argument"), - ): - factory(s3_staging_dir="", work_group="") - session.assert_not_called() + with pytest.raises(ProgrammingError, match="Required argument"): + factory(region_name="us-east-1", s3_staging_dir="", work_group="") @pytest.mark.parametrize("factory", [pyathena.connect, Connection]) From 86f4443c6f3f742d250cf4a044a09c0056e05e0c Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 10 Oct 2026 12:14:30 +0900 Subject: [PATCH 3/5] Document the empty work_group fallback in connect() Co-Authored-By: Claude Opus 5.5 --- pyathena/__init__.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyathena/__init__.py b/pyathena/__init__.py index 85469a959..5739e89b4 100644 --- a/pyathena/__init__.py +++ b/pyathena/__init__.py @@ -100,6 +100,8 @@ def connect(**kwargs) -> Connection[Any]: catalog_name: Athena data catalog name. Defaults to "awsdatacatalog". work_group: Athena workgroup name. Can be used instead of s3_staging_dir if the workgroup has a result location configured. + Pass an empty string to skip the ``AWS_ATHENA_WORK_GROUP`` + environment variable fallback. poll_interval: Time in seconds between polling for query completion. Defaults to 1.0. encryption_option: S3 encryption option for query results. Can be From 169c8a03cf4ea0d1ad9242351eaef15171c23c7c Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 10 Oct 2026 12:52:13 +0900 Subject: [PATCH 4/5] Move Spark execute argument tests to each cursor's test class The checks live in each cursor's execute(), not in spark/common.py, so the tests belong in TestSparkCursor, TestAsyncSparkCursor, and TestAioSparkCursor. Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/aio/spark/test_cursor.py | 44 ++++++++++++++ tests/pyathena/spark/test_async_cursor.py | 56 +++++++++++++++++- tests/pyathena/spark/test_common.py | 72 +---------------------- tests/pyathena/spark/test_spark_cursor.py | 54 ++++++++++++++++- 4 files changed, 153 insertions(+), 73 deletions(-) diff --git a/tests/pyathena/aio/spark/test_cursor.py b/tests/pyathena/aio/spark/test_cursor.py index 228e07d83..0a175bf81 100644 --- a/tests/pyathena/aio/spark/test_cursor.py +++ b/tests/pyathena/aio/spark/test_cursor.py @@ -504,6 +504,50 @@ async def test_execute_cancelled_while_starting_without_kill_on_interrupt(self): cancel.assert_not_awaited() assert cursor.calculation_id is None + @pytest.mark.parametrize("parameters", [{}, [], {"value": 1}]) + async def test_execute_rejects_parameters(self, parameters): + """Any parameters other than None raise before a calculation starts (no AWS).""" + cursor = AioSparkCursor.__new__(AioSparkCursor) # bypass __init__ to avoid AWS calls + cursor._calculation_id = "previous" + cursor._calculate = AsyncMock() + + with pytest.raises(NotSupportedError, match="do not support parameters"): + await cursor.execute("code", parameters) + + cursor._calculate.assert_not_awaited() + assert cursor.calculation_id == "previous" + + async def test_execute_rejects_work_group(self): + """work_group is not an execute() argument, even when None (no AWS).""" + cursor = AioSparkCursor.__new__(AioSparkCursor) # bypass __init__ to avoid AWS calls + cursor._calculate = AsyncMock() + + with pytest.raises(TypeError, match="unexpected keyword argument 'work_group'"): + await cursor.execute("code", work_group=None) + + cursor._calculate.assert_not_awaited() + + @pytest.mark.parametrize( + ("args", "kwargs"), [((), {}), ((None,), {}), ((), {"parameters": None})] + ) + async def test_execute_accepts_none_parameters(self, args, kwargs): + """Omitted or None parameters run the calculation (no AWS).""" + cursor = AioSparkCursor.__new__(AioSparkCursor) # bypass __init__ to avoid AWS calls + cursor._session_id = "session_id" + cursor._calculate = AsyncMock(return_value="calculation_id") + cursor._poll = AsyncMock( + return_value=MagicMock(state=AthenaCalculationExecutionStatus.STATE_COMPLETED) + ) + + await cursor.execute("code", *args, **kwargs) + + cursor._calculate.assert_awaited_once_with( + session_id="session_id", + code_block="code", + description=None, + client_request_token=None, + ) + async def test_executemany(self, aio_spark_cursor): with pytest.raises(NotSupportedError): await aio_spark_cursor.executemany("SELECT 1", []) diff --git a/tests/pyathena/spark/test_async_cursor.py b/tests/pyathena/spark/test_async_cursor.py index 97e726852..154ffe34f 100644 --- a/tests/pyathena/spark/test_async_cursor.py +++ b/tests/pyathena/spark/test_async_cursor.py @@ -8,11 +8,11 @@ import textwrap import threading from concurrent.futures import ThreadPoolExecutor -from unittest.mock import MagicMock +from unittest.mock import MagicMock, patch import pytest -from pyathena import OperationalError +from pyathena import NotSupportedError, OperationalError from pyathena.model import AthenaCalculationExecutionStatus, AthenaSessionStatus from pyathena.spark.async_cursor import AsyncSparkCursor from tests import ENV @@ -256,3 +256,55 @@ def test_close_retries_termination_after_failure(self): cursor.close() assert cursor._terminate_session.call_count == 2 + + @pytest.mark.parametrize("parameters", [{}, [], {"value": 1}]) + def test_execute_rejects_parameters(self, parameters): + """Any parameters other than None raise before a calculation starts (no AWS).""" + cursor = AsyncSparkCursor.__new__(AsyncSparkCursor) # bypass __init__ to avoid AWS calls + cursor._executor = MagicMock() + + with ( + patch.object(AsyncSparkCursor, "_calculate") as calculate, + pytest.raises(NotSupportedError, match="do not support parameters"), + ): + cursor.execute("code", parameters) + + calculate.assert_not_called() + cursor._executor.submit.assert_not_called() + + def test_execute_rejects_work_group(self): + """work_group is not an execute() argument, even when None (no AWS).""" + cursor = AsyncSparkCursor.__new__(AsyncSparkCursor) # bypass __init__ to avoid AWS calls + cursor._executor = MagicMock() + + with ( + patch.object(AsyncSparkCursor, "_calculate") as calculate, + pytest.raises(TypeError, match="unexpected keyword argument 'work_group'"), + ): + cursor.execute("code", work_group=None) + + calculate.assert_not_called() + cursor._executor.submit.assert_not_called() + + @pytest.mark.parametrize( + ("args", "kwargs"), [((), {}), ((None,), {}), ((), {"parameters": None})] + ) + def test_execute_accepts_none_parameters(self, args, kwargs): + """Omitted or None parameters start the calculation and its polling (no AWS).""" + cursor = AsyncSparkCursor.__new__(AsyncSparkCursor) # bypass __init__ to avoid AWS calls + cursor._session_id = "session_id" + cursor._executor = MagicMock() + + with patch.object( + AsyncSparkCursor, "_calculate", return_value="calculation_id" + ) as calculate: + query_id, future = cursor.execute("code", *args, **kwargs) + + calculate.assert_called_once_with( + session_id="session_id", + code_block="code", + description=None, + client_request_token=None, + ) + assert query_id == "calculation_id" + assert future is cursor._executor.submit.return_value diff --git a/tests/pyathena/spark/test_common.py b/tests/pyathena/spark/test_common.py index 8e6de24a3..ad83598b3 100644 --- a/tests/pyathena/spark/test_common.py +++ b/tests/pyathena/spark/test_common.py @@ -9,12 +9,12 @@ import logging import threading import uuid -from unittest.mock import AsyncMock, MagicMock, PropertyMock, patch +from unittest.mock import MagicMock, PropertyMock, patch import pytest from botocore.exceptions import ClientError -from pyathena import DatabaseError, NotSupportedError, OperationalError +from pyathena import DatabaseError, OperationalError from pyathena.aio.spark.cursor import AioSparkCursor from pyathena.model import AthenaCalculationExecutionStatus, AthenaSessionStatus from pyathena.spark.async_cursor import AsyncSparkCursor @@ -142,74 +142,6 @@ def _close(cursor) -> None: cursor.close() -@pytest.fixture(params=SPARK_CURSOR_CLASSES, ids=lambda cls: cls.__name__) -def execute_cursor(request): - """A Spark cursor whose calculation requests are mocked. - - Args: - request: The pytest fixture request, selecting the cursor class. - - Returns: - The cursor, with ``previous`` as its last calculation ID. - """ - cursor = request.param.__new__(request.param) # bypass __init__ to avoid AWS calls - cursor._session_id = "session_id" - cursor._calculation_id = "previous" - cursor._calculation_execution = None - mock = AsyncMock if request.param is AioSparkCursor else MagicMock - cursor._calculate = mock(return_value="calculation_id") - cursor._poll = mock( - return_value=MagicMock(state=AthenaCalculationExecutionStatus.STATE_COMPLETED) - ) - cursor._executor = MagicMock() - return cursor - - -async def _execute(cursor, *args, **kwargs): - """Call ``execute()``, awaiting it on ``AioSparkCursor``. - - Args: - cursor: The Spark cursor. - *args: Positional arguments for ``execute()``. - **kwargs: Keyword arguments for ``execute()``. - - Returns: - The result of ``execute()``. - """ - if isinstance(cursor, AioSparkCursor): - return await cursor.execute(*args, **kwargs) - return cursor.execute(*args, **kwargs) - - -class TestSparkExecute: - @pytest.mark.parametrize("parameters", [{}, [], {"value": 1}]) - async def test_parameters_rejected(self, execute_cursor, parameters): - with pytest.raises(NotSupportedError, match="do not support parameters"): - await _execute(execute_cursor, "code", parameters) - - execute_cursor._calculate.assert_not_called() - assert execute_cursor.calculation_id == "previous" - - async def test_work_group_rejected(self, execute_cursor): - with pytest.raises(TypeError, match="unexpected keyword argument 'work_group'"): - await _execute(execute_cursor, "code", work_group=None) - - execute_cursor._calculate.assert_not_called() - - @pytest.mark.parametrize( - ("args", "kwargs"), [((), {}), ((None,), {}), ((), {"parameters": None})] - ) - async def test_none_parameters_accepted(self, execute_cursor, args, kwargs): - await _execute(execute_cursor, "code", *args, **kwargs) - - execute_cursor._calculate.assert_called_once_with( - session_id="session_id", - code_block="code", - description=None, - client_request_token=None, - ) - - class TestSparkBaseCursor: @pytest.mark.parametrize( "state", diff --git a/tests/pyathena/spark/test_spark_cursor.py b/tests/pyathena/spark/test_spark_cursor.py index b8e6433ba..2676ad689 100644 --- a/tests/pyathena/spark/test_spark_cursor.py +++ b/tests/pyathena/spark/test_spark_cursor.py @@ -12,7 +12,7 @@ import pytest -from pyathena import OperationalError +from pyathena import NotSupportedError, OperationalError from pyathena.model import AthenaCalculationExecutionStatus, AthenaSessionStatus from pyathena.spark.cursor import SparkCursor from tests import ENV @@ -268,6 +268,58 @@ def test_execute_interrupt_without_kill_on_interrupt(self): cancel.assert_not_called() assert cursor.calculation_execution is None + @pytest.mark.parametrize("parameters", [{}, [], {"value": 1}]) + def test_execute_rejects_parameters(self, parameters): + """Any parameters other than None raise before a calculation starts (no AWS).""" + cursor = SparkCursor.__new__(SparkCursor) # bypass __init__ to avoid AWS calls + cursor._calculation_id = "previous" + + with ( + patch.object(SparkCursor, "_calculate") as calculate, + pytest.raises(NotSupportedError, match="do not support parameters"), + ): + cursor.execute("code", parameters) + + calculate.assert_not_called() + assert cursor.calculation_id == "previous" + + def test_execute_rejects_work_group(self): + """work_group is not an execute() argument, even when None (no AWS).""" + cursor = SparkCursor.__new__(SparkCursor) # bypass __init__ to avoid AWS calls + + with ( + patch.object(SparkCursor, "_calculate") as calculate, + pytest.raises(TypeError, match="unexpected keyword argument 'work_group'"), + ): + cursor.execute("code", work_group=None) + + calculate.assert_not_called() + + @pytest.mark.parametrize( + ("args", "kwargs"), [((), {}), ((None,), {}), ((), {"parameters": None})] + ) + def test_execute_accepts_none_parameters(self, args, kwargs): + """Omitted or None parameters run the calculation (no AWS).""" + cursor = SparkCursor.__new__(SparkCursor) # bypass __init__ to avoid AWS calls + cursor._session_id = "session_id" + + with ( + patch.object(SparkCursor, "_calculate", return_value="calculation_id") as calculate, + patch.object( + SparkCursor, + "_poll", + return_value=MagicMock(state=AthenaCalculationExecutionStatus.STATE_COMPLETED), + ), + ): + cursor.execute("code", *args, **kwargs) + + calculate.assert_called_once_with( + session_id="session_id", + code_block="code", + description=None, + client_request_token=None, + ) + def test_spark_on_poll_invoked_each_iteration(): """on_poll fires once per Spark calculation poll iteration with the status (no AWS).""" From 6e9149616721c2b4e7a6ba356ab2b4ce5cb10108 Mon Sep 17 00:00:00 2001 From: laughingman7743 Date: Sat, 10 Oct 2026 12:54:55 +0900 Subject: [PATCH 5/5] Check that rejected Spark arguments keep the previous calculation ID Co-Authored-By: Claude Opus 5.5 --- tests/pyathena/aio/spark/test_cursor.py | 2 ++ tests/pyathena/spark/test_async_cursor.py | 4 ++++ tests/pyathena/spark/test_spark_cursor.py | 2 ++ 3 files changed, 8 insertions(+) diff --git a/tests/pyathena/aio/spark/test_cursor.py b/tests/pyathena/aio/spark/test_cursor.py index 0a175bf81..8be05d274 100644 --- a/tests/pyathena/aio/spark/test_cursor.py +++ b/tests/pyathena/aio/spark/test_cursor.py @@ -520,12 +520,14 @@ async def test_execute_rejects_parameters(self, parameters): async def test_execute_rejects_work_group(self): """work_group is not an execute() argument, even when None (no AWS).""" cursor = AioSparkCursor.__new__(AioSparkCursor) # bypass __init__ to avoid AWS calls + cursor._calculation_id = "previous" cursor._calculate = AsyncMock() with pytest.raises(TypeError, match="unexpected keyword argument 'work_group'"): await cursor.execute("code", work_group=None) cursor._calculate.assert_not_awaited() + assert cursor.calculation_id == "previous" @pytest.mark.parametrize( ("args", "kwargs"), [((), {}), ((None,), {}), ((), {"parameters": None})] diff --git a/tests/pyathena/spark/test_async_cursor.py b/tests/pyathena/spark/test_async_cursor.py index 154ffe34f..bcc332ba1 100644 --- a/tests/pyathena/spark/test_async_cursor.py +++ b/tests/pyathena/spark/test_async_cursor.py @@ -261,6 +261,7 @@ def test_close_retries_termination_after_failure(self): def test_execute_rejects_parameters(self, parameters): """Any parameters other than None raise before a calculation starts (no AWS).""" cursor = AsyncSparkCursor.__new__(AsyncSparkCursor) # bypass __init__ to avoid AWS calls + cursor._calculation_id = "previous" cursor._executor = MagicMock() with ( @@ -271,10 +272,12 @@ def test_execute_rejects_parameters(self, parameters): calculate.assert_not_called() cursor._executor.submit.assert_not_called() + assert cursor.calculation_id == "previous" def test_execute_rejects_work_group(self): """work_group is not an execute() argument, even when None (no AWS).""" cursor = AsyncSparkCursor.__new__(AsyncSparkCursor) # bypass __init__ to avoid AWS calls + cursor._calculation_id = "previous" cursor._executor = MagicMock() with ( @@ -285,6 +288,7 @@ def test_execute_rejects_work_group(self): calculate.assert_not_called() cursor._executor.submit.assert_not_called() + assert cursor.calculation_id == "previous" @pytest.mark.parametrize( ("args", "kwargs"), [((), {}), ((None,), {}), ((), {"parameters": None})] diff --git a/tests/pyathena/spark/test_spark_cursor.py b/tests/pyathena/spark/test_spark_cursor.py index 2676ad689..39a045d76 100644 --- a/tests/pyathena/spark/test_spark_cursor.py +++ b/tests/pyathena/spark/test_spark_cursor.py @@ -286,6 +286,7 @@ def test_execute_rejects_parameters(self, parameters): def test_execute_rejects_work_group(self): """work_group is not an execute() argument, even when None (no AWS).""" cursor = SparkCursor.__new__(SparkCursor) # bypass __init__ to avoid AWS calls + cursor._calculation_id = "previous" with ( patch.object(SparkCursor, "_calculate") as calculate, @@ -294,6 +295,7 @@ def test_execute_rejects_work_group(self): cursor.execute("code", work_group=None) calculate.assert_not_called() + assert cursor.calculation_id == "previous" @pytest.mark.parametrize( ("args", "kwargs"), [((), {}), ((None,), {}), ((), {"parameters": None})]