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.