This is an automated email from the ASF dual-hosted git repository.

jackye pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/iceberg.git


The following commit(s) were added to refs/heads/master by this push:
     new a86f08c  AWS: Add TagSession support in AssumeRoleAwsClientFactory 
(#4358)
a86f08c is described below

commit a86f08ca1cbafd3e29decd9f60041af7b088d964
Author: Xiaoxuan <[email protected]>
AuthorDate: Thu Mar 17 20:35:20 2022 -0700

    AWS: Add TagSession support in AssumeRoleAwsClientFactory (#4358)
---
 .../iceberg/aws/TestAssumeRoleAwsClientFactory.java |  5 ++++-
 .../iceberg/aws/AssumeRoleAwsClientFactory.java     | 21 +++++++++++++++++----
 .../java/org/apache/iceberg/aws/AwsProperties.java  |  6 ++++++
 3 files changed, 27 insertions(+), 5 deletions(-)

diff --git 
a/aws/src/integration/java/org/apache/iceberg/aws/TestAssumeRoleAwsClientFactory.java
 
b/aws/src/integration/java/org/apache/iceberg/aws/TestAssumeRoleAwsClientFactory.java
index 2502776..30b0056 100644
--- 
a/aws/src/integration/java/org/apache/iceberg/aws/TestAssumeRoleAwsClientFactory.java
+++ 
b/aws/src/integration/java/org/apache/iceberg/aws/TestAssumeRoleAwsClientFactory.java
@@ -68,13 +68,16 @@ public class TestAssumeRoleAwsClientFactory {
             "\"Effect\":\"Allow\"," +
             "\"Principal\":{" +
             "\"AWS\":\"arn:aws:iam::" + AwsIntegTestUtil.testAccountId() + 
":root\"}," +
-            "\"Action\": \"sts:AssumeRole\"}]}")
+            "\"Action\": [\"sts:AssumeRole\"," +
+            "\"sts:TagSession\"]}]}")
         .maxSessionDuration(3600)
         .build());
     assumeRoleProperties = Maps.newHashMap();
     assumeRoleProperties.put(AwsProperties.CLIENT_FACTORY, 
AssumeRoleAwsClientFactory.class.getName());
     assumeRoleProperties.put(AwsProperties.CLIENT_ASSUME_ROLE_REGION, 
"us-east-1");
     assumeRoleProperties.put(AwsProperties.CLIENT_ASSUME_ROLE_ARN, 
response.role().arn());
+    assumeRoleProperties.put(AwsProperties.CLIENT_ASSUME_ROLE_TAGS_PREFIX + 
"key1", "value1");
+    assumeRoleProperties.put(AwsProperties.CLIENT_ASSUME_ROLE_TAGS_PREFIX + 
"key2", "value2");
     policyName = UUID.randomUUID().toString();
   }
 
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 1d3a53d..1433761 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,9 @@
 package org.apache.iceberg.aws;
 
 import java.util.Map;
+import java.util.Set;
 import java.util.UUID;
+import java.util.stream.Collectors;
 import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
 import org.apache.iceberg.util.PropertyUtil;
 import software.amazon.awssdk.awscore.client.builder.AwsClientBuilder;
@@ -34,11 +36,13 @@ import software.amazon.awssdk.services.s3.S3Client;
 import software.amazon.awssdk.services.sts.StsClient;
 import 
software.amazon.awssdk.services.sts.auth.StsAssumeRoleCredentialsProvider;
 import software.amazon.awssdk.services.sts.model.AssumeRoleRequest;
+import software.amazon.awssdk.services.sts.model.Tag;
 
 public class AssumeRoleAwsClientFactory implements AwsClientFactory {
 
   private String roleArn;
   private String externalId;
+  private Set<Tag> tags;
   private int timeout;
   private String region;
   private String s3Endpoint;
@@ -68,17 +72,18 @@ public class AssumeRoleAwsClientFactory implements 
AwsClientFactory {
 
   @Override
   public void initialize(Map<String, String> properties) {
-    roleArn = properties.get(AwsProperties.CLIENT_ASSUME_ROLE_ARN);
+    this.roleArn = properties.get(AwsProperties.CLIENT_ASSUME_ROLE_ARN);
     Preconditions.checkNotNull(roleArn,
         "Cannot initialize AssumeRoleClientConfigFactory with null role ARN");
-    timeout = PropertyUtil.propertyAsInt(properties,
+    this.timeout = PropertyUtil.propertyAsInt(properties,
         AwsProperties.CLIENT_ASSUME_ROLE_TIMEOUT_SEC, 
AwsProperties.CLIENT_ASSUME_ROLE_TIMEOUT_SEC_DEFAULT);
-    externalId = properties.get(AwsProperties.CLIENT_ASSUME_ROLE_EXTERNAL_ID);
+    this.externalId = 
properties.get(AwsProperties.CLIENT_ASSUME_ROLE_EXTERNAL_ID);
 
-    region = properties.get(AwsProperties.CLIENT_ASSUME_ROLE_REGION);
+    this.region = properties.get(AwsProperties.CLIENT_ASSUME_ROLE_REGION);
     Preconditions.checkNotNull(region, "Cannot initialize 
AssumeRoleClientConfigFactory with null region");
 
     this.s3Endpoint = properties.get(AwsProperties.S3FILEIO_ENDPOINT);
+    this.tags = toTags(properties);
   }
 
   private <T extends AwsClientBuilder & AwsSyncClientBuilder> T configure(T 
clientBuilder) {
@@ -87,6 +92,7 @@ public class AssumeRoleAwsClientFactory implements 
AwsClientFactory {
         .roleSessionName(genSessionName())
         .durationSeconds(timeout)
         .externalId(externalId)
+        .tags(tags)
         .build();
 
     clientBuilder.credentialsProvider(
@@ -104,4 +110,11 @@ public class AssumeRoleAwsClientFactory implements 
AwsClientFactory {
   private String genSessionName() {
     return String.format("iceberg-aws-%s", UUID.randomUUID());
   }
+
+  private static Set<Tag> toTags(Map<String, String> properties) {
+    return PropertyUtil.propertiesWithPrefix(properties, 
AwsProperties.CLIENT_ASSUME_ROLE_TAGS_PREFIX)
+        .entrySet().stream()
+        .map(e -> Tag.builder().key(e.getKey()).value(e.getValue()).build())
+        .collect(Collectors.toSet());
+  }
 }
diff --git a/aws/src/main/java/org/apache/iceberg/aws/AwsProperties.java 
b/aws/src/main/java/org/apache/iceberg/aws/AwsProperties.java
index 590ccda..520de94 100644
--- a/aws/src/main/java/org/apache/iceberg/aws/AwsProperties.java
+++ b/aws/src/main/java/org/apache/iceberg/aws/AwsProperties.java
@@ -215,6 +215,12 @@ public class AwsProperties implements Serializable {
   public static final String CLIENT_ASSUME_ROLE_ARN = "client.assume-role.arn";
 
   /**
+   * Used by {@link AssumeRoleAwsClientFactory} to pass a list of sessions.
+   * Each session tag consists of a key name and an associated value.
+   */
+  public static final String CLIENT_ASSUME_ROLE_TAGS_PREFIX = 
"client.assume-role.tags.";
+
+  /**
    * Used by {@link AssumeRoleAwsClientFactory}.
    * The timeout of the assume role session in seconds, default to 1 hour.
    * At the end of the timeout, a new set of role session credentials will be 
fetched through a STS client.

Reply via email to