AWS: Enable RetryMode for AWS KMS client

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)

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<String, String> properties) {
       this.awsClientProperties = new AwsClientProperties(properties);
       this.s3FileIOProperties = new S3FileIOProperties(properties);
       this.httpClientProperties = new HttpClientProperties(properties);
+      this.kmsClientProperties = new KmsClientProperties(properties);
     }
   }
 

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<String, String> 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 <T extends KmsClientBuilder> void applyRetryConfigurations(T builder) {
+    ClientOverrideConfiguration.Builder configBuilder =
+        null != builder.overrideConfiguration()
+            ? builder.overrideConfiguration().toBuilder()
+            : ClientOverrideConfiguration.builder();
+
+    builder.overrideConfiguration(configBuilder.retryStrategy(kmsRetryMode).build());
+  }
+}

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()))

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<String, String> 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);
+  }
+}