AWS: Add S3 Tagging

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,14 @@ 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.
+   * Example in Spark: --conf spark.sql.catalog.my_catalog.s3.write.tags.my_key=my_val
+   * <p>
+   * For more details, see https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-tagging.html
+   */
+  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
    */

aws/src/main/java/org/apache/iceberg/aws/s3/BaseS3File.java

@@ -19,26 +19,40 @@
 
 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;
   private final S3URI uri;
   private final AwsProperties awsProperties;
   private HeadObjectResponse metadata;
   private final MetricsContext metrics;
+  private final Set<Tag> 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<Tag> 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<Tag> writeTags() {
+    return writeTags;
+  }
+
   /**
    * Note: this may be stale if file was deleted since metadata is cached for size/existence checks.
    *

aws/src/main/java/org/apache/iceberg/aws/s3/S3FileIO.java

@@ -20,19 +20,23 @@
 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;
 import org.apache.iceberg.io.FileIO;
 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<Tag> 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,8 @@ private S3Client client() {
   @Override
   public void initialize(Map<String, String> properties) {
     this.awsProperties = new AwsProperties(properties);
+    this.writeTags = toTags(
+        PropertyUtil.propertiesWithPrefix(properties, AwsProperties.S3_WRITE_TAGS_PREFIX));
 
     // Do not override s3 client if it was provided
     if (s3 == null) {
@@ -137,4 +144,10 @@ public void close() {
       }
     }
   }
+
+  private Set<Tag> toTags(Map<String, String> tagKeyToTagValue) {
+    return tagKeyToTagValue.entrySet().stream()
+        .map(e -> Tag.builder().key(e.getKey()).value(e.getValue()).build())
+        .collect(Collectors.toSet());
+  }
 }

aws/src/main/java/org/apache/iceberg/aws/s3/S3OutputFile.java

@@ -21,22 +21,24 @@
 
 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 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);
+      MetricsContext metrics, Set<Tag> writeTags) {
+    return new S3OutputFile(client, new S3URI(location), awsProperties, metrics, writeTags);
   }
 
-  S3OutputFile(S3Client client, S3URI uri, AwsProperties awsProperties, MetricsContext metrics) {
-    super(client, uri, awsProperties, metrics);
+  S3OutputFile(S3Client client, S3URI uri, AwsProperties awsProperties, MetricsContext metrics, Set<Tag> writeTags) {
+    super(client, uri, awsProperties, metrics, writeTags);
   }
 
   /**
@@ -57,7 +59,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);
     }

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<Tag> writeTags;
 
   private CountingOutputStream stream;
   private final List<FileAndDigest> 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<Tag> 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()));
       }

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

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