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