diff --git a/mkdocs/docs/configuration.md b/mkdocs/docs/configuration.md index c01f92fa4b..2b280d2e2b 100644 --- a/mkdocs/docs/configuration.md +++ b/mkdocs/docs/configuration.md @@ -805,6 +805,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. | @@ -863,6 +865,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 | @@ -907,8 +911,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. | diff --git a/pyiceberg/catalog/__init__.py b/pyiceberg/catalog/__init__.py index 95ceaa539f..2f4c095ba0 100644 --- a/pyiceberg/catalog/__init__.py +++ b/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,46 @@ 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: + """Assume an IAM role via STS and return a new boto3 session using the temporary credentials. + + 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. + + Args: + session: The base boto3 session 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 the assumed-role temporary credentials. + """ + import boto3 + + sts_client = session.client("sts", region_name=region_name) + response = sts_client.assume_role( + RoleArn=role_arn, + RoleSessionName=role_session_name or DEFAULT_ROLE_SESSION_NAME, + ) + credentials = response["Credentials"] + return boto3.Session( + aws_access_key_id=credentials["AccessKeyId"], + aws_secret_access_key=credentials["SecretAccessKey"], + aws_session_token=credentials["SessionToken"], + 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 10f74fb9a0..a1b631b497 100644 --- a/pyiceberg/catalog/dynamodb.py +++ b/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) diff --git a/pyiceberg/catalog/glue.py b/pyiceberg/catalog/glue.py index 12b36efc5c..ce47266271 100644 --- a/pyiceberg/catalog/glue.py +++ b/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), diff --git a/tests/catalog/test_dynamodb.py b/tests/catalog/test_dynamodb.py index 5933e7d472..b470ade341 100644 --- a/tests/catalog/test_dynamodb.py +++ b/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, @@ -613,6 +613,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 diff --git a/tests/catalog/test_glue.py b/tests/catalog/test_glue.py index 89646f5719..a1aa456cc4 100644 --- a/tests/catalog/test_glue.py +++ b/tests/catalog/test_glue.py @@ -21,6 +21,7 @@ import pytest 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, @@ -751,6 +752,83 @@ 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_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_commit_table_update_schema( _glue: boto3.client,