diff --git a/docs/arrow.md b/docs/arrow.md index 5a280d6e3..2cf250a49 100644 --- a/docs/arrow.md +++ b/docs/arrow.md @@ -256,18 +256,16 @@ pyathena.error.OperationalError: SYNTAX_ERROR: line 1:1: Column name not specifi ArrowCursor supports configuring S3 connection and request timeouts through `connect_timeout` and `request_timeout` parameters. These parameters are particularly useful when experiencing timeout errors due to: -- Role assumption with AWS STS (cross-account access) - High network latency between your environment and S3 - Connecting from regions far from the S3 bucket By default, PyArrow uses AWS SDK default timeouts (typically 1 second for connection, 3 seconds for requests). -You can increase these values to accommodate slower authentication or network conditions. +You can increase these values to accommodate slower network conditions. ```python from pyathena import connect from pyathena.arrow.cursor import ArrowCursor -# Configure higher timeouts for role assumption scenarios cursor = connect( s3_staging_dir="s3://YOUR_S3_BUCKET/path/to/", region_name="us-west-2", @@ -291,6 +289,7 @@ cursor = connect( The timeout parameters accept float values in seconds and apply to all S3 operations performed by the cursor, including HeadObject and GetObject operations when retrieving query results. +The pyarrow S3 filesystem uses the credentials of the connection's boto3 session, including the credentials of an assumed role. ```{note} These timeout parameters require PyArrow >= 10.0.0, which added support for configuring S3FileSystem timeouts. diff --git a/pyathena/arrow/result_set.py b/pyathena/arrow/result_set.py index d57838203..40e28218c 100644 --- a/pyathena/arrow/result_set.py +++ b/pyathena/arrow/result_set.py @@ -127,69 +127,18 @@ def __init__( def __s3_file_system(self): from pyarrow import fs - connection = self.connection - - # Build timeout parameters dict - timeout_kwargs = {} + kwargs: dict[str, Any] = {"region": self.connection.region_name} if self._connect_timeout is not None: - timeout_kwargs["connect_timeout"] = self._connect_timeout + kwargs["connect_timeout"] = self._connect_timeout if self._request_timeout is not None: - timeout_kwargs["request_timeout"] = self._request_timeout - - if connection._kwargs.get("role_arn"): - external_id = connection._kwargs.get("external_id") - fs = fs.S3FileSystem( - role_arn=connection._kwargs["role_arn"], - session_name=connection._kwargs["role_session_name"], - external_id="" if external_id is None else external_id, - load_frequency=connection._kwargs["duration_seconds"], - region=connection.region_name, - **timeout_kwargs, - ) - elif connection.profile_name: - profile = connection.session._session.full_config["profiles"][connection.profile_name] - fs = fs.S3FileSystem( - access_key=profile.get("aws_access_key_id", None), - secret_key=profile.get("aws_secret_access_key", None), - session_token=profile.get("aws_session_token", None), - region=connection.region_name, - **timeout_kwargs, - ) - else: - # Try explicit credentials first - explicit_access_key = connection._kwargs.get("aws_access_key_id") - explicit_secret_key = connection._kwargs.get("aws_secret_access_key") - - if explicit_access_key and explicit_secret_key: - # Use explicitly provided credentials - fs = fs.S3FileSystem( - access_key=explicit_access_key, - secret_key=explicit_secret_key, - session_token=connection._kwargs.get("aws_session_token"), - region=connection.region_name, - **timeout_kwargs, - ) - else: - # Fall back to dynamic credentials from boto3 session - # This handles EC2 instance profiles, temporary credentials, etc. - try: - credentials = connection.session._session.get_credentials() - if credentials: - fs = fs.S3FileSystem( - access_key=credentials.access_key, - secret_key=credentials.secret_key, - session_token=credentials.token, - region=connection.region_name, - **timeout_kwargs, - ) - else: - # Fall back to default (no explicit credentials) - fs = fs.S3FileSystem(region=connection.region_name, **timeout_kwargs) - except Exception: - # Fall back to default if credential retrieval fails - fs = fs.S3FileSystem(region=connection.region_name, **timeout_kwargs) - - return fs + kwargs["request_timeout"] = self._request_timeout + credentials = self.connection.session.get_credentials() + if credentials: + frozen_credentials = credentials.get_frozen_credentials() + kwargs["access_key"] = frozen_credentials.access_key + kwargs["secret_key"] = frozen_credentials.secret_key + kwargs["session_token"] = frozen_credentials.token + return fs.S3FileSystem(**kwargs) @property def timestamp_parsers(self) -> list[str]: diff --git a/tests/pyathena/arrow/test_result_set.py b/tests/pyathena/arrow/test_result_set.py new file mode 100644 index 000000000..cd9baf3c1 --- /dev/null +++ b/tests/pyathena/arrow/test_result_set.py @@ -0,0 +1,73 @@ +# Copyright 2026 The PyAthena authors +# +# Licensed under the MIT License. +# See LICENSE or https://opensource.org/licenses/MIT. +# +# SPDX-License-Identifier: MIT +import os +import sys +from unittest.mock import patch + +import pytest + +from pyathena.arrow.result_set import AthenaArrowResultSet +from pyathena.connection import Connection + + +class TestAthenaArrowResultSet: + @pytest.mark.parametrize( + ("connection_kwargs", "expected_access_key"), + [ + ({}, "DEFAULTKEY"), + ({"profile_name": "static"}, "STATICKEY"), + ({"profile_name": "process"}, "PROCESSKEY"), + ( + {"profile_name": "process", "role_arn": "arn:aws:iam::123456789012:role/r"}, + "ROLEKEY", + ), + ], + ids=["default", "static", "credential_process", "role_arn"], + ) + def test_create_s3_file_system_uses_session_credentials( + self, monkeypatch, tmp_path, connection_kwargs, expected_access_key + ): + """The pyarrow filesystem reads with the credentials of the connection's session. + + No AWS calls; the profiles come from temporary config files, and the role + assumption is mocked. + """ + for key in list(os.environ): + if key.startswith("AWS_"): + monkeypatch.delenv(key) + process = tmp_path / "credential_process.py" + process.write_text( + 'print(\'{"Version": 1, "AccessKeyId": "PROCESSKEY", "SecretAccessKey": "secret"}\')\n' + ) + config = tmp_path / "config" + config.write_text( + "[default]\naws_access_key_id = DEFAULTKEY\naws_secret_access_key = secret\n" + "[profile static]\naws_access_key_id = STATICKEY\naws_secret_access_key = secret\n" + f'[profile process]\ncredential_process = "{sys.executable}" "{process}"\n' + ) + monkeypatch.setenv("AWS_CONFIG_FILE", str(config)) + monkeypatch.setenv("AWS_SHARED_CREDENTIALS_FILE", str(tmp_path / "credentials")) + role_credentials = { + "AccessKeyId": "ROLEKEY", + "SecretAccessKey": "secret", + "SessionToken": "token", + } + with patch.object(Connection, "_assume_role", return_value=role_credentials): + connection = Connection( + region_name="us-east-1", s3_staging_dir="s3://bucket/path/", **connection_kwargs + ) + with connection: + result_set = AthenaArrowResultSet.__new__(AthenaArrowResultSet) # bypass __init__ + result_set._connection = connection + result_set._connect_timeout = None + result_set._request_timeout = None + fs = result_set._AthenaArrowResultSet__s3_file_system() + + # pyarrow exposes the filesystem options only through its pickle support. + options = fs.__reduce__()[1][0] + assert options["access_key"] == expected_access_key + assert options["role_arn"] == ""