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 59a4d8d3ac38..dd955d579761 100644 --- a/aws/src/main/java/org/apache/iceberg/aws/AssumeRoleAwsClientFactory.java +++ b/aws/src/main/java/org/apache/iceberg/aws/AssumeRoleAwsClientFactory.java @@ -81,6 +81,7 @@ public GlueClient glue() { return GlueClient.builder() .applyMutation(this::applyAssumeRoleConfigurations) .applyMutation(httpClientProperties::applyHttpClientConfigurations) + .applyMutation(awsClientProperties::applyRetryConfigurations) .build(); } @@ -99,6 +100,7 @@ public DynamoDbClient dynamo() { .applyMutation(this::applyAssumeRoleConfigurations) .applyMutation(httpClientProperties::applyHttpClientConfigurations) .applyMutation(awsProperties::applyDynamoDbEndpointConfigurations) + .applyMutation(awsClientProperties::applyRetryConfigurations) .build(); } diff --git a/aws/src/main/java/org/apache/iceberg/aws/AwsClientFactories.java b/aws/src/main/java/org/apache/iceberg/aws/AwsClientFactories.java index 0ee4bf26e68e..e6eb8d28d01a 100644 --- a/aws/src/main/java/org/apache/iceberg/aws/AwsClientFactories.java +++ b/aws/src/main/java/org/apache/iceberg/aws/AwsClientFactories.java @@ -148,6 +148,7 @@ public GlueClient glue() { .applyMutation(httpClientProperties::applyHttpClientConfigurations) .applyMutation(awsProperties::applyGlueEndpointConfigurations) .applyMutation(awsClientProperties::applyClientCredentialConfigurations) + .applyMutation(awsClientProperties::applyRetryConfigurations) .build(); } @@ -169,6 +170,7 @@ public DynamoDbClient dynamo() { .applyMutation(httpClientProperties::applyHttpClientConfigurations) .applyMutation(awsClientProperties::applyClientCredentialConfigurations) .applyMutation(awsProperties::applyDynamoDbEndpointConfigurations) + .applyMutation(awsClientProperties::applyRetryConfigurations) .build(); } 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 fe95f9364673..94ef60e05d68 100644 --- a/aws/src/test/java/org/apache/iceberg/aws/TestAwsClientFactories.java +++ b/aws/src/test/java/org/apache/iceberg/aws/TestAwsClientFactories.java @@ -23,6 +23,7 @@ import java.io.IOException; import java.util.Map; +import java.util.Optional; import org.apache.iceberg.TestHelpers; import org.apache.iceberg.aws.lakeformation.LakeFormationAwsClientFactory; import org.apache.iceberg.aws.s3.S3FileIOProperties; @@ -38,7 +39,10 @@ import software.amazon.awssdk.auth.credentials.AwsBasicCredentials; import software.amazon.awssdk.auth.credentials.AwsCredentials; import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.awscore.AwsClient; import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.retries.api.RetryStrategy; +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.GetTablesRequest; @@ -245,6 +249,42 @@ public void testWithClassDoesNotImplementCredentialsProvider() { testProviderAndAssertThrownBy(providerClassName, containsMessage); } + @Test + public void testGlueClientSetsAdaptiveRetryPolicy() { + AwsClientFactory factory = + getAwsClientFactoryByCredentialsProvider(DummyValidProvider.class.getName()); + GlueClient glueClient = factory.glue(); + assertAwsClientSetsAdaptiveRetryPolicy(glueClient); + } + + @Test + public void testKmsClientSetsAdaptiveRetryPolicy() { + AwsClientFactory factory = + getAwsClientFactoryByCredentialsProvider(DummyValidProvider.class.getName()); + KmsClient kmsClient = factory.kms(); + assertAwsClientSetsAdaptiveRetryPolicy(kmsClient); + } + + @Test + public void testDynamoClientSetsAdaptiveRetryPolicy() { + AwsClientFactory factory = + getAwsClientFactoryByCredentialsProvider(DummyValidProvider.class.getName()); + DynamoDbClient dynamoClient = factory.dynamo(); + assertAwsClientSetsAdaptiveRetryPolicy(dynamoClient); + } + + /** + * 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 + * object. + */ + private void assertAwsClientSetsAdaptiveRetryPolicy(AwsClient client) { + Optional retryStrategy = + client.serviceClientConfiguration().overrideConfiguration().retryStrategy(); + assertThat(retryStrategy).isPresent(); + assertThat(retryStrategy.get()).isInstanceOf(DefaultAdaptiveRetryStrategy.class); + } + private void testProviderAndAssertThrownBy(String providerClassName, String containsMessage) { AwsClientFactory defaultAwsClientFactory = getAwsClientFactoryByCredentialsProvider(providerClassName);