From 28b3b31361e5aeb2b536555c17141004082e40e9 Mon Sep 17 00:00:00 2001 From: hsiang-c Date: Mon, 28 Oct 2024 15:53:16 +0800 Subject: [PATCH 1/5] Enable RetryMode for AWS KMS client --- .../aws/AssumeRoleAwsClientFactory.java | 7 +++ .../iceberg/aws/AwsClientFactories.java | 5 ++ .../iceberg/aws/kms/KmsClientProperties.java | 63 +++++++++++++++++++ .../LakeFormationAwsClientFactory.java | 1 + .../aws/kms/TestKmsClientProperties.java | 43 +++++++++++++ 5 files changed, 119 insertions(+) create mode 100644 aws/src/main/java/org/apache/iceberg/aws/kms/KmsClientProperties.java create mode 100644 aws/src/test/java/org/apache/iceberg/aws/kms/TestKmsClientProperties.java 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 4c3c305d4bae..7c659efcf4bd 100644 --- a/aws/src/main/java/org/apache/iceberg/aws/AssumeRoleAwsClientFactory.java +++ b/aws/src/main/java/org/apache/iceberg/aws/AssumeRoleAwsClientFactory.java @@ -20,6 +20,7 @@ import java.util.Map; import java.util.UUID; +import org.apache.iceberg.aws.kms.KmsClientProperties; import org.apache.iceberg.aws.s3.S3FileIOProperties; import org.apache.iceberg.relocated.com.google.common.base.Preconditions; import software.amazon.awssdk.awscore.client.builder.AwsClientBuilder; @@ -38,6 +39,7 @@ public class AssumeRoleAwsClientFactory implements AwsClientFactory { private HttpClientProperties httpClientProperties; private S3FileIOProperties s3FileIOProperties; private String roleSessionName; + private KmsClientProperties kmsClientProperties; @Override public S3Client s3() { @@ -64,6 +66,7 @@ public KmsClient kms() { return KmsClient.builder() .applyMutation(this::applyAssumeRoleConfigurations) .applyMutation(httpClientProperties::applyHttpClientConfigurations) + .applyMutation(kmsClientProperties::applyRetryConfigurations) .build(); } @@ -126,6 +129,10 @@ protected S3FileIOProperties s3FileIOProperties() { return s3FileIOProperties; } + protected KmsClientProperties kmsClientProperties() { + return kmsClientProperties; + } + private StsClient sts() { return StsClient.builder() .applyMutation(httpClientProperties::applyHttpClientConfigurations) 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 5974e21209e3..9a5c97384baf 100644 --- a/aws/src/main/java/org/apache/iceberg/aws/AwsClientFactories.java +++ b/aws/src/main/java/org/apache/iceberg/aws/AwsClientFactories.java @@ -20,6 +20,7 @@ import java.net.URI; import java.util.Map; +import org.apache.iceberg.aws.kms.KmsClientProperties; import org.apache.iceberg.aws.s3.S3FileIOProperties; import org.apache.iceberg.common.DynConstructors; import org.apache.iceberg.relocated.com.google.common.base.Strings; @@ -94,12 +95,14 @@ static class DefaultAwsClientFactory implements AwsClientFactory { private AwsClientProperties awsClientProperties; private S3FileIOProperties s3FileIOProperties; private HttpClientProperties httpClientProperties; + private KmsClientProperties kmsClientProperties; DefaultAwsClientFactory() { awsProperties = new AwsProperties(); awsClientProperties = new AwsClientProperties(); s3FileIOProperties = new S3FileIOProperties(); httpClientProperties = new HttpClientProperties(); + kmsClientProperties = new KmsClientProperties(); } @Override @@ -134,6 +137,7 @@ public KmsClient kms() { .applyMutation(awsClientProperties::applyClientRegionConfiguration) .applyMutation(httpClientProperties::applyHttpClientConfigurations) .applyMutation(awsClientProperties::applyClientCredentialConfigurations) + .applyMutation(kmsClientProperties::applyRetryConfigurations) .build(); } @@ -153,6 +157,7 @@ public void initialize(Map properties) { this.awsClientProperties = new AwsClientProperties(properties); this.s3FileIOProperties = new S3FileIOProperties(properties); this.httpClientProperties = new HttpClientProperties(properties); + this.kmsClientProperties = new KmsClientProperties(properties); } } diff --git a/aws/src/main/java/org/apache/iceberg/aws/kms/KmsClientProperties.java b/aws/src/main/java/org/apache/iceberg/aws/kms/KmsClientProperties.java new file mode 100644 index 000000000000..0b152b8613f2 --- /dev/null +++ b/aws/src/main/java/org/apache/iceberg/aws/kms/KmsClientProperties.java @@ -0,0 +1,63 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.iceberg.aws.kms; + +import java.io.Serializable; +import java.util.Map; +import org.apache.iceberg.util.PropertyUtil; +import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration; +import software.amazon.awssdk.core.retry.RetryMode; +import software.amazon.awssdk.services.kms.KmsClientBuilder; + +public class KmsClientProperties implements Serializable { + + /** This property is used to configure {@link RetryMode} for AWS KMS client */ + public static final String KMS_RETRY_MODE = "kms.retry.mode"; + + public static final String KMS_RETRY_MODE_DEFAULT = "LEGACY"; + + private RetryMode kmsRetryMode; + + public KmsClientProperties() { + this.kmsRetryMode = RetryMode.valueOf(KMS_RETRY_MODE_DEFAULT); + } + + public KmsClientProperties(Map properties) { + this.kmsRetryMode = + RetryMode.valueOf( + PropertyUtil.propertyAsString(properties, KMS_RETRY_MODE, KMS_RETRY_MODE_DEFAULT)); + } + + public RetryMode kmsRetryMode() { + return kmsRetryMode; + } + + public void setKmsRetryMode(RetryMode kmsRetryMode) { + this.kmsRetryMode = kmsRetryMode; + } + + public void applyRetryConfigurations(T builder) { + ClientOverrideConfiguration.Builder configBuilder = + null != builder.overrideConfiguration() + ? builder.overrideConfiguration().toBuilder() + : ClientOverrideConfiguration.builder(); + + builder.overrideConfiguration(configBuilder.retryStrategy(kmsRetryMode).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 5d37470066d2..7230d3213823 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 @@ -96,6 +96,7 @@ public KmsClient kms() { if (isTableRegisteredWithLakeFormation()) { return KmsClient.builder() .applyMutation(httpClientProperties()::applyHttpClientConfigurations) + .applyMutation(kmsClientProperties()::applyRetryConfigurations) .credentialsProvider( new LakeFormationCredentialsProvider(lakeFormation(), buildTableArn())) .region(Region.of(region())) diff --git a/aws/src/test/java/org/apache/iceberg/aws/kms/TestKmsClientProperties.java b/aws/src/test/java/org/apache/iceberg/aws/kms/TestKmsClientProperties.java new file mode 100644 index 000000000000..93e33af5b9c3 --- /dev/null +++ b/aws/src/test/java/org/apache/iceberg/aws/kms/TestKmsClientProperties.java @@ -0,0 +1,43 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.iceberg.aws.kms; + +import static org.assertj.core.api.Assertions.assertThat; + +import java.util.Map; +import org.apache.iceberg.relocated.com.google.common.collect.Maps; +import org.junit.jupiter.api.Test; +import software.amazon.awssdk.core.retry.RetryMode; +import software.amazon.awssdk.services.kms.KmsClient; +import software.amazon.awssdk.services.kms.KmsClientBuilder; + +public class TestKmsClientProperties { + @Test + public void testApplyRetryConfiguration() { + Map properties = Maps.newHashMap(); + properties.put(KmsClientProperties.KMS_RETRY_MODE, "ADAPTIVE_V2"); + KmsClientProperties kmsClientProperties = new KmsClientProperties(properties); + + KmsClientBuilder builder = KmsClient.builder(); + kmsClientProperties.applyRetryConfigurations(builder); + + RetryMode retryPolicy = builder.overrideConfiguration().retryMode().get(); + assertThat(retryPolicy).as("retry mode was not set").isEqualTo(RetryMode.ADAPTIVE_V2); + } +} From 13321b7f52b9e6f060585ea5ea6bf7a761ef7fb8 Mon Sep 17 00:00:00 2001 From: hsiang-c Date: Wed, 30 Oct 2024 08:33:34 +0800 Subject: [PATCH 2/5] Comment on RetryMode --- .../org/apache/iceberg/aws/kms/KmsClientProperties.java | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/aws/src/main/java/org/apache/iceberg/aws/kms/KmsClientProperties.java b/aws/src/main/java/org/apache/iceberg/aws/kms/KmsClientProperties.java index 0b152b8613f2..38f04997f7cc 100644 --- a/aws/src/main/java/org/apache/iceberg/aws/kms/KmsClientProperties.java +++ b/aws/src/main/java/org/apache/iceberg/aws/kms/KmsClientProperties.java @@ -52,6 +52,14 @@ public void setKmsRetryMode(RetryMode kmsRetryMode) { this.kmsRetryMode = kmsRetryMode; } + /** + * Configure RetryMode + * for a KMS client. + * + *

For more details, see + * https://sdk.amazonaws.com/java/api/latest/software/amazon/awssdk/core/client/config/ClientOverrideConfiguration.Builder.html#retryStrategy(software.amazon.awssdk.core.retry.RetryMode) + */ public void applyRetryConfigurations(T builder) { ClientOverrideConfiguration.Builder configBuilder = null != builder.overrideConfiguration() From 943f6b040b435bb8f809d71a044579b0624e5aa8 Mon Sep 17 00:00:00 2001 From: hsiang-c Date: Fri, 8 Nov 2024 10:27:43 +0800 Subject: [PATCH 3/5] Set retry mode to ADAPTIVE_V2 --- .../aws/AssumeRoleAwsClientFactory.java | 1 + .../iceberg/aws/AwsClientFactories.java | 2 +- .../iceberg/aws/kms/KmsClientProperties.java | 32 ++----------------- .../aws/kms/TestKmsClientProperties.java | 6 +--- 4 files changed, 6 insertions(+), 35 deletions(-) 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 7c659efcf4bd..3d0fe57a266b 100644 --- a/aws/src/main/java/org/apache/iceberg/aws/AssumeRoleAwsClientFactory.java +++ b/aws/src/main/java/org/apache/iceberg/aws/AssumeRoleAwsClientFactory.java @@ -84,6 +84,7 @@ public void initialize(Map properties) { this.awsProperties = new AwsProperties(properties); this.s3FileIOProperties = new S3FileIOProperties(properties); this.httpClientProperties = new HttpClientProperties(properties); + this.kmsClientProperties = new KmsClientProperties(); this.roleSessionName = genSessionName(); Preconditions.checkNotNull( awsProperties.clientAssumeRoleArn(), 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 9a5c97384baf..5abc822a0223 100644 --- a/aws/src/main/java/org/apache/iceberg/aws/AwsClientFactories.java +++ b/aws/src/main/java/org/apache/iceberg/aws/AwsClientFactories.java @@ -157,7 +157,7 @@ public void initialize(Map properties) { this.awsClientProperties = new AwsClientProperties(properties); this.s3FileIOProperties = new S3FileIOProperties(properties); this.httpClientProperties = new HttpClientProperties(properties); - this.kmsClientProperties = new KmsClientProperties(properties); + this.kmsClientProperties = new KmsClientProperties(); } } diff --git a/aws/src/main/java/org/apache/iceberg/aws/kms/KmsClientProperties.java b/aws/src/main/java/org/apache/iceberg/aws/kms/KmsClientProperties.java index 38f04997f7cc..da0967067ed7 100644 --- a/aws/src/main/java/org/apache/iceberg/aws/kms/KmsClientProperties.java +++ b/aws/src/main/java/org/apache/iceberg/aws/kms/KmsClientProperties.java @@ -19,41 +19,15 @@ package org.apache.iceberg.aws.kms; import java.io.Serializable; -import java.util.Map; -import org.apache.iceberg.util.PropertyUtil; import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration; import software.amazon.awssdk.core.retry.RetryMode; import software.amazon.awssdk.services.kms.KmsClientBuilder; public class KmsClientProperties implements Serializable { - - /** This property is used to configure {@link RetryMode} for AWS KMS client */ - public static final String KMS_RETRY_MODE = "kms.retry.mode"; - - public static final String KMS_RETRY_MODE_DEFAULT = "LEGACY"; - - private RetryMode kmsRetryMode; - - public KmsClientProperties() { - this.kmsRetryMode = RetryMode.valueOf(KMS_RETRY_MODE_DEFAULT); - } - - public KmsClientProperties(Map properties) { - this.kmsRetryMode = - RetryMode.valueOf( - PropertyUtil.propertyAsString(properties, KMS_RETRY_MODE, KMS_RETRY_MODE_DEFAULT)); - } - - public RetryMode kmsRetryMode() { - return kmsRetryMode; - } - - public void setKmsRetryMode(RetryMode kmsRetryMode) { - this.kmsRetryMode = kmsRetryMode; - } + public KmsClientProperties() {} /** - * Configure RetryMode * for a KMS client. * @@ -66,6 +40,6 @@ public void applyRetryConfigurations(T builder) { ? builder.overrideConfiguration().toBuilder() : ClientOverrideConfiguration.builder(); - builder.overrideConfiguration(configBuilder.retryStrategy(kmsRetryMode).build()); + builder.overrideConfiguration(configBuilder.retryStrategy(RetryMode.ADAPTIVE_V2).build()); } } diff --git a/aws/src/test/java/org/apache/iceberg/aws/kms/TestKmsClientProperties.java b/aws/src/test/java/org/apache/iceberg/aws/kms/TestKmsClientProperties.java index 93e33af5b9c3..68cafba5f4d2 100644 --- a/aws/src/test/java/org/apache/iceberg/aws/kms/TestKmsClientProperties.java +++ b/aws/src/test/java/org/apache/iceberg/aws/kms/TestKmsClientProperties.java @@ -20,8 +20,6 @@ import static org.assertj.core.api.Assertions.assertThat; -import java.util.Map; -import org.apache.iceberg.relocated.com.google.common.collect.Maps; import org.junit.jupiter.api.Test; import software.amazon.awssdk.core.retry.RetryMode; import software.amazon.awssdk.services.kms.KmsClient; @@ -30,9 +28,7 @@ public class TestKmsClientProperties { @Test public void testApplyRetryConfiguration() { - Map properties = Maps.newHashMap(); - properties.put(KmsClientProperties.KMS_RETRY_MODE, "ADAPTIVE_V2"); - KmsClientProperties kmsClientProperties = new KmsClientProperties(properties); + KmsClientProperties kmsClientProperties = new KmsClientProperties(); KmsClientBuilder builder = KmsClient.builder(); kmsClientProperties.applyRetryConfigurations(builder); From cfa192d6db0ee29a8f17f7363445b9d62355b34b Mon Sep 17 00:00:00 2001 From: hsiang-c Date: Mon, 2 Dec 2024 10:22:23 +0800 Subject: [PATCH 4/5] Move retry config to the generic AwsClientProperties class --- .../aws/AssumeRoleAwsClientFactory.java | 11 +++-- .../iceberg/aws/AwsClientFactories.java | 6 +-- .../iceberg/aws/AwsClientProperties.java | 19 ++++++++ .../iceberg/aws/kms/KmsClientProperties.java | 45 ------------------- .../LakeFormationAwsClientFactory.java | 2 +- .../aws/kms/TestKmsClientProperties.java | 7 +-- 6 files changed, 30 insertions(+), 60 deletions(-) delete mode 100644 aws/src/main/java/org/apache/iceberg/aws/kms/KmsClientProperties.java 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 3d0fe57a266b..4149d795d30d 100644 --- a/aws/src/main/java/org/apache/iceberg/aws/AssumeRoleAwsClientFactory.java +++ b/aws/src/main/java/org/apache/iceberg/aws/AssumeRoleAwsClientFactory.java @@ -20,7 +20,6 @@ import java.util.Map; import java.util.UUID; -import org.apache.iceberg.aws.kms.KmsClientProperties; import org.apache.iceberg.aws.s3.S3FileIOProperties; import org.apache.iceberg.relocated.com.google.common.base.Preconditions; import software.amazon.awssdk.awscore.client.builder.AwsClientBuilder; @@ -39,7 +38,7 @@ public class AssumeRoleAwsClientFactory implements AwsClientFactory { private HttpClientProperties httpClientProperties; private S3FileIOProperties s3FileIOProperties; private String roleSessionName; - private KmsClientProperties kmsClientProperties; + private AwsClientProperties awsClientProperties; @Override public S3Client s3() { @@ -66,7 +65,7 @@ public KmsClient kms() { return KmsClient.builder() .applyMutation(this::applyAssumeRoleConfigurations) .applyMutation(httpClientProperties::applyHttpClientConfigurations) - .applyMutation(kmsClientProperties::applyRetryConfigurations) + .applyMutation(awsClientProperties::applyRetryConfigurations) .build(); } @@ -84,7 +83,7 @@ public void initialize(Map properties) { this.awsProperties = new AwsProperties(properties); this.s3FileIOProperties = new S3FileIOProperties(properties); this.httpClientProperties = new HttpClientProperties(properties); - this.kmsClientProperties = new KmsClientProperties(); + this.awsClientProperties = new AwsClientProperties(properties); this.roleSessionName = genSessionName(); Preconditions.checkNotNull( awsProperties.clientAssumeRoleArn(), @@ -130,8 +129,8 @@ protected S3FileIOProperties s3FileIOProperties() { return s3FileIOProperties; } - protected KmsClientProperties kmsClientProperties() { - return kmsClientProperties; + protected AwsClientProperties awsClientProperties() { + return awsClientProperties; } private StsClient sts() { 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 5abc822a0223..7554b5629be4 100644 --- a/aws/src/main/java/org/apache/iceberg/aws/AwsClientFactories.java +++ b/aws/src/main/java/org/apache/iceberg/aws/AwsClientFactories.java @@ -20,7 +20,6 @@ import java.net.URI; import java.util.Map; -import org.apache.iceberg.aws.kms.KmsClientProperties; import org.apache.iceberg.aws.s3.S3FileIOProperties; import org.apache.iceberg.common.DynConstructors; import org.apache.iceberg.relocated.com.google.common.base.Strings; @@ -95,14 +94,12 @@ static class DefaultAwsClientFactory implements AwsClientFactory { private AwsClientProperties awsClientProperties; private S3FileIOProperties s3FileIOProperties; private HttpClientProperties httpClientProperties; - private KmsClientProperties kmsClientProperties; DefaultAwsClientFactory() { awsProperties = new AwsProperties(); awsClientProperties = new AwsClientProperties(); s3FileIOProperties = new S3FileIOProperties(); httpClientProperties = new HttpClientProperties(); - kmsClientProperties = new KmsClientProperties(); } @Override @@ -137,7 +134,7 @@ public KmsClient kms() { .applyMutation(awsClientProperties::applyClientRegionConfiguration) .applyMutation(httpClientProperties::applyHttpClientConfigurations) .applyMutation(awsClientProperties::applyClientCredentialConfigurations) - .applyMutation(kmsClientProperties::applyRetryConfigurations) + .applyMutation(awsClientProperties::applyRetryConfigurations) .build(); } @@ -157,7 +154,6 @@ public void initialize(Map properties) { this.awsClientProperties = new AwsClientProperties(properties); this.s3FileIOProperties = new S3FileIOProperties(properties); this.httpClientProperties = new HttpClientProperties(properties); - this.kmsClientProperties = new KmsClientProperties(); } } diff --git a/aws/src/main/java/org/apache/iceberg/aws/AwsClientProperties.java b/aws/src/main/java/org/apache/iceberg/aws/AwsClientProperties.java index 0c91f8685ae9..add09d0fab4a 100644 --- a/aws/src/main/java/org/apache/iceberg/aws/AwsClientProperties.java +++ b/aws/src/main/java/org/apache/iceberg/aws/AwsClientProperties.java @@ -31,6 +31,8 @@ import software.amazon.awssdk.auth.credentials.DefaultCredentialsProvider; import software.amazon.awssdk.auth.credentials.StaticCredentialsProvider; import software.amazon.awssdk.awscore.client.builder.AwsClientBuilder; +import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration; +import software.amazon.awssdk.core.retry.RetryMode; import software.amazon.awssdk.regions.Region; public class AwsClientProperties implements Serializable { @@ -154,6 +156,23 @@ public AwsCredentialsProvider credentialsProvider( return DefaultCredentialsProvider.builder().build(); } + /** + * Set ADAPTIVE_V2 RetryMode + * for an AWS client builder. + * + *

For more details, see + * https://sdk.amazonaws.com/java/api/latest/software/amazon/awssdk/core/client/config/ClientOverrideConfiguration.Builder.html#retryStrategy(software.amazon.awssdk.core.retry.RetryMode) + */ + public void applyRetryConfigurations(T builder) { + ClientOverrideConfiguration.Builder configBuilder = + null != builder.overrideConfiguration() + ? builder.overrideConfiguration().toBuilder() + : ClientOverrideConfiguration.builder(); + + builder.overrideConfiguration(configBuilder.retryStrategy(RetryMode.ADAPTIVE_V2).build()); + } + private AwsCredentialsProvider credentialsProvider(String credentialsProviderClass) { Class providerClass; try { diff --git a/aws/src/main/java/org/apache/iceberg/aws/kms/KmsClientProperties.java b/aws/src/main/java/org/apache/iceberg/aws/kms/KmsClientProperties.java deleted file mode 100644 index da0967067ed7..000000000000 --- a/aws/src/main/java/org/apache/iceberg/aws/kms/KmsClientProperties.java +++ /dev/null @@ -1,45 +0,0 @@ -/* - * Licensed to the Apache Software Foundation (ASF) under one - * or more contributor license agreements. See the NOTICE file - * distributed with this work for additional information - * regarding copyright ownership. The ASF licenses this file - * to you under the Apache License, Version 2.0 (the - * "License"); you may not use this file except in compliance - * with the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, - * software distributed under the License is distributed on an - * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY - * KIND, either express or implied. See the License for the - * specific language governing permissions and limitations - * under the License. - */ -package org.apache.iceberg.aws.kms; - -import java.io.Serializable; -import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration; -import software.amazon.awssdk.core.retry.RetryMode; -import software.amazon.awssdk.services.kms.KmsClientBuilder; - -public class KmsClientProperties implements Serializable { - public KmsClientProperties() {} - - /** - * Set ADAPTIVE_V2 RetryMode - * for a KMS client. - * - *

For more details, see - * https://sdk.amazonaws.com/java/api/latest/software/amazon/awssdk/core/client/config/ClientOverrideConfiguration.Builder.html#retryStrategy(software.amazon.awssdk.core.retry.RetryMode) - */ - public void applyRetryConfigurations(T builder) { - ClientOverrideConfiguration.Builder configBuilder = - null != builder.overrideConfiguration() - ? builder.overrideConfiguration().toBuilder() - : ClientOverrideConfiguration.builder(); - - builder.overrideConfiguration(configBuilder.retryStrategy(RetryMode.ADAPTIVE_V2).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 7230d3213823..aaab2c2bc769 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 @@ -96,7 +96,7 @@ public KmsClient kms() { if (isTableRegisteredWithLakeFormation()) { return KmsClient.builder() .applyMutation(httpClientProperties()::applyHttpClientConfigurations) - .applyMutation(kmsClientProperties()::applyRetryConfigurations) + .applyMutation(awsClientProperties()::applyRetryConfigurations) .credentialsProvider( new LakeFormationCredentialsProvider(lakeFormation(), buildTableArn())) .region(Region.of(region())) diff --git a/aws/src/test/java/org/apache/iceberg/aws/kms/TestKmsClientProperties.java b/aws/src/test/java/org/apache/iceberg/aws/kms/TestKmsClientProperties.java index 68cafba5f4d2..0dc416b683b8 100644 --- a/aws/src/test/java/org/apache/iceberg/aws/kms/TestKmsClientProperties.java +++ b/aws/src/test/java/org/apache/iceberg/aws/kms/TestKmsClientProperties.java @@ -20,6 +20,7 @@ import static org.assertj.core.api.Assertions.assertThat; +import org.apache.iceberg.aws.AwsClientProperties; import org.junit.jupiter.api.Test; import software.amazon.awssdk.core.retry.RetryMode; import software.amazon.awssdk.services.kms.KmsClient; @@ -28,12 +29,12 @@ public class TestKmsClientProperties { @Test public void testApplyRetryConfiguration() { - KmsClientProperties kmsClientProperties = new KmsClientProperties(); + AwsClientProperties awsClientProperties = new AwsClientProperties(); KmsClientBuilder builder = KmsClient.builder(); - kmsClientProperties.applyRetryConfigurations(builder); - + awsClientProperties.applyRetryConfigurations(builder); RetryMode retryPolicy = builder.overrideConfiguration().retryMode().get(); + assertThat(retryPolicy).as("retry mode was not set").isEqualTo(RetryMode.ADAPTIVE_V2); } } From fd2f8e89cfb9ce3a6dd402d3ec3a7e1d87730cec Mon Sep 17 00:00:00 2001 From: hsiang-c Date: Tue, 10 Dec 2024 16:41:06 +0800 Subject: [PATCH 5/5] Address Amogh's review --- .../org/apache/iceberg/aws/AwsClientProperties.java | 11 +++++++---- .../iceberg/aws/kms/TestKmsClientProperties.java | 2 +- 2 files changed, 8 insertions(+), 5 deletions(-) diff --git a/aws/src/main/java/org/apache/iceberg/aws/AwsClientProperties.java b/aws/src/main/java/org/apache/iceberg/aws/AwsClientProperties.java index add09d0fab4a..6a275737c9f0 100644 --- a/aws/src/main/java/org/apache/iceberg/aws/AwsClientProperties.java +++ b/aws/src/main/java/org/apache/iceberg/aws/AwsClientProperties.java @@ -157,12 +157,15 @@ public AwsCredentialsProvider credentialsProvider( } /** - * Set ADAPTIVE_V2 RetryMode - * for an AWS client builder. + * to ADAPTIVE_V2 for AWS clients * - *

For more details, see - * https://sdk.amazonaws.com/java/api/latest/software/amazon/awssdk/core/client/config/ClientOverrideConfiguration.Builder.html#retryStrategy(software.amazon.awssdk.core.retry.RetryMode) + *

Sample usage: + * + * 

+   *   KmsClient.builder().applyMutation(awsClientProperties::applyRetryConfigurations)
+   *
*/ public void applyRetryConfigurations(T builder) { ClientOverrideConfiguration.Builder configBuilder = diff --git a/aws/src/test/java/org/apache/iceberg/aws/kms/TestKmsClientProperties.java b/aws/src/test/java/org/apache/iceberg/aws/kms/TestKmsClientProperties.java index 0dc416b683b8..9d1fb8db51eb 100644 --- a/aws/src/test/java/org/apache/iceberg/aws/kms/TestKmsClientProperties.java +++ b/aws/src/test/java/org/apache/iceberg/aws/kms/TestKmsClientProperties.java @@ -35,6 +35,6 @@ public void testApplyRetryConfiguration() { awsClientProperties.applyRetryConfigurations(builder); RetryMode retryPolicy = builder.overrideConfiguration().retryMode().get(); - assertThat(retryPolicy).as("retry mode was not set").isEqualTo(RetryMode.ADAPTIVE_V2); + assertThat(retryPolicy).as("retry mode should be ADAPTIVE_V2").isEqualTo(RetryMode.ADAPTIVE_V2); } }