This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-4132-1e6e94a44f84dd3035dcff19546306f89ca2f74f in repository https://gitbox.apache.org/repos/asf/iceberg-python.git
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",
