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/__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 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/spark/test_cursor.py b/tests/pyathena/aio/spark/test_cursor.py index 228e07d83..8be05d274 100644 --- a/tests/pyathena/aio/spark/test_cursor.py +++ b/tests/pyathena/aio/spark/test_cursor.py @@ -504,6 +504,52 @@ 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._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})] + ) + 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/aio/test_connection.py b/tests/pyathena/aio/test_connection.py index 6fb2cb59a..624cecaf1 100644 --- a/tests/pyathena/aio/test_connection.py +++ b/tests/pyathena/aio/test_connection.py @@ -32,6 +32,22 @@ READER_CURSOR_CLASSES = [AioPandasCursor, AioPolarsCursor] +@pytest.mark.parametrize("factory", [pyathena.aio_connect, AioConnection.create]) +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", + ) + try: + assert conn.work_group is None + finally: + conn.close() + + @pytest.mark.parametrize("factory", [pyathena.aio_connect, AioConnection.create]) async def test_connect_keyword_only(factory): assert all( @@ -73,9 +89,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_async_cursor.py b/tests/pyathena/spark/test_async_cursor.py index 97e726852..bcc332ba1 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,59 @@ 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._calculation_id = "previous" + 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() + 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 ( + 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() + assert cursor.calculation_id == "previous" + + @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_spark_cursor.py b/tests/pyathena/spark/test_spark_cursor.py index b8e6433ba..39a045d76 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,60 @@ 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 + cursor._calculation_id = "previous" + + 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() + assert cursor.calculation_id == "previous" + + @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).""" diff --git a/tests/pyathena/test_connection.py b/tests/pyathena/test_connection.py index 8b9171a36..6df0b1686 100644 --- a/tests/pyathena/test_connection.py +++ b/tests/pyathena/test_connection.py @@ -55,6 +55,43 @@ READER_CURSOR_CLASSES = [PandasCursor, PolarsCursor, AsyncPandasCursor, AsyncPolarsCursor] +@pytest.mark.parametrize("factory", [Connection, pyathena.connect]) +@pytest.mark.parametrize( + ("setting", "kwargs", "expected"), + [ + ("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, 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", + **kwargs, + ) + try: + assert getattr(conn, setting) == expected + finally: + 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 pytest.raises(ProgrammingError, match="Required argument"): + factory(region_name="us-east-1", s3_staging_dir="", work_group="") + + @pytest.mark.parametrize("factory", [pyathena.connect, Connection]) def test_connect_keyword_only(factory): assert all( @@ -91,9 +128,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")