Backport #1116 to 3.x: Read Arrow results with the connection session's credentials

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.

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]:

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