Support assume-role for Glue and DynamoDB catalogs

mkdocs/docs/configuration.md

@@ -838,6 +838,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.             |
 
@@ -896,6 +898,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-->
 
@@ -940,8 +944,8 @@ 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.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:...    | 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. |
 
 <!-- prettier-ignore-start -->
 

pyiceberg/catalog/__init__.py

@@ -74,6 +74,7 @@
 from pyiceberg.view.metadata import ViewVersion
 
 if TYPE_CHECKING:
+    import boto3
     import pyarrow as pa
 
 logger = logging.getLogger(__name__)
@@ -97,6 +98,54 @@
 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

pyiceberg/catalog/dynamodb.py

@@ -34,6 +34,7 @@
     TABLE_TYPE,
     MetastoreCatalog,
     PropertiesUpdateSummary,
+    _get_aws_session_with_assumed_role,
 )
 from pyiceberg.exceptions import (
     ConditionalCheckFailedException,
@@ -46,7 +47,15 @@
     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_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 +101,8 @@
 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 +118,22 @@ def __init__(self, name: str, client: Optional["DynamoDBClient"] = None, **prope
         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_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_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)

pyiceberg/catalog/glue.py

@@ -39,6 +39,7 @@
     TABLE_TYPE,
     MetastoreCatalog,
     PropertiesUpdateSummary,
+    _get_aws_session_with_assumed_role,
 )
 from pyiceberg.exceptions import (
     CommitFailedException,
@@ -50,7 +51,16 @@
     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_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 +141,8 @@
 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"
@@ -335,14 +347,22 @@ def __init__(self, name: str, client: Optional["GlueClient"] = None, **propertie
         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_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_ROLE_SESSION_NAME),
+                    region_name=region_name,
+                )
             self.glue: GlueClient = session.client(
                 "glue",
                 endpoint_url=properties.get(GLUE_CATALOG_ENDPOINT),

tests/catalog/test_dynamodb.py

@@ -21,7 +21,7 @@
 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,61 @@ 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_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