This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/iceberg-python.git
The following commit(s) were added to refs/heads/main by this push:
new 27a516274 Support assume-role for Glue and DynamoDB catalogs (#4132)
27a516274 is described below
commit 27a5162749bb670d00b23b8360869410b6e5b3a1
Author: David Banquet <[email protected]>
AuthorDate: Tue Oct 27 23:08:26 2026 +0000
Support assume-role for Glue and DynamoDB catalogs (#4132)
* Support assume-role for Glue and DynamoDB catalogs
Extend the existing S3 FileIO assume-role mechanism (client.role-arn /
client.role-session-name) to the boto3 sessions built by GlueCatalog and
DynamoDbCatalog. When a role ARN is configured, an STS AssumeRole is
performed from the base session and the resulting temporary credentials
are used for the Glue/DynamoDB client.
Role configuration follows the same precedence already used for the other
client.* properties: a service-prefixed override (glue.role-arn /
dynamodb.role-arn, glue.role-session-name / dynamodb.role-session-name)
falls back to the unified client.role-arn / client.role-session-name.
Default behavior (no role ARN configured) is unchanged.
Closes #2747
Signed-off-by: lemon0333 <[email protected]>
* Refresh assumed-role credentials for the Glue and DynamoDB catalogs
The catalog session was built from the static credentials of a single
AssumeRole call, so a long-lived catalog failed with ExpiredToken once the
role session ended (one hour by default). Use botocore's
AssumeRoleCredentialFetcher with DeferredRefreshableCredentials: the role is
assumed on first use and again before the credentials expire.
Co-authored-by: lemon0333 <[email protected]>
Co-Authored-By: Claude Opus 5.5 <[email protected]>
* Read the Java assume-role keys before client.role-*
Java uses client.assume-role.arn and client.assume-role.session-name. The
Glue and DynamoDB catalogs and the S3 FileIO now read them first, so the
same properties work for both when a catalog passes them. The existing
client.role-arn / client.role-session-name keys stay as a fallback, and the
service-prefixed overrides keep their precedence.
Co-Authored-By: Claude Opus 5.5 <[email protected]>
---------
Signed-off-by: lemon0333 <[email protected]>
Co-authored-by: lemon0333 <[email protected]>
Co-authored-by: Claude Opus 5.5 <[email protected]>
---
mkdocs/docs/configuration.md | 10 +++-
pyiceberg/catalog/__init__.py | 49 +++++++++++++++
pyiceberg/catalog/dynamodb.py | 27 ++++++++-
pyiceberg/catalog/glue.py | 28 ++++++++-
pyiceberg/io/__init__.py | 2 +
pyiceberg/io/pyarrow.py | 14 +++--
tests/catalog/test_dynamodb.py | 81 ++++++++++++++++++++++++-
tests/catalog/test_glue.py | 132 +++++++++++++++++++++++++++++++++++++++++
tests/io/test_pyarrow.py | 17 ++++++
9 files changed, 349 insertions(+), 11 deletions(-)
diff --git a/mkdocs/docs/configuration.md b/mkdocs/docs/configuration.md
index da336cdcf..eba97908b 100644
--- a/mkdocs/docs/configuration.md
+++ b/mkdocs/docs/configuration.md
@@ -845,6 +845,8 @@ catalog:
| glue.access-key-id | admin | Configure
the static access key id used to access the Glue Catalog |
| glue.secret-access-key | password | Configure
the static secret access key used to access the Glue Catalog |
| glue.session-token | AQoDYXdzEJr... | Configure
the static session token used to access the Glue Catalog |
+| glue.role-arn | arn:aws:... | AWS Role
ARN. If provided, temporary credentials will be fetched by assuming this role
before accessing the Glue Catalog |
+| glue.role-session-name | session | An
optional identifier for the assumed role session used to access the Glue
Catalog |
| glue.max-retries | 10 | Configure
the maximum number of retries for the Glue service calls |
| glue.retry-mode | standard | Configure
the retry mode for the Glue service. Default to standard. |
@@ -903,6 +905,8 @@ catalog:
| dynamodb.access-key-id | admin | Configure the static access
key id used to access the DynamoDB Catalog |
| dynamodb.secret-access-key | password | Configure the static secret
access key used to access the DynamoDB Catalog |
| dynamodb.session-token | AQoDYXdzEJr... | Configure the static session
token used to access the DynamoDB Catalog |
+| dynamodb.role-arn | arn:aws:... | AWS Role ARN. If provided,
temporary credentials will be fetched by assuming this role before accessing
the DynamoDB Catalog |
+| dynamodb.role-session-name | session | An optional identifier for the
assumed role session used to access the DynamoDB Catalog |
<!-- markdown-link-check-enable-->
@@ -947,8 +951,10 @@ configures the AWS credentials for both Glue Catalog and
S3 FileIO.
| client.secret-access-key | password | Configure the static secret
access key used to access both the Glue/DynamoDB Catalog and the S3 FileIO
|
| client.session-token | AQoDYXdzEJr... | Configure the static session
token used to access both the Glue/DynamoDB Catalog and the S3 FileIO
|
| client.profile-name | default | Configure the AWS profile used
to access both the Glue/DynamoDB Catalog and the S3 FileIO (only supported by
`FsspecFileIO` currently) |
-| client.role-session-name | session | An optional identifier for the
assumed role session.
|
-| client.role-arn | arn:aws:... | AWS Role ARN. If provided
instead of access_key and secret_key, temporary credentials will be fetched by
assuming this role. |
+| client.assume-role.arn | arn:aws:... | AWS Role ARN. If provided
instead of access_key and secret_key, temporary credentials will be fetched by
assuming this role before accessing both the Glue/DynamoDB Catalog and the S3
FileIO. |
+| client.assume-role.session-name | session | An optional identifier
for the assumed role session used to access both the Glue/DynamoDB Catalog and
the S3 FileIO. |
+| client.role-arn | arn:aws:... | Same as
`client.assume-role.arn`, used when that key is not set. |
+| client.role-session-name | session | Same as
`client.assume-role.session-name`, used when that key is not set. |
<!-- prettier-ignore-start -->
diff --git a/pyiceberg/catalog/__init__.py b/pyiceberg/catalog/__init__.py
index 022783b08..cdb4d26b3 100644
--- a/pyiceberg/catalog/__init__.py
+++ b/pyiceberg/catalog/__init__.py
@@ -74,6 +74,7 @@ from pyiceberg.view import View
from pyiceberg.view.metadata import ViewVersion
if TYPE_CHECKING:
+ import boto3
import pyarrow as pa
logger = logging.getLogger(__name__)
@@ -97,6 +98,54 @@ LOCATION = "location"
EXTERNAL_TABLE = "EXTERNAL_TABLE"
BOTOCORE_SESSION = "botocore_session"
+DEFAULT_ROLE_SESSION_NAME = "pyiceberg"
+
+
+def _get_aws_session_with_assumed_role(
+ session: boto3.Session,
+ role_arn: str,
+ role_session_name: str | None,
+ region_name: str | None,
+) -> boto3.Session:
+ """Return a boto3 session whose credentials come from assuming an IAM role
via STS.
+
+ This mirrors the assume-role behavior already available for the S3 FileIO
(``client.role-arn`` /
+ ``client.role-session-name``) so that the AWS catalog clients (Glue,
DynamoDB) can honor the same
+ properties. The role is assumed on first use and again before the
temporary credentials expire,
+ so a long-lived catalog keeps working past the role session duration.
+
+ Args:
+ session: The base boto3 session whose credentials are used to call STS.
+ role_arn: The ARN of the role to assume.
+ role_session_name: An identifier for the assumed-role session.
Defaults to ``pyiceberg``.
+ region_name: The region for the returned session.
+
+ Returns:
+ A new boto3 session authenticated with refreshable assumed-role
credentials.
+ """
+ import boto3
+ import botocore.session
+ from botocore.credentials import AssumeRoleCredentialFetcher,
DeferredRefreshableCredentials
+
+ source_credentials = session.get_credentials()
+ if source_credentials is None:
+ raise ValueError(f"No AWS credentials found to assume role {role_arn}")
+
+ fetcher = AssumeRoleCredentialFetcher(
+ client_creator=session.client,
+ source_credentials=source_credentials,
+ role_arn=role_arn,
+ extra_args={"RoleSessionName": role_session_name or
DEFAULT_ROLE_SESSION_NAME},
+ )
+ botocore_session = botocore.session.Session()
+ # botocore has no public setter for refreshable credentials on a session.
+ botocore_session._credentials = DeferredRefreshableCredentials(
+ refresh_using=fetcher.fetch_credentials,
+ method="assume-role",
+ )
+ return boto3.Session(botocore_session=botocore_session,
region_name=region_name)
+
+
TABLE_METADATA_FILE_NAME_REGEX = re.compile(
r"""
(\d+) # version number
diff --git a/pyiceberg/catalog/dynamodb.py b/pyiceberg/catalog/dynamodb.py
index d20c0268b..220313ab7 100644
--- a/pyiceberg/catalog/dynamodb.py
+++ b/pyiceberg/catalog/dynamodb.py
@@ -34,6 +34,7 @@ from pyiceberg.catalog import (
TABLE_TYPE,
MetastoreCatalog,
PropertiesUpdateSummary,
+ _get_aws_session_with_assumed_role,
)
from pyiceberg.exceptions import (
ConditionalCheckFailedException,
@@ -46,7 +47,17 @@ from pyiceberg.exceptions import (
NoSuchTableError,
TableAlreadyExistsError,
)
-from pyiceberg.io import AWS_ACCESS_KEY_ID, AWS_REGION, AWS_SECRET_ACCESS_KEY,
AWS_SESSION_TOKEN, load_file_io
+from pyiceberg.io import (
+ AWS_ACCESS_KEY_ID,
+ AWS_ASSUME_ROLE_ARN,
+ AWS_ASSUME_ROLE_SESSION_NAME,
+ AWS_REGION,
+ AWS_ROLE_ARN,
+ AWS_ROLE_SESSION_NAME,
+ AWS_SECRET_ACCESS_KEY,
+ AWS_SESSION_TOKEN,
+ load_file_io,
+)
from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC, PartitionSpec
from pyiceberg.schema import Schema
from pyiceberg.serializers import FromInputFile
@@ -92,6 +103,8 @@ DYNAMODB_REGION = "dynamodb.region"
DYNAMODB_ACCESS_KEY_ID = "dynamodb.access-key-id"
DYNAMODB_SECRET_ACCESS_KEY = "dynamodb.secret-access-key"
DYNAMODB_SESSION_TOKEN = "dynamodb.session-token"
+DYNAMODB_ROLE_ARN = "dynamodb.role-arn"
+DYNAMODB_ROLE_SESSION_NAME = "dynamodb.role-session-name"
class DynamoDbCatalog(MetastoreCatalog):
@@ -107,14 +120,24 @@ class DynamoDbCatalog(MetastoreCatalog):
if client is not None:
self.dynamodb = client
else:
+ region_name = get_first_property_value(properties,
DYNAMODB_REGION, AWS_REGION)
session = boto3.Session(
profile_name=properties.get(DYNAMODB_PROFILE_NAME),
- region_name=get_first_property_value(properties,
DYNAMODB_REGION, AWS_REGION),
+ region_name=region_name,
botocore_session=properties.get(BOTOCORE_SESSION),
aws_access_key_id=get_first_property_value(properties,
DYNAMODB_ACCESS_KEY_ID, AWS_ACCESS_KEY_ID),
aws_secret_access_key=get_first_property_value(properties,
DYNAMODB_SECRET_ACCESS_KEY, AWS_SECRET_ACCESS_KEY),
aws_session_token=get_first_property_value(properties,
DYNAMODB_SESSION_TOKEN, AWS_SESSION_TOKEN),
)
+ if role_arn := get_first_property_value(properties,
DYNAMODB_ROLE_ARN, AWS_ASSUME_ROLE_ARN, AWS_ROLE_ARN):
+ session = _get_aws_session_with_assumed_role(
+ session=session,
+ role_arn=role_arn,
+ role_session_name=get_first_property_value(
+ properties, DYNAMODB_ROLE_SESSION_NAME,
AWS_ASSUME_ROLE_SESSION_NAME, AWS_ROLE_SESSION_NAME
+ ),
+ region_name=region_name,
+ )
self.dynamodb = session.client(DYNAMODB_CLIENT)
self.dynamodb_table_name = self.properties.get(DYNAMODB_TABLE_NAME,
DYNAMODB_TABLE_NAME_DEFAULT)
diff --git a/pyiceberg/catalog/glue.py b/pyiceberg/catalog/glue.py
index e678a8b6f..bbb954fcd 100644
--- a/pyiceberg/catalog/glue.py
+++ b/pyiceberg/catalog/glue.py
@@ -39,6 +39,7 @@ from pyiceberg.catalog import (
TABLE_TYPE,
MetastoreCatalog,
PropertiesUpdateSummary,
+ _get_aws_session_with_assumed_role,
)
from pyiceberg.exceptions import (
CommitFailedException,
@@ -50,7 +51,18 @@ from pyiceberg.exceptions import (
NoSuchTableError,
TableAlreadyExistsError,
)
-from pyiceberg.io import AWS_ACCESS_KEY_ID, AWS_PROFILE_NAME, AWS_REGION,
AWS_SECRET_ACCESS_KEY, AWS_SESSION_TOKEN, FileIO
+from pyiceberg.io import (
+ AWS_ACCESS_KEY_ID,
+ AWS_ASSUME_ROLE_ARN,
+ AWS_ASSUME_ROLE_SESSION_NAME,
+ AWS_PROFILE_NAME,
+ AWS_REGION,
+ AWS_ROLE_ARN,
+ AWS_ROLE_SESSION_NAME,
+ AWS_SECRET_ACCESS_KEY,
+ AWS_SESSION_TOKEN,
+ FileIO,
+)
from pyiceberg.partitioning import UNPARTITIONED_PARTITION_SPEC, PartitionSpec
from pyiceberg.schema import Schema, SchemaVisitor, visit
from pyiceberg.serializers import FromInputFile, ToOutputFile
@@ -131,6 +143,8 @@ GLUE_REGION = "glue.region"
GLUE_ACCESS_KEY_ID = "glue.access-key-id"
GLUE_SECRET_ACCESS_KEY = "glue.secret-access-key"
GLUE_SESSION_TOKEN = "glue.session-token"
+GLUE_ROLE_ARN = "glue.role-arn"
+GLUE_ROLE_SESSION_NAME = "glue.role-session-name"
GLUE_MAX_RETRIES = "glue.max-retries"
GLUE_RETRY_MODE = "glue.retry-mode"
GLUE_CONNECTION_S3_TABLES = "aws:s3tables"
@@ -349,14 +363,24 @@ class GlueCatalog(MetastoreCatalog):
else:
retry_mode_prop_value = get_first_property_value(properties,
GLUE_RETRY_MODE)
+ region_name = get_first_property_value(properties, GLUE_REGION,
AWS_REGION)
session = boto3.Session(
profile_name=get_first_property_value(properties,
GLUE_PROFILE_NAME, AWS_PROFILE_NAME),
- region_name=get_first_property_value(properties, GLUE_REGION,
AWS_REGION),
+ region_name=region_name,
botocore_session=properties.get(BOTOCORE_SESSION),
aws_access_key_id=get_first_property_value(properties,
GLUE_ACCESS_KEY_ID, AWS_ACCESS_KEY_ID),
aws_secret_access_key=get_first_property_value(properties,
GLUE_SECRET_ACCESS_KEY, AWS_SECRET_ACCESS_KEY),
aws_session_token=get_first_property_value(properties,
GLUE_SESSION_TOKEN, AWS_SESSION_TOKEN),
)
+ if role_arn := get_first_property_value(properties, GLUE_ROLE_ARN,
AWS_ASSUME_ROLE_ARN, AWS_ROLE_ARN):
+ session = _get_aws_session_with_assumed_role(
+ session=session,
+ role_arn=role_arn,
+ role_session_name=get_first_property_value(
+ properties, GLUE_ROLE_SESSION_NAME,
AWS_ASSUME_ROLE_SESSION_NAME, AWS_ROLE_SESSION_NAME
+ ),
+ region_name=region_name,
+ )
self.glue: GlueClient = session.client(
"glue",
endpoint_url=properties.get(GLUE_CATALOG_ENDPOINT),
diff --git a/pyiceberg/io/__init__.py b/pyiceberg/io/__init__.py
index 9cae219f7..ec232a02e 100644
--- a/pyiceberg/io/__init__.py
+++ b/pyiceberg/io/__init__.py
@@ -61,6 +61,8 @@ AWS_SECRET_ACCESS_KEY = "client.secret-access-key"
AWS_SESSION_TOKEN = "client.session-token"
AWS_ROLE_ARN = "client.role-arn"
AWS_ROLE_SESSION_NAME = "client.role-session-name"
+AWS_ASSUME_ROLE_ARN = "client.assume-role.arn"
+AWS_ASSUME_ROLE_SESSION_NAME = "client.assume-role.session-name"
S3_PROFILE_NAME = "s3.profile-name"
S3_ANONYMOUS = "s3.anonymous"
S3_ENDPOINT = "s3.endpoint"
diff --git a/pyiceberg/io/pyarrow.py b/pyiceberg/io/pyarrow.py
index 0cab26309..1dd816d5d 100644
--- a/pyiceberg/io/pyarrow.py
+++ b/pyiceberg/io/pyarrow.py
@@ -90,6 +90,8 @@ from pyiceberg.io import (
ADLS_SAS_TOKEN,
ADLS_TENANT_ID,
AWS_ACCESS_KEY_ID,
+ AWS_ASSUME_ROLE_ARN,
+ AWS_ASSUME_ROLE_SESSION_NAME,
AWS_REGION,
AWS_ROLE_ARN,
AWS_ROLE_SESSION_NAME,
@@ -485,10 +487,12 @@ class PyArrowFileIO(FileIO):
if request_timeout := self.properties.get(S3_REQUEST_TIMEOUT):
client_kwargs["request_timeout"] = float(request_timeout)
- if role_arn := get_first_property_value(self.properties, S3_ROLE_ARN,
AWS_ROLE_ARN):
+ if role_arn := get_first_property_value(self.properties, S3_ROLE_ARN,
AWS_ASSUME_ROLE_ARN, AWS_ROLE_ARN):
client_kwargs["role_arn"] = role_arn
- if session_name := get_first_property_value(self.properties,
S3_ROLE_SESSION_NAME, AWS_ROLE_SESSION_NAME):
+ if session_name := get_first_property_value(
+ self.properties, S3_ROLE_SESSION_NAME,
AWS_ASSUME_ROLE_SESSION_NAME, AWS_ROLE_SESSION_NAME
+ ):
client_kwargs["session_name"] = session_name
if s3_anonymous := self.properties.get(S3_ANONYMOUS):
@@ -533,10 +537,12 @@ class PyArrowFileIO(FileIO):
if request_timeout := self.properties.get(S3_REQUEST_TIMEOUT):
client_kwargs["request_timeout"] = float(request_timeout)
- if role_arn := get_first_property_value(self.properties, S3_ROLE_ARN,
AWS_ROLE_ARN):
+ if role_arn := get_first_property_value(self.properties, S3_ROLE_ARN,
AWS_ASSUME_ROLE_ARN, AWS_ROLE_ARN):
client_kwargs["role_arn"] = role_arn
- if session_name := get_first_property_value(self.properties,
S3_ROLE_SESSION_NAME, AWS_ROLE_SESSION_NAME):
+ if session_name := get_first_property_value(
+ self.properties, S3_ROLE_SESSION_NAME,
AWS_ASSUME_ROLE_SESSION_NAME, AWS_ROLE_SESSION_NAME
+ ):
client_kwargs["session_name"] = session_name
if self.properties.get(S3_FORCE_VIRTUAL_ADDRESSING) is not None:
diff --git a/tests/catalog/test_dynamodb.py b/tests/catalog/test_dynamodb.py
index dabaf5e41..803cb2cd8 100644
--- a/tests/catalog/test_dynamodb.py
+++ b/tests/catalog/test_dynamodb.py
@@ -21,7 +21,7 @@ import pyarrow as pa
import pytest
from moto import mock_aws
-from pyiceberg.catalog import METADATA_LOCATION, TABLE_TYPE
+from pyiceberg.catalog import METADATA_LOCATION, TABLE_TYPE,
_get_aws_session_with_assumed_role
from pyiceberg.catalog.dynamodb import (
ACTIVE,
DYNAMODB_COL_CREATED_AT,
@@ -637,6 +637,85 @@ def test_passing_unified_session_properties_to_dynamodb()
-> None:
assert test_catalog.dynamodb is mock_session().client()
+@mock_aws
+def test_dynamodb_assume_role_with_unified_role_properties() -> None:
+ role_arn = "arn:aws:iam::123456789012:role/my-dynamodb-role"
+ session_properties: Properties = {
+ "client.role-arn": role_arn,
+ "client.role-session-name": "my-session",
+ "client.region": "us-east-1",
+ }
+
+ with mock.patch(
+ "pyiceberg.catalog.dynamodb._get_aws_session_with_assumed_role",
+ wraps=_get_aws_session_with_assumed_role,
+ ) as mock_assume_role:
+ DynamoDbCatalog("dynamodb", **session_properties)
+
+ mock_assume_role.assert_called_once_with(
+ session=mock.ANY,
+ role_arn=role_arn,
+ role_session_name="my-session",
+ region_name="us-east-1",
+ )
+
+
+@mock_aws
+def test_dynamodb_assume_role_prefixed_properties_take_precedence() -> None:
+ session_properties: Properties = {
+ "dynamodb.role-arn":
"arn:aws:iam::123456789012:role/dynamodb-specific-role",
+ "dynamodb.role-session-name": "dynamodb-session",
+ "client.role-arn": "arn:aws:iam::123456789012:role/unified-role",
+ "client.role-session-name": "unified-session",
+ "client.region": "us-east-1",
+ }
+
+ with mock.patch(
+ "pyiceberg.catalog.dynamodb._get_aws_session_with_assumed_role",
+ wraps=_get_aws_session_with_assumed_role,
+ ) as mock_assume_role:
+ DynamoDbCatalog("dynamodb", **session_properties)
+
+ mock_assume_role.assert_called_once_with(
+ session=mock.ANY,
+ role_arn="arn:aws:iam::123456789012:role/dynamodb-specific-role",
+ role_session_name="dynamodb-session",
+ region_name="us-east-1",
+ )
+
+
+@mock_aws
+def
test_dynamodb_assume_role_properties_take_precedence_over_role_properties() ->
None:
+ session_properties: Properties = {
+ "client.assume-role.arn": "arn:aws:iam::123456789012:role/assume-role",
+ "client.assume-role.session-name": "assume-role-session",
+ "client.role-arn": "arn:aws:iam::123456789012:role/role",
+ "client.role-session-name": "role-session",
+ "client.region": "us-east-1",
+ }
+
+ with mock.patch(
+ "pyiceberg.catalog.dynamodb._get_aws_session_with_assumed_role",
+ wraps=_get_aws_session_with_assumed_role,
+ ) as mock_assume_role:
+ DynamoDbCatalog("dynamodb", **session_properties)
+
+ mock_assume_role.assert_called_once_with(
+ session=mock.ANY,
+ role_arn="arn:aws:iam::123456789012:role/assume-role",
+ role_session_name="assume-role-session",
+ region_name="us-east-1",
+ )
+
+
+@mock_aws
+def test_dynamodb_no_role_arn_does_not_assume_role() -> None:
+ with
mock.patch("pyiceberg.catalog.dynamodb._get_aws_session_with_assumed_role") as
mock_assume_role:
+ DynamoDbCatalog("dynamodb", **{"client.region": "us-east-1"})
+
+ mock_assume_role.assert_not_called()
+
+
@mock_aws
def test_table_exists(
_bucket_initialize: None, moto_endpoint_url: str, table_schema_nested:
Schema, database_name: str, table_name: str
diff --git a/tests/catalog/test_glue.py b/tests/catalog/test_glue.py
index 8943858b8..63b63e930 100644
--- a/tests/catalog/test_glue.py
+++ b/tests/catalog/test_glue.py
@@ -14,13 +14,16 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+from datetime import datetime, timedelta, timezone
from unittest import mock
import boto3
import pyarrow as pa
import pytest
+from botocore.credentials import AssumeRoleCredentialFetcher
from moto import mock_aws
+from pyiceberg.catalog import _get_aws_session_with_assumed_role
from pyiceberg.catalog.glue import GLUE_CONNECTION_S3_TABLES, GlueCatalog
from pyiceberg.exceptions import (
NamespaceAlreadyExistsError,
@@ -761,6 +764,135 @@ def test_passing_unified_session_properties_to_glue() ->
None:
assert test_catalog.glue is mock_session().client()
+@mock_aws
+def test_glue_assume_role_with_unified_role_properties() -> None:
+ role_arn = "arn:aws:iam::123456789012:role/my-glue-role"
+ session_properties: Properties = {
+ "client.role-arn": role_arn,
+ "client.role-session-name": "my-session",
+ "client.region": "us-east-1",
+ }
+
+ with mock.patch(
+ "pyiceberg.catalog.glue._get_aws_session_with_assumed_role",
+ wraps=_get_aws_session_with_assumed_role,
+ ) as mock_assume_role:
+ GlueCatalog("glue", **session_properties)
+
+ mock_assume_role.assert_called_once_with(
+ session=mock.ANY,
+ role_arn=role_arn,
+ role_session_name="my-session",
+ region_name="us-east-1",
+ )
+
+
+@mock_aws
+def test_glue_assume_role_prefixed_properties_take_precedence() -> None:
+ session_properties: Properties = {
+ "glue.role-arn": "arn:aws:iam::123456789012:role/glue-specific-role",
+ "glue.role-session-name": "glue-session",
+ "client.role-arn": "arn:aws:iam::123456789012:role/unified-role",
+ "client.role-session-name": "unified-session",
+ "client.region": "us-east-1",
+ }
+
+ with mock.patch(
+ "pyiceberg.catalog.glue._get_aws_session_with_assumed_role",
+ wraps=_get_aws_session_with_assumed_role,
+ ) as mock_assume_role:
+ GlueCatalog("glue", **session_properties)
+
+ mock_assume_role.assert_called_once_with(
+ session=mock.ANY,
+ role_arn="arn:aws:iam::123456789012:role/glue-specific-role",
+ role_session_name="glue-session",
+ region_name="us-east-1",
+ )
+
+
+@mock_aws
+def test_glue_assume_role_properties_take_precedence_over_role_properties() ->
None:
+ session_properties: Properties = {
+ "client.assume-role.arn": "arn:aws:iam::123456789012:role/assume-role",
+ "client.assume-role.session-name": "assume-role-session",
+ "client.role-arn": "arn:aws:iam::123456789012:role/role",
+ "client.role-session-name": "role-session",
+ "client.region": "us-east-1",
+ }
+
+ with mock.patch(
+ "pyiceberg.catalog.glue._get_aws_session_with_assumed_role",
+ wraps=_get_aws_session_with_assumed_role,
+ ) as mock_assume_role:
+ GlueCatalog("glue", **session_properties)
+
+ mock_assume_role.assert_called_once_with(
+ session=mock.ANY,
+ role_arn="arn:aws:iam::123456789012:role/assume-role",
+ role_session_name="assume-role-session",
+ region_name="us-east-1",
+ )
+
+
+@mock_aws
+def test_glue_no_role_arn_does_not_assume_role() -> None:
+ with
mock.patch("pyiceberg.catalog.glue._get_aws_session_with_assumed_role") as
mock_assume_role:
+ GlueCatalog("glue", **{"client.region": "us-east-1"})
+
+ mock_assume_role.assert_not_called()
+
+
+@mock_aws
+def test_get_aws_session_with_assumed_role_returns_temporary_credentials() ->
None:
+ base_session = boto3.Session(
+ aws_access_key_id="base-access-key",
+ aws_secret_access_key="base-secret-key",
+ region_name="us-east-1",
+ )
+
+ assumed_session = _get_aws_session_with_assumed_role(
+ session=base_session,
+ role_arn="arn:aws:iam::123456789012:role/my-role",
+ role_session_name="test-session",
+ region_name="us-east-1",
+ )
+
+ credentials = assumed_session.get_credentials()
+ assert credentials is not None
+ # Temporary credentials from STS AssumeRole always carry a session token.
+ assert credentials.token is not None
+ assert assumed_session.region_name == "us-east-1"
+
+
+@mock_aws
+def test_glue_assume_role_credentials_are_refreshed_before_expiry() -> None:
+ catalog = GlueCatalog(
+ "glue",
+ **{"client.role-arn": "arn:aws:iam::123456789012:role/my-glue-role",
"client.region": "us-east-1"},
+ )
+ credentials = catalog.glue._request_signer._credentials
+ now = datetime.now(timezone.utc)
+
+ with mock.patch.object(
+ AssumeRoleCredentialFetcher,
+ "_get_credentials",
+ autospec=True,
+ side_effect=AssumeRoleCredentialFetcher._get_credentials,
+ ) as mock_assume_role:
+ catalog.list_namespaces()
+ assert mock_assume_role.call_count == 1
+
+ # moto issues one-hour credentials: 48 minutes later they are inside
botocore's 15-minute refresh window.
+ later = now + timedelta(minutes=48)
+ with (
+ mock.patch("botocore.credentials._local_now", return_value=later),
+ mock.patch.object(credentials, "_time_fetcher",
return_value=later),
+ ):
+ catalog.list_namespaces()
+ assert mock_assume_role.call_count == 2
+
+
@mock_aws
def test_commit_table_update_schema(
_glue: boto3.client,
diff --git a/tests/io/test_pyarrow.py b/tests/io/test_pyarrow.py
index 0ba8117ef..db336acd7 100644
--- a/tests/io/test_pyarrow.py
+++ b/tests/io/test_pyarrow.py
@@ -419,6 +419,23 @@ def test_pyarrow_s3_session_properties() -> None:
)
+def
test_pyarrow_s3_assume_role_properties_take_precedence_over_role_properties()
-> None:
+ session_properties: Properties = {
+ "s3.region": "us-east-1",
+ "client.assume-role.arn": "arn:aws:iam::123456789012:role/assume-role",
+ "client.assume-role.session-name": "assume-role-session",
+ "client.role-arn": "arn:aws:iam::123456789012:role/role",
+ "client.role-session-name": "role-session",
+ }
+
+ with patch("pyarrow.fs.S3FileSystem") as mock_s3fs,
patch("pyarrow.fs.resolve_s3_region") as mock_s3_region_resolver:
+ mock_s3_region_resolver.side_effect = OSError("S3 bucket is not found")
+
PyArrowFileIO(properties=session_properties).new_input(location=f"s3://warehouse/{uuid.uuid4()}")
+
+ assert mock_s3fs.call_args.kwargs["role_arn"] ==
"arn:aws:iam::123456789012:role/assume-role"
+ assert mock_s3fs.call_args.kwargs["session_name"] ==
"assume-role-session"
+
+
def test_pyarrow_s3_ssl_ca_cert_requires_supported_pyarrow_version() -> None:
session_properties: Properties = {
S3_SSL_CA_CERT: "/path/to/ca.pem",
