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 fef845a0a005..e023f7c3d331 100644
--- a/aws/src/main/java/org/apache/iceberg/aws/AwsProperties.java
+++ b/aws/src/main/java/org/apache/iceberg/aws/AwsProperties.java
@@ -22,6 +22,7 @@
import java.io.Serializable;
import java.util.Map;
import org.apache.iceberg.aws.dynamodb.DynamoDbCatalog;
+import org.apache.iceberg.aws.s3.S3FileIO;
import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
import org.apache.iceberg.util.PropertyUtil;
import software.amazon.awssdk.services.s3.model.ObjectCannedACL;
@@ -215,6 +216,15 @@ public class AwsProperties implements Serializable {
*/
public static final String CLIENT_ASSUME_ROLE_REGION = "client.assume-role.region";
+ /**
+ * Used by {@link S3FileIO} to tag objects when writing. To set, we can pass a catalog property.
+ *
+ * For more details, see https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-tagging.html
+ *
+ * Example in Spark: --conf spark.sql.catalog.my_catalog.s3.write.tags.my_key=my_val
+ */
+ public static final String S3_WRITE_TAGS_PREFIX = "s3.write.tags.";
+
/**
* @deprecated will be removed at 0.15.0, please use {@link #S3_CHECKSUM_ENABLED_DEFAULT} instead
*/
diff --git a/aws/src/main/java/org/apache/iceberg/aws/s3/BaseS3File.java b/aws/src/main/java/org/apache/iceberg/aws/s3/BaseS3File.java
index d5381733a8ee..2fc043fcf5c7 100644
--- a/aws/src/main/java/org/apache/iceberg/aws/s3/BaseS3File.java
+++ b/aws/src/main/java/org/apache/iceberg/aws/s3/BaseS3File.java
@@ -19,13 +19,16 @@
package org.apache.iceberg.aws.s3;
+import java.util.Set;
import org.apache.iceberg.aws.AwsProperties;
import org.apache.iceberg.metrics.MetricsContext;
+import org.apache.iceberg.relocated.com.google.common.collect.Sets;
import software.amazon.awssdk.http.HttpStatusCode;
import software.amazon.awssdk.services.s3.S3Client;
import software.amazon.awssdk.services.s3.model.HeadObjectRequest;
import software.amazon.awssdk.services.s3.model.HeadObjectResponse;
import software.amazon.awssdk.services.s3.model.S3Exception;
+import software.amazon.awssdk.services.s3.model.Tag;
abstract class BaseS3File {
private final S3Client client;
@@ -33,12 +36,23 @@ abstract class BaseS3File {
private final AwsProperties awsProperties;
private HeadObjectResponse metadata;
private final MetricsContext metrics;
+ private final Set writeTags;
BaseS3File(S3Client client, S3URI uri, AwsProperties awsProperties, MetricsContext metrics) {
this.client = client;
this.uri = uri;
this.awsProperties = awsProperties;
this.metrics = metrics;
+ this.writeTags = Sets.newHashSet();
+ }
+
+ BaseS3File(S3Client client, S3URI uri, AwsProperties awsProperties, MetricsContext metrics,
+ Set writeTags) {
+ this.client = client;
+ this.uri = uri;
+ this.awsProperties = awsProperties;
+ this.metrics = metrics;
+ this.writeTags = writeTags;
}
public String location() {
@@ -61,6 +75,10 @@ protected MetricsContext metrics() {
return metrics;
}
+ public Set writeTags() {
+ return writeTags;
+ }
+
/**
* Note: this may be stale if file was deleted since metadata is cached for size/existence checks.
*
diff --git a/aws/src/main/java/org/apache/iceberg/aws/s3/S3FileIO.java b/aws/src/main/java/org/apache/iceberg/aws/s3/S3FileIO.java
index 469a1607746c..7f908fa26e32 100644
--- a/aws/src/main/java/org/apache/iceberg/aws/s3/S3FileIO.java
+++ b/aws/src/main/java/org/apache/iceberg/aws/s3/S3FileIO.java
@@ -20,7 +20,9 @@
package org.apache.iceberg.aws.s3;
import java.util.Map;
+import java.util.Set;
import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.stream.Collectors;
import org.apache.iceberg.aws.AwsClientFactories;
import org.apache.iceberg.aws.AwsProperties;
import org.apache.iceberg.common.DynConstructors;
@@ -28,11 +30,13 @@
import org.apache.iceberg.io.InputFile;
import org.apache.iceberg.io.OutputFile;
import org.apache.iceberg.metrics.MetricsContext;
+import org.apache.iceberg.util.PropertyUtil;
import org.apache.iceberg.util.SerializableSupplier;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import software.amazon.awssdk.services.s3.S3Client;
import software.amazon.awssdk.services.s3.model.DeleteObjectRequest;
+import software.amazon.awssdk.services.s3.model.Tag;
/**
* FileIO implementation backed by S3.
@@ -50,6 +54,7 @@ public class S3FileIO implements FileIO {
private transient S3Client client;
private MetricsContext metrics = MetricsContext.nullMetrics();
private final AtomicBoolean isResourceClosed = new AtomicBoolean(false);
+ private Set writeTags;
/**
* No-arg constructor to load the FileIO dynamically.
@@ -88,7 +93,7 @@ public InputFile newInputFile(String path) {
@Override
public OutputFile newOutputFile(String path) {
- return S3OutputFile.fromLocation(path, client(), awsProperties, metrics);
+ return S3OutputFile.fromLocation(path, client(), awsProperties, metrics, writeTags);
}
@Override
@@ -110,6 +115,7 @@ private S3Client client() {
@Override
public void initialize(Map properties) {
this.awsProperties = new AwsProperties(properties);
+ this.writeTags = toTags(properties);
// Do not override s3 client if it was provided
if (s3 == null) {
@@ -137,4 +143,11 @@ public void close() {
}
}
}
+
+ private Set toTags(Map properties) {
+ return PropertyUtil.propertiesWithPrefix(properties, AwsProperties.S3_WRITE_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/s3/S3OutputFile.java b/aws/src/main/java/org/apache/iceberg/aws/s3/S3OutputFile.java
index 48b2ee4f0cae..5c70dbabe8fe 100644
--- a/aws/src/main/java/org/apache/iceberg/aws/s3/S3OutputFile.java
+++ b/aws/src/main/java/org/apache/iceberg/aws/s3/S3OutputFile.java
@@ -21,22 +21,30 @@
import java.io.IOException;
import java.io.UncheckedIOException;
+import java.util.Set;
import org.apache.iceberg.aws.AwsProperties;
import org.apache.iceberg.exceptions.AlreadyExistsException;
import org.apache.iceberg.io.InputFile;
import org.apache.iceberg.io.OutputFile;
import org.apache.iceberg.io.PositionOutputStream;
import org.apache.iceberg.metrics.MetricsContext;
+import org.apache.iceberg.relocated.com.google.common.collect.Sets;
import software.amazon.awssdk.services.s3.S3Client;
+import software.amazon.awssdk.services.s3.model.Tag;
public class S3OutputFile extends BaseS3File implements OutputFile {
public static S3OutputFile fromLocation(String location, S3Client client, AwsProperties awsProperties,
MetricsContext metrics) {
- return new S3OutputFile(client, new S3URI(location), awsProperties, metrics);
+ return new S3OutputFile(client, new S3URI(location), awsProperties, metrics, Sets.newHashSet());
}
- S3OutputFile(S3Client client, S3URI uri, AwsProperties awsProperties, MetricsContext metrics) {
- super(client, uri, awsProperties, metrics);
+ public static S3OutputFile fromLocation(String location, S3Client client, AwsProperties awsProperties,
+ MetricsContext metrics, Set writeTags) {
+ return new S3OutputFile(client, new S3URI(location), awsProperties, metrics, writeTags);
+ }
+
+ S3OutputFile(S3Client client, S3URI uri, AwsProperties awsProperties, MetricsContext metrics, Set writeTags) {
+ super(client, uri, awsProperties, metrics, writeTags);
}
/**
@@ -57,7 +65,7 @@ public PositionOutputStream create() {
@Override
public PositionOutputStream createOrOverwrite() {
try {
- return new S3OutputStream(client(), uri(), awsProperties(), metrics());
+ return new S3OutputStream(client(), uri(), awsProperties(), metrics(), writeTags());
} catch (IOException e) {
throw new UncheckedIOException("Failed to create output stream for location: " + uri(), e);
}
diff --git a/aws/src/main/java/org/apache/iceberg/aws/s3/S3OutputStream.java b/aws/src/main/java/org/apache/iceberg/aws/s3/S3OutputStream.java
index 4d5605b147c6..bf84c4230e5b 100644
--- a/aws/src/main/java/org/apache/iceberg/aws/s3/S3OutputStream.java
+++ b/aws/src/main/java/org/apache/iceberg/aws/s3/S3OutputStream.java
@@ -37,6 +37,7 @@
import java.util.Comparator;
import java.util.List;
import java.util.Map;
+import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
@@ -67,6 +68,8 @@
import software.amazon.awssdk.services.s3.model.CompletedPart;
import software.amazon.awssdk.services.s3.model.CreateMultipartUploadRequest;
import software.amazon.awssdk.services.s3.model.PutObjectRequest;
+import software.amazon.awssdk.services.s3.model.Tag;
+import software.amazon.awssdk.services.s3.model.Tagging;
import software.amazon.awssdk.services.s3.model.UploadPartRequest;
import software.amazon.awssdk.services.s3.model.UploadPartResponse;
import software.amazon.awssdk.utils.BinaryUtils;
@@ -81,6 +84,7 @@ class S3OutputStream extends PositionOutputStream {
private final S3Client s3;
private final S3URI location;
private final AwsProperties awsProperties;
+ private final Set writeTags;
private CountingOutputStream stream;
private final List stagingFiles = Lists.newArrayList();
@@ -101,7 +105,8 @@ class S3OutputStream extends PositionOutputStream {
private boolean closed = false;
@SuppressWarnings("StaticAssignmentInConstructor")
- S3OutputStream(S3Client s3, S3URI location, AwsProperties awsProperties, MetricsContext metrics) throws IOException {
+ S3OutputStream(S3Client s3, S3URI location, AwsProperties awsProperties, MetricsContext metrics, Set writeTags)
+ throws IOException {
if (executorService == null) {
synchronized (S3OutputStream.class) {
if (executorService == null) {
@@ -119,6 +124,7 @@ class S3OutputStream extends PositionOutputStream {
this.s3 = s3;
this.location = location;
this.awsProperties = awsProperties;
+ this.writeTags = writeTags;
this.createStack = Thread.currentThread().getStackTrace();
@@ -252,7 +258,12 @@ public void close() throws IOException {
private void initializeMultiPartUpload() {
CreateMultipartUploadRequest.Builder requestBuilder = CreateMultipartUploadRequest.builder()
- .bucket(location.bucket()).key(location.key());
+ .bucket(location.bucket())
+ .key(location.key());
+ if (!writeTags.isEmpty()) {
+ requestBuilder.tagging(Tagging.builder().tagSet(writeTags).build());
+ }
+
S3RequestUtil.configureEncryption(awsProperties, requestBuilder);
S3RequestUtil.configurePermission(awsProperties, requestBuilder);
@@ -367,6 +378,10 @@ private void completeUploads() {
.bucket(location.bucket())
.key(location.key());
+ if (!writeTags.isEmpty()) {
+ requestBuilder.tagging(Tagging.builder().tagSet(writeTags).build());
+ }
+
if (isChecksumEnabled) {
requestBuilder.contentMD5(BinaryUtils.toBase64(completeMessageDigest.digest()));
}
diff --git a/aws/src/test/java/org/apache/iceberg/aws/glue/TestGlueCatalog.java b/aws/src/test/java/org/apache/iceberg/aws/glue/TestGlueCatalog.java
index b607ca6b3b23..6268c9fc5da2 100644
--- a/aws/src/test/java/org/apache/iceberg/aws/glue/TestGlueCatalog.java
+++ b/aws/src/test/java/org/apache/iceberg/aws/glue/TestGlueCatalog.java
@@ -94,8 +94,8 @@ public void testConstructorEmptyWarehousePath() {
@Test
public void testConstructorWarehousePathWithEndSlash() {
GlueCatalog catalogWithSlash = new GlueCatalog();
- catalogWithSlash.initialize(
- CATALOG_NAME, WAREHOUSE_PATH + "/", new AwsProperties(), glue, LockManagers.defaultLockManager(), null);
+ catalogWithSlash.initialize(CATALOG_NAME, WAREHOUSE_PATH + "/", new AwsProperties(), glue,
+ LockManagers.defaultLockManager(), null);
Mockito.doReturn(GetDatabaseResponse.builder()
.database(Database.builder().name("db").build()).build())
.when(glue).getDatabase(Mockito.any(GetDatabaseRequest.class));
diff --git a/aws/src/test/java/org/apache/iceberg/aws/s3/TestS3FileIO.java b/aws/src/test/java/org/apache/iceberg/aws/s3/TestS3FileIO.java
index 2cca2f6c2d12..a08e3d9e7fda 100644
--- a/aws/src/test/java/org/apache/iceberg/aws/s3/TestS3FileIO.java
+++ b/aws/src/test/java/org/apache/iceberg/aws/s3/TestS3FileIO.java
@@ -23,11 +23,14 @@
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
+import java.util.Map;
import java.util.Random;
import org.apache.commons.io.IOUtils;
import org.apache.commons.lang3.SerializationUtils;
import org.apache.iceberg.io.InputFile;
import org.apache.iceberg.io.OutputFile;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableSet;
import org.apache.iceberg.util.SerializableSupplier;
import org.junit.Before;
import org.junit.ClassRule;
@@ -36,6 +39,7 @@
import software.amazon.awssdk.regions.Region;
import software.amazon.awssdk.services.s3.S3Client;
import software.amazon.awssdk.services.s3.model.CreateBucketRequest;
+import software.amazon.awssdk.services.s3.model.Tag;
import static org.junit.Assert.assertArrayEquals;
import static org.junit.Assert.assertEquals;
@@ -49,10 +53,13 @@ public class TestS3FileIO {
private final Random random = new Random(1);
private S3FileIO s3FileIO;
+ private final Map properties = ImmutableMap.of(
+ "s3.write.tags.tagKey1", "TagValue1");
@Before
public void before() {
s3FileIO = new S3FileIO(s3);
+ s3FileIO.initialize(properties);
s3.get().createBucket(CreateBucketRequest.builder().bucket("bucket").build());
}
@@ -94,4 +101,30 @@ public void testSerializeClient() {
assertEquals("s3", post.get().serviceName());
}
+
+ @Test
+ public void testWriteTags() throws IOException {
+ String location = "s3://bucket/path/to/file.txt";
+ byte[] expected = new byte[1024 * 1024];
+ random.nextBytes(expected);
+
+ InputFile in = s3FileIO.newInputFile(location);
+ assertFalse(in.exists());
+
+ OutputFile out = s3FileIO.newOutputFile(location);
+ try (OutputStream os = out.createOrOverwrite()) {
+ IOUtils.write(expected, os);
+ }
+
+ assertTrue(in.exists());
+
+ // Assert for writeTags
+ assertTrue(((S3InputFile) in).writeTags().isEmpty());
+ assertEquals(((S3OutputFile) out).writeTags().size(), properties.size());
+ assertEquals(((S3OutputFile) out).writeTags(), ImmutableSet.of(
+ Tag.builder().key("tagKey1").value("TagValue1").build()));
+
+ s3FileIO.deleteFile(in);
+ assertFalse(s3FileIO.newInputFile(location).exists());
+ }
}
diff --git a/aws/src/test/java/org/apache/iceberg/aws/s3/TestS3OutputStream.java b/aws/src/test/java/org/apache/iceberg/aws/s3/TestS3OutputStream.java
index 6fda6732f24c..921b0f59e11e 100644
--- a/aws/src/test/java/org/apache/iceberg/aws/s3/TestS3OutputStream.java
+++ b/aws/src/test/java/org/apache/iceberg/aws/s3/TestS3OutputStream.java
@@ -30,11 +30,13 @@
import java.util.Comparator;
import java.util.List;
import java.util.Random;
+import java.util.Set;
import java.util.UUID;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import org.apache.iceberg.aws.AwsProperties;
import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableSet;
import org.junit.After;
import org.junit.Before;
import org.junit.ClassRule;
@@ -54,6 +56,7 @@
import software.amazon.awssdk.services.s3.model.GetObjectRequest;
import software.amazon.awssdk.services.s3.model.GetObjectResponse;
import software.amazon.awssdk.services.s3.model.PutObjectRequest;
+import software.amazon.awssdk.services.s3.model.Tag;
import software.amazon.awssdk.services.s3.model.UploadPartRequest;
import software.amazon.awssdk.utils.BinaryUtils;
@@ -85,6 +88,9 @@ public class TestS3OutputStream {
private final Random random = new Random(1);
private final Path tmpDir = Files.createTempDirectory("s3fileio-test-");
private final String newTmpDirectory = "/tmp/newStagingDirectory";
+ private final Set tags = ImmutableSet.of(
+ Tag.builder().key("abc").value("123").build(),
+ Tag.builder().key("def").value("789").build());
private final AwsProperties properties = new AwsProperties(ImmutableMap.of(
AwsProperties.S3FILEIO_MULTIPART_SIZE, Integer.toString(5 * 1024 * 1024),
@@ -116,7 +122,8 @@ public void testWrite() {
public void testAbortAfterFailedPartUpload() {
doThrow(new RuntimeException()).when(s3mock).uploadPart((UploadPartRequest) any(), (RequestBody) any());
- try (S3OutputStream stream = new S3OutputStream(s3mock, randomURI(), properties, nullMetrics())) {
+ try (S3OutputStream stream = new S3OutputStream(
+ s3mock, randomURI(), properties, nullMetrics(), tags)) {
stream.write(randomData(10 * 1024 * 1024));
} catch (Exception e) {
verify(s3mock, atLeastOnce()).abortMultipartUpload((AbortMultipartUploadRequest) any());
@@ -127,7 +134,8 @@ public void testAbortAfterFailedPartUpload() {
public void testAbortMultipart() {
doThrow(new RuntimeException()).when(s3mock).completeMultipartUpload((CompleteMultipartUploadRequest) any());
- try (S3OutputStream stream = new S3OutputStream(s3mock, randomURI(), properties, nullMetrics())) {
+ try (S3OutputStream stream = new S3OutputStream(
+ s3mock, randomURI(), properties, nullMetrics(), tags)) {
stream.write(randomData(10 * 1024 * 1024));
} catch (Exception e) {
verify(s3mock).abortMultipartUpload((AbortMultipartUploadRequest) any());
@@ -136,7 +144,7 @@ public void testAbortMultipart() {
@Test
public void testMultipleClose() throws IOException {
- S3OutputStream stream = new S3OutputStream(s3, randomURI(), properties, nullMetrics());
+ S3OutputStream stream = new S3OutputStream(s3, randomURI(), properties, nullMetrics(), tags);
stream.close();
stream.close();
}
@@ -145,7 +153,8 @@ public void testMultipleClose() throws IOException {
public void testStagingDirectoryCreation() throws IOException {
AwsProperties newStagingDirectoryAwsProperties = new AwsProperties(ImmutableMap.of(
AwsProperties.S3FILEIO_STAGING_DIRECTORY, newTmpDirectory));
- S3OutputStream stream = new S3OutputStream(s3, randomURI(), newStagingDirectoryAwsProperties, nullMetrics());
+ S3OutputStream stream = new S3OutputStream(
+ s3, randomURI(), newStagingDirectoryAwsProperties, nullMetrics(), tags);
stream.close();
}
@@ -166,6 +175,7 @@ private void writeTest() {
verify(s3mock, times(1)).putObject(putObjectRequestArgumentCaptor.capture(),
(RequestBody) any());
checkPutObjectRequestContent(data, putObjectRequestArgumentCaptor);
+ checkTags(putObjectRequestArgumentCaptor);
reset(s3mock);
// Test file larger than part size but less than multipart threshold
@@ -175,6 +185,7 @@ private void writeTest() {
verify(s3mock, times(1)).putObject(putObjectRequestArgumentCaptor.capture(),
(RequestBody) any());
checkPutObjectRequestContent(data, putObjectRequestArgumentCaptor);
+ checkTags(putObjectRequestArgumentCaptor);
reset(s3mock);
// Test file large enough to trigger multipart upload
@@ -224,6 +235,20 @@ private void checkPutObjectRequestContent(
}
}
+ private void checkTags(ArgumentCaptor putObjectRequestArgumentCaptor) {
+ if (properties.isS3ChecksumEnabled()) {
+ List putObjectRequests = putObjectRequestArgumentCaptor.getAllValues();
+ String tagging = putObjectRequests.get(0).tagging();
+ assertEquals(getTags(tags), tagging);
+ }
+ }
+
+ private String getTags(Set objectTags) {
+ return objectTags.stream()
+ .map(e -> e.key() + "=" + e.value())
+ .collect(Collectors.joining("&"));
+ }
+
private String getDigest(byte[] data, int offset, int length) {
try {
MessageDigest md5 = MessageDigest.getInstance("MD5");
@@ -236,7 +261,7 @@ private String getDigest(byte[] data, int offset, int length) {
}
private void writeAndVerify(S3Client client, S3URI uri, byte [] data, boolean arrayWrite) {
- try (S3OutputStream stream = new S3OutputStream(client, uri, properties, nullMetrics())) {
+ try (S3OutputStream stream = new S3OutputStream(client, uri, properties, nullMetrics(), tags)) {
if (arrayWrite) {
stream.write(data);
assertEquals(data.length, stream.getPos());
diff --git a/core/src/main/java/org/apache/iceberg/util/PropertyUtil.java b/core/src/main/java/org/apache/iceberg/util/PropertyUtil.java
index ae66b67592ec..601f88fc0745 100644
--- a/core/src/main/java/org/apache/iceberg/util/PropertyUtil.java
+++ b/core/src/main/java/org/apache/iceberg/util/PropertyUtil.java
@@ -20,6 +20,9 @@
package org.apache.iceberg.util;
import java.util.Map;
+import java.util.stream.Collectors;
+import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
public class PropertyUtil {
@@ -70,4 +73,26 @@ public static String propertyAsString(Map properties,
}
return defaultValue;
}
+
+ /**
+ * Returns subset of provided map with keys matching the provided prefix. Matching is case-sensitive and the matching
+ * prefix is removed from the keys in returned map.
+ *
+ * @param properties input map
+ * @param prefix prefix to choose keys from input map
+ * @return subset of input map with keys starting with provided prefix and prefix trimmed out
+ */
+ public static Map propertiesWithPrefix(
+ Map properties, String prefix) {
+ if (properties == null || properties.isEmpty()) {
+ return ImmutableMap.of();
+ }
+
+ Preconditions.checkArgument(prefix != null, "prefix can't be null.");
+
+ return properties.entrySet().stream()
+ .filter(e -> e.getKey().startsWith(prefix))
+ .collect(Collectors.toMap(
+ e -> e.getKey().replaceFirst(prefix, ""), Map.Entry::getValue));
+ }
}