From 0002a13d1b26d5fcc3543007b18330d0977f7cc6 Mon Sep 17 00:00:00 2001 From: yusufyildiz34 Date: Mon, 19 Oct 2026 15:45:15 -0400 Subject: [PATCH] AWS: Apply Glue and KMS endpoint overrides in AssumeRoleAwsClientFactory AssumeRoleAwsClientFactory built Glue and KMS clients without the glue.endpoint and kms.endpoint overrides, and LakeFormationAwsClientFactory did the same for KMS clients of registered tables, so these clients always used the default AWS endpoints. Closes #17047 Co-authored-by: Eunbin Son --- .../aws/AssumeRoleAwsClientFactory.java | 2 + .../LakeFormationAwsClientFactory.java | 1 + .../iceberg/aws/TestAwsClientFactories.java | 71 +++++++++++++++++++ 3 files changed, 74 insertions(+) diff --git a/aws/src/main/java/org/apache/iceberg/aws/AssumeRoleAwsClientFactory.java b/aws/src/main/java/org/apache/iceberg/aws/AssumeRoleAwsClientFactory.java index cf9678e54d83..e8168f142db6 100644 --- a/aws/src/main/java/org/apache/iceberg/aws/AssumeRoleAwsClientFactory.java +++ b/aws/src/main/java/org/apache/iceberg/aws/AssumeRoleAwsClientFactory.java @@ -76,6 +76,7 @@ public GlueClient glue() { return GlueClient.builder() .applyMutation(this::applyAssumeRoleConfigurations) .applyMutation(httpClientProperties::applyHttpClientConfigurations) + .applyMutation(awsProperties::applyGlueEndpointConfigurations) .applyMutation(awsClientProperties::applyRetryConfigurations) .build(); } @@ -85,6 +86,7 @@ public KmsClient kms() { return KmsClient.builder() .applyMutation(this::applyAssumeRoleConfigurations) .applyMutation(httpClientProperties::applyHttpClientConfigurations) + .applyMutation(awsProperties::applyKmsEndpointConfigurations) .applyMutation(awsClientProperties::applyRetryConfigurations) .build(); } diff --git a/aws/src/main/java/org/apache/iceberg/aws/lakeformation/LakeFormationAwsClientFactory.java b/aws/src/main/java/org/apache/iceberg/aws/lakeformation/LakeFormationAwsClientFactory.java index 51c579ca1741..0f700357110a 100644 --- a/aws/src/main/java/org/apache/iceberg/aws/lakeformation/LakeFormationAwsClientFactory.java +++ b/aws/src/main/java/org/apache/iceberg/aws/lakeformation/LakeFormationAwsClientFactory.java @@ -97,6 +97,7 @@ public KmsClient kms() { if (isTableRegisteredWithLakeFormation()) { return KmsClient.builder() .applyMutation(httpClientProperties()::applyHttpClientConfigurations) + .applyMutation(awsProperties()::applyKmsEndpointConfigurations) .applyMutation(awsClientProperties()::applyRetryConfigurations) .credentialsProvider( new LakeFormationCredentialsProvider(lakeFormation(), buildTableArn())) diff --git a/aws/src/test/java/org/apache/iceberg/aws/TestAwsClientFactories.java b/aws/src/test/java/org/apache/iceberg/aws/TestAwsClientFactories.java index 94ef60e05d68..65a690982c5b 100644 --- a/aws/src/test/java/org/apache/iceberg/aws/TestAwsClientFactories.java +++ b/aws/src/test/java/org/apache/iceberg/aws/TestAwsClientFactories.java @@ -22,6 +22,7 @@ import static org.assertj.core.api.Assertions.assertThatThrownBy; import java.io.IOException; +import java.net.URI; import java.util.Map; import java.util.Optional; import org.apache.iceberg.TestHelpers; @@ -36,6 +37,7 @@ import org.junit.jupiter.api.Test; import org.junit.jupiter.params.ParameterizedTest; import org.junit.jupiter.params.provider.MethodSource; +import org.mockito.Mockito; import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; import software.amazon.awssdk.auth.credentials.AwsCredentials; import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; @@ -45,7 +47,10 @@ import software.amazon.awssdk.retries.internal.DefaultAdaptiveRetryStrategy; import software.amazon.awssdk.services.dynamodb.DynamoDbClient; import software.amazon.awssdk.services.glue.GlueClient; +import software.amazon.awssdk.services.glue.model.GetTableRequest; +import software.amazon.awssdk.services.glue.model.GetTableResponse; import software.amazon.awssdk.services.glue.model.GetTablesRequest; +import software.amazon.awssdk.services.glue.model.Table; import software.amazon.awssdk.services.kms.KmsClient; import software.amazon.awssdk.services.s3.S3AsyncClient; import software.amazon.awssdk.services.s3.S3Client; @@ -273,6 +278,58 @@ public void testDynamoClientSetsAdaptiveRetryPolicy() { assertAwsClientSetsAdaptiveRetryPolicy(dynamoClient); } + @Test + void assumeRoleGlueClientAppliesEndpointOverride() { + String endpoint = "https://glue.example.com"; + AwsClientFactory factory = assumeRoleFactory(AwsProperties.GLUE_CATALOG_ENDPOINT, endpoint); + + assertThat(factory.glue().serviceClientConfiguration().endpointOverride()) + .contains(URI.create(endpoint)); + } + + @Test + void assumeRoleKmsClientAppliesEndpointOverride() { + String endpoint = "https://kms.example.com"; + AwsClientFactory factory = assumeRoleFactory(AwsProperties.KMS_ENDPOINT, endpoint); + + assertThat(factory.kms().serviceClientConfiguration().endpointOverride()) + .contains(URI.create(endpoint)); + } + + @Test + void lakeFormationKmsClientAppliesEndpointOverride() { + String endpoint = "https://kms.example.com"; + Map properties = Maps.newHashMap(); + properties.put(AwsProperties.CLIENT_ASSUME_ROLE_ARN, "arn::test"); + properties.put(AwsProperties.CLIENT_ASSUME_ROLE_REGION, "us-east-1"); + properties.put( + AwsProperties.CLIENT_ASSUME_ROLE_TAGS_PREFIX + + LakeFormationAwsClientFactory.LF_AUTHORIZED_CALLER, + "emr"); + properties.put(AwsProperties.LAKE_FORMATION_DB_NAME, "db"); + properties.put(AwsProperties.LAKE_FORMATION_TABLE_NAME, "table"); + properties.put(AwsProperties.GLUE_ACCOUNT_ID, "123456789012"); + properties.put(AwsProperties.KMS_ENDPOINT, endpoint); + LakeFormationAwsClientFactory factory = new RegisteredTableLakeFormationAwsClientFactory(); + factory.initialize(properties); + + assertThat(factory.kms().serviceClientConfiguration().endpointOverride()) + .contains(URI.create(endpoint)); + } + + private static AwsClientFactory assumeRoleFactory(String property, String value) { + return AwsClientFactories.from( + ImmutableMap.of( + AwsProperties.CLIENT_FACTORY, + AssumeRoleAwsClientFactory.class.getName(), + AwsProperties.CLIENT_ASSUME_ROLE_ARN, + "arn::test", + AwsProperties.CLIENT_ASSUME_ROLE_REGION, + "us-east-1", + property, + value)); + } + /** * Extract the retry strategy from an AwsClient object, then assert that it's set to the correct * strategy. This enforces that we correctly applied the retry configurations to the client @@ -415,4 +472,18 @@ public DynamoDbClient dynamo() { @Override public void initialize(Map properties) {} } + + private static class RegisteredTableLakeFormationAwsClientFactory + extends LakeFormationAwsClientFactory { + @Override + public GlueClient glue() { + GlueClient glue = Mockito.mock(GlueClient.class); + Mockito.when(glue.getTable(Mockito.any(GetTableRequest.class))) + .thenReturn( + GetTableResponse.builder() + .table(Table.builder().isRegisteredWithLakeFormation(true).build()) + .build()); + return glue; + } + } }