Support tagging for multipart upload in AWS and GCP

blob/blob-aws/src/main/java/com/salesforce/multicloudj/blob/aws/AwsBlobStore.java

@@ -370,6 +370,7 @@ protected MultipartUpload doInitiateMultipartUpload(final MultipartUploadRequest
                 .key(createMultipartUploadResponse.key())
                 .id(createMultipartUploadResponse.uploadId())
                 .metadata(request.getMetadata())
+                .tags(request.getTags())
                 .kmsKeyId(request.getKmsKeyId())
                 .build();
     }

blob/blob-aws/src/main/java/com/salesforce/multicloudj/blob/aws/AwsTransformer.java

@@ -300,6 +300,13 @@ public CreateMultipartUploadRequest toCreateMultipartUploadRequest(MultipartUplo
                 .key(request.getKey())
                 .metadata(request.getMetadata());
 
+        if (request.getTags() != null && !request.getTags().isEmpty()) {
+            List<Tag> tags = request.getTags().entrySet().stream()
+                    .map(entry -> Tag.builder().key(entry.getKey()).value(entry.getValue()).build())
+                    .collect(Collectors.toList());
+            builder.tagging(Tagging.builder().tagSet(tags).build());
+        }
+
         if (request.getKmsKeyId() != null && !request.getKmsKeyId().isEmpty()) {
             builder.serverSideEncryption(ServerSideEncryption.AWS_KMS)
                    .ssekmsKeyId(request.getKmsKeyId());

blob/blob-aws/src/main/java/com/salesforce/multicloudj/blob/aws/async/AwsAsyncBlobStore.java

@@ -275,6 +275,7 @@ protected CompletableFuture<MultipartUpload> doInitiateMultipartUpload(Multipart
                         .key(response.key())
                         .id(response.uploadId())
                         .metadata(request.getMetadata())
+                        .tags(request.getTags())
                         .build());
     }
 

blob/blob-aws/src/test/java/com/salesforce/multicloudj/blob/aws/AwsBlobStoreTest.java

@@ -737,6 +737,41 @@ void testDoInitiateMultipartUploadWithKms() {
         assertEquals(kmsKeyId, response.getKmsKeyId());
     }
 
+    @Test
+    void testDoInitiateMultipartUploadWithTags() {
+        CreateMultipartUploadResponse mockResponse = mock(CreateMultipartUploadResponse.class);
+        doReturn("bucket-1").when(mockResponse).bucket();
+        doReturn("object-1").when(mockResponse).key();
+        doReturn("mpu-id").when(mockResponse).uploadId();
+        when(mockS3Client.createMultipartUpload((CreateMultipartUploadRequest) any())).thenReturn(mockResponse);
+        Map<String, String> metadata = Map.of("key-1", "value-1");
+        Map<String, String> tags = Map.of("tag-1", "tag-value-1", "tag-2", "tag-value-2");
+        MultipartUploadRequest request = new MultipartUploadRequest.Builder()
+                .withKey("object-1")
+                .withMetadata(metadata)
+                .withTags(tags)
+                .build();
+
+        MultipartUpload response = aws.initiateMultipartUpload(request);
+
+        // Verify the request is mapped to the SDK with tags
+        ArgumentCaptor<CreateMultipartUploadRequest> requestCaptor = ArgumentCaptor.forClass(CreateMultipartUploadRequest.class);
+        verify(mockS3Client, times(1)).createMultipartUpload(requestCaptor.capture());
+        CreateMultipartUploadRequest actualRequest = requestCaptor.getValue();
+        assertEquals("object-1", actualRequest.key());
+        assertEquals("bucket-1", actualRequest.bucket());
+        assertEquals(metadata, actualRequest.metadata());
+        // Verify tagging header is set (tagging() returns String in AWS SDK)
+        assertNotNull(actualRequest.tagging());
+        assertFalse(actualRequest.tagging().isEmpty());
+
+        // Verify the response is mapped back properly with tags
+        assertEquals("object-1", response.getKey());
+        assertEquals("bucket-1", response.getBucket());
+        assertEquals("mpu-id", response.getId());
+        assertEquals(tags, response.getTags());
+    }
+
     @Test
     void testDoUploadMultipartPart() {
         UploadPartResponse mockResponse = mock(UploadPartResponse.class);

blob/blob-aws/src/test/java/com/salesforce/multicloudj/blob/aws/AwsTransformerTest.java

@@ -391,6 +391,24 @@ void testToCreateMultipartUploadRequest() {
         assertEquals(metadata, request.metadata());
     }
 
+    @Test
+    void testToCreateMultipartUploadRequestWithTags() {
+        Map<String, String> metadata = Map.of("key1", "value1", "key2", "value2");
+        Map<String, String> tags = Map.of("tag1", "value1", "tag2", "value2");
+        MultipartUploadRequest mpuRequest = new MultipartUploadRequest.Builder()
+                .withKey("object-1")
+                .withMetadata(metadata)
+                .withTags(tags)
+                .build();
+        CreateMultipartUploadRequest request = transformer.toCreateMultipartUploadRequest(mpuRequest);
+        assertEquals("object-1", request.key());
+        assertEquals(BUCKET, request.bucket());
+        assertEquals(metadata, request.metadata());
+        // Verify tagging header is set (tagging() returns String in AWS SDK)
+        assertNotNull(request.tagging());
+        assertFalse(request.tagging().isEmpty());
+    }
+
     @Test
     void testToUploadPartRequest() {
         Map<String, String> metadata = Map.of("key1", "value1", "key2", "value2");

blob/blob-aws/src/test/java/com/salesforce/multicloudj/blob/aws/async/AwsAsyncBlobStoreTest.java

@@ -866,6 +866,41 @@ void testDoInitiateMultipartUpload() throws ExecutionException, InterruptedExcep
         assertEquals(metadata, response.getMetadata());
     }
 
+    @Test
+    void testDoInitiateMultipartUploadWithTags() throws ExecutionException, InterruptedException {
+        CreateMultipartUploadResponse mockResponse = mock(CreateMultipartUploadResponse.class);
+        doReturn("bucket-1").when(mockResponse).bucket();
+        doReturn("object-1").when(mockResponse).key();
+        doReturn("mpu-id").when(mockResponse).uploadId();
+        doReturn(future(mockResponse)).when(mockS3Client).createMultipartUpload((CreateMultipartUploadRequest) any());
+        Map<String, String> metadata = Map.of("key-1", "value-1");
+        Map<String, String> tags = Map.of("tag-1", "tag-value-1", "tag-2", "tag-value-2");
+        MultipartUploadRequest request = new MultipartUploadRequest.Builder()
+                .withKey("object-1")
+                .withMetadata(metadata)
+                .withTags(tags)
+                .build();
+
+        MultipartUpload response = aws.initiateMultipartUpload(request).get();
+
+        // Verify the request is mapped to the SDK with tags
+        ArgumentCaptor<CreateMultipartUploadRequest> requestCaptor = ArgumentCaptor.forClass(CreateMultipartUploadRequest.class);
+        verify(mockS3Client, times(1)).createMultipartUpload(requestCaptor.capture());
+        CreateMultipartUploadRequest actualRequest = requestCaptor.getValue();
+        assertEquals("object-1", actualRequest.key());
+        assertEquals("bucket-1", actualRequest.bucket());
+        assertEquals(metadata, actualRequest.metadata());
+        // Verify tagging header is set (tagging() returns String in AWS SDK)
+        assertNotNull(actualRequest.tagging());
+        assertFalse(actualRequest.tagging().isEmpty());
+
+        // Verify the response is mapped back properly with tags
+        assertEquals("object-1", response.getKey());
+        assertEquals("bucket-1", response.getBucket());
+        assertEquals("mpu-id", response.getId());
+        assertEquals(tags, response.getTags());
+    }
+
     @Test
     void testDoUploadMultipartPart() throws ExecutionException, InterruptedException {
         UploadPartResponse mockResponse = mock(UploadPartResponse.class);

blob/blob-client/src/main/java/com/salesforce/multicloudj/blob/driver/MultipartUpload.java

@@ -16,5 +16,6 @@ public class MultipartUpload {
     private final String key;
     private final String id;
     private final Map<String, String> metadata;
+    private final Map<String, String> tags;
     private final String kmsKeyId;
 }

blob/blob-client/src/main/java/com/salesforce/multicloudj/blob/driver/MultipartUploadRequest.java

@@ -2,6 +2,7 @@
 
 import lombok.Getter;
 
+import java.util.Collections;
 import java.util.Map;
 
 import static java.util.Collections.unmodifiableMap;
@@ -14,17 +15,28 @@ public class MultipartUploadRequest {
 
     private final String key;
     private Map<String, String> metadata;
+    private final Map<String, String> tags;
     private final String kmsKeyId;
 
     private MultipartUploadRequest(final Builder builder){
         this.key = builder.key;
         this.metadata = builder.metadata;
+        this.tags = builder.tags;
         this.kmsKeyId = builder.kmsKeyId;
     }
 
+    public Map<String, String> getMetadata() {
+        return metadata == null ? Map.of() : unmodifiableMap(metadata);
+    }
+
+    public Map<String, String> getTags() {
+        return tags == null ? Map.of() : unmodifiableMap(tags);
+    }
+
     public static class Builder {
         private String key;
-        private Map<String, String> metadata;
+        private Map<String, String> metadata = Collections.emptyMap();
+        private Map<String, String> tags = Collections.emptyMap();
         private String kmsKeyId;
 
         public Builder withKey(String key) {
@@ -37,6 +49,11 @@ public Builder withMetadata(final Map<String, String> metadata) {
             return this;
         }
 
+        public Builder withTags(final Map<String, String> tags) {
+            this.tags = unmodifiableMap(tags);
+            return this;
+        }
+
         public Builder withKmsKeyId(String kmsKeyId) {
             this.kmsKeyId = kmsKeyId;
             return this;

blob/blob-client/src/test/java/com/salesforce/multicloudj/blob/client/AbstractBlobStoreIT.java

@@ -1638,6 +1638,7 @@ static class MultipartUploadTestConfig {
         final boolean abortUpload;
         final boolean wantCompletionError;
         final String kmsKeyId;
+        final Map<String, String> tags;
 
         public MultipartUploadTestConfig(String testName,
                           String key,
@@ -1646,7 +1647,7 @@ public MultipartUploadTestConfig(String testName,
                           List<MultipartUploadPartResult> partsToComplete,
                           boolean abortUpload,
                           boolean wantCompletionError) {
-            this(testName, key, metadata, partsToUpload, partsToComplete, abortUpload, wantCompletionError, null);
+            this(testName, key, metadata, partsToUpload, partsToComplete, abortUpload, wantCompletionError, null, null);
         }
 
         public MultipartUploadTestConfig(String testName,
@@ -1657,6 +1658,18 @@ public MultipartUploadTestConfig(String testName,
                           boolean abortUpload,
                           boolean wantCompletionError,
                           String kmsKeyId) {
+            this(testName, key, metadata, partsToUpload, partsToComplete, abortUpload, wantCompletionError, kmsKeyId, null);
+        }
+
+        public MultipartUploadTestConfig(String testName,
+                          String key,
+                          Map<String,String> metadata,
+                          List<MultipartUploadTestPart> partsToUpload,
+                          List<MultipartUploadPartResult> partsToComplete,
+                          boolean abortUpload,
+                          boolean wantCompletionError,
+                          String kmsKeyId,
+                          Map<String, String> tags) {
             this.testName = testName;
             this.key = key;
             this.metadata = metadata;
@@ -1665,6 +1678,7 @@ public MultipartUploadTestConfig(String testName,
             this.abortUpload = abortUpload;
             this.wantCompletionError = wantCompletionError;
             this.kmsKeyId = kmsKeyId;
+            this.tags = tags;
         }
     }
 
@@ -1684,6 +1698,9 @@ private void runMultipartUploadTest(MultipartUploadTestConfig testConfig) throws
             if (testConfig.kmsKeyId != null) {
                 requestBuilder.withKmsKeyId(testConfig.kmsKeyId);
             }
+            if (testConfig.tags != null && !testConfig.tags.isEmpty()) {
+                requestBuilder.withTags(testConfig.tags);
+            }
             MultipartUploadRequest multipartUploadRequest = requestBuilder.build();
             mpu = bucketClient.initiateMultipartUpload(multipartUploadRequest);
 
@@ -1739,7 +1756,24 @@ private void runMultipartUploadTest(MultipartUploadTestConfig testConfig) throws
             }
 
             BlobMetadata blobMetadata = bucketClient.getMetadata(testConfig.key, null);
-            Assertions.assertEquals(testConfig.metadata, blobMetadata.getMetadata(), testConfig.testName + ": Downloaded metadata did not match");
+            Map<String, String> actualMetadata = blobMetadata.getMetadata();
+
+            // For GCP, tags are stored as metadata with "gcp-tag-" prefix, so we need to filter them out
+            // when comparing metadata
+            if (GCP_PROVIDER_ID.equals(harness.getProviderId()) && testConfig.tags != null && !testConfig.tags.isEmpty()) {
+                String tagPrefix = "gcp-tag-";
+                actualMetadata = actualMetadata.entrySet().stream()
+                        .filter(entry -> !entry.getKey().startsWith(tagPrefix))
+                        .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
+            }
+
+            Assertions.assertEquals(testConfig.metadata, actualMetadata, testConfig.testName + ": Downloaded metadata did not match");
+
+            // Verify tags if they were provided
+            if (testConfig.tags != null && !testConfig.tags.isEmpty()) {
+                Map<String, String> tagResults = bucketClient.getTags(testConfig.key);
+                Assertions.assertEquals(testConfig.tags, tagResults, testConfig.testName + ": Tags did not match what was uploaded");
+            }
         }
         finally {
             // Now delete all blobs that were created
@@ -2004,7 +2038,23 @@ public void testMultipartUpload_withKms() throws IOException {
     }
 
     @Test
-    @Disabled
+    public void testMultipartUpload_withTags() throws IOException {
+        // Exclude Ali provider as it doesn't support tags in multipart upload
+        Assumptions.assumeFalse("ali".equals(harness.getProviderId()));
+        Map<String, String> tags = Map.of("tag1", "value1", "tag2", "value2");
+        runMultipartUploadTest(new MultipartUploadTestConfig(
+                "multipart with tags", DEFAULT_MULTIPART_KEY_PREFIX + "withTags",
+                Map.of("key1", "value1"),
+                List.of(
+                        new MultipartUploadTestPart(1, multipartBytes1),
+                        new MultipartUploadTestPart(2, multipartBytes2)),
+                List.of(
+                        new MultipartUploadPartResult(1, true),
+                        new MultipartUploadPartResult(2, true)),
+                false, false, null, tags));
+    }
+
+    @Test
     public void testTagging() throws IOException {
 
         AbstractBlobStore blobStore = harness.createBlobStore(true, true, false);

blob/blob-gcp/src/main/java/com/salesforce/multicloudj/blob/gcp/GcpBlobStore.java

@@ -326,6 +326,7 @@ protected MultipartUpload doInitiateMultipartUpload(MultipartUploadRequest reque
                 .key(request.getKey())
                 .id(uploadId)
                 .metadata(request.getMetadata())
+                .tags(request.getTags())
                 .kmsKeyId(request.getKmsKeyId())
                 .build();
     }

blob/blob-gcp/src/main/java/com/salesforce/multicloudj/blob/gcp/GcpTransformer.java

@@ -268,7 +268,17 @@ public UploadRequest toUploadRequest(MultipartUpload mpu, MultipartPart mpp) {
     }
 
     public BlobInfo toBlobInfo(MultipartUpload mpu) {
-        return toBlobInfo(mpu.getKey(), mpu.getMetadata());
+        Map<String, String> metadata = new HashMap<>();
+        if(mpu.getMetadata() != null) {
+            metadata.putAll(mpu.getMetadata());
+        }
+
+        // Add tags to metadata with TAG_PREFIX
+        if(mpu.getTags() != null && !mpu.getTags().isEmpty()) {
+            mpu.getTags().forEach((tagName, tagValue) -> metadata.put(TAG_PREFIX + tagName, tagValue));
+        }
+
+        return toBlobInfo(mpu.getKey(), metadata);
     }
 
     public List<BlobId> toBlobIdList(Page<Blob> blobs) {

blob/blob-gcp/src/test/java/com/salesforce/multicloudj/blob/gcp/GcpBlobStoreTest.java

@@ -1033,6 +1033,23 @@ void testDoInitiateMultipartUploadWithKms() {
         assertNotNull(mpu.getId());
     }
 
+    @Test
+    void testDoInitiateMultipartUploadWithTags() {
+        Map<String, String> tags = Map.of("tag1", "value1", "tag2", "value2");
+        MultipartUploadRequest request = new MultipartUploadRequest.Builder()
+                .withKey(TEST_KEY)
+                .withMetadata(Map.of("key1","value1","key2","value2"))
+                .withTags(tags)
+                .build();
+
+        MultipartUpload mpu = gcpBlobStore.doInitiateMultipartUpload(request);
+
+        assertEquals(TEST_BUCKET, mpu.getBucket());
+        assertEquals(TEST_KEY, mpu.getKey());
+        assertEquals(tags, mpu.getTags());
+        assertNotNull(mpu.getId());
+    }
+
     @Test
     void testDoUploadMultipartPart() {
         try (MockedStatic<ByteStreams> mockedStatic = Mockito.mockStatic(ByteStreams.class)) {
@@ -1086,6 +1103,34 @@ void testDoCompleteMultipartUpload() {
         assertEquals(TEST_ETAG, response.getEtag());
     }
 
+    @Test
+    void testDoCompleteMultipartUploadWithTags() {
+        Map<String, String> tags = Map.of("tag1", "value1", "tag2", "value2");
+        MultipartUpload mpu = MultipartUpload.builder()
+                .bucket(TEST_BUCKET)
+                .key(TEST_KEY)
+                .id("id")
+                .tags(tags)
+                .build();
+        List<UploadPartResponse> parts = List.of(new UploadPartResponse(1,"etag", 100));
+        Blob mockBlob = mock(Blob.class);
+        doReturn(TEST_ETAG).when(mockBlob).getEtag();
+        doReturn(mockBlob).when(mockStorage).compose(any());
+        when(mockTransformer.toPartName(any(MultipartUpload.class), anyInt())).thenCallRealMethod();
+        BlobInfo blobInfoWithTags = BlobInfo.newBuilder(TEST_BUCKET, TEST_KEY)
+                .setMetadata(Map.of("gcp-tag-tag1", "value1", "gcp-tag-tag2", "value2"))
+                .build();
+        when(mockTransformer.toBlobInfo(mpu)).thenReturn(blobInfoWithTags);
+
+        ArgumentCaptor<Storage.ComposeRequest> composeRequestCaptor = ArgumentCaptor.forClass(Storage.ComposeRequest.class);
+        MultipartUploadResponse response = gcpBlobStore.doCompleteMultipartUpload(mpu, parts);
+
+        verify(mockStorage).compose(composeRequestCaptor.capture());
+        Storage.ComposeRequest composeRequest = composeRequestCaptor.getValue();
+        assertEquals(blobInfoWithTags, composeRequest.getTarget());
+        assertEquals(TEST_ETAG, response.getEtag());
+    }
+
     @Test
     void testDoListMultipartUpload() {
         MultipartUpload mpu = MultipartUpload.builder()