Implement interface changes for s3 plugin to read/write blob with obj… (

#13113) * Implement interface changes for s3 plugin to read/write blob with object metadata --------- Signed-off-by: Sandeep Kumawat <[email protected]> Signed-off-by: Sandeep Kumawat <[email protected]> Co-authored-by: Sandeep Kumawat <[email protected]>
opensearch-project · Apr 12, 2024 · 7345371 · 7345371
1 parent 02e3c56
commit 7345371
Show file tree

Hide file tree

Showing 14 changed files with 220 additions and 42 deletions.
diff --git a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java
@@ -62,6 +62,7 @@
 import software.amazon.awssdk.services.s3.model.UploadPartRequest;
 import software.amazon.awssdk.services.s3.model.UploadPartResponse;
 import software.amazon.awssdk.services.s3.paginators.ListObjectsV2Iterable;
+import software.amazon.awssdk.utils.CollectionUtils;
 
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
@@ -77,6 +78,7 @@
 import org.opensearch.common.blobstore.BlobPath;
 import org.opensearch.common.blobstore.BlobStoreException;
 import org.opensearch.common.blobstore.DeleteResult;
+import org.opensearch.common.blobstore.FetchBlobResult;
 import org.opensearch.common.blobstore.stream.read.ReadContext;
 import org.opensearch.common.blobstore.stream.write.WriteContext;
 import org.opensearch.common.blobstore.stream.write.WritePriority;
@@ -138,6 +140,13 @@ public boolean blobExists(String blobName) {
         }
     }
 
+    @ExperimentalApi
+    @Override
+    public FetchBlobResult readBlobWithMetadata(String blobName) throws IOException {
+        S3RetryingInputStream s3RetryingInputStream = new S3RetryingInputStream(blobStore, buildKey(blobName));
+        return new FetchBlobResult(s3RetryingInputStream, s3RetryingInputStream.getMetadata());
+    }
+
     @Override
     public InputStream readBlob(String blobName) throws IOException {
         return new S3RetryingInputStream(blobStore, buildKey(blobName));
@@ -169,12 +178,27 @@ public long readBlobPreferredLength() {
      */
     @Override
     public void writeBlob(String blobName, InputStream inputStream, long blobSize, boolean failIfAlreadyExists) throws IOException {
+        writeBlobWithMetadata(blobName, inputStream, blobSize, failIfAlreadyExists, null);
+    }
+
+    /**
+     * Write blob with its object metadata.
+     */
+    @ExperimentalApi
+    @Override
+    public void writeBlobWithMetadata(
+        String blobName,
+        InputStream inputStream,
+        long blobSize,
+        boolean failIfAlreadyExists,
+        @Nullable Map<String, String> metadata
+    ) throws IOException {
         assert inputStream.markSupported() : "No mark support on inputStream breaks the S3 SDK's ability to retry requests";
         SocketAccess.doPrivilegedIOException(() -> {
             if (blobSize <= getLargeBlobThresholdInBytes()) {
-                executeSingleUpload(blobStore, buildKey(blobName), inputStream, blobSize);
+                executeSingleUpload(blobStore, buildKey(blobName), inputStream, blobSize, metadata);
             } else {
-                executeMultipartUpload(blobStore, buildKey(blobName), inputStream, blobSize);
+                executeMultipartUpload(blobStore, buildKey(blobName), inputStream, blobSize, metadata);
             }
             return null;
         });
@@ -190,7 +214,8 @@ public void asyncBlobUpload(WriteContext writeContext, ActionListener<Void> comp
             writeContext.getUploadFinalizer(),
             writeContext.doRemoteDataIntegrityCheck(),
             writeContext.getExpectedChecksum(),
-            blobStore.isUploadRetryEnabled()
+            blobStore.isUploadRetryEnabled(),
+            writeContext.getMetadata()
         );
         try {
             if (uploadRequest.getContentLength() > ByteSizeUnit.GB.toBytes(10) && blobStore.isRedirectLargeUploads()) {
@@ -203,7 +228,8 @@ public void asyncBlobUpload(WriteContext writeContext, ActionListener<Void> comp
                         blobStore,
                         uploadRequest.getKey(),
                         inputStream.getInputStream(),
-                        uploadRequest.getContentLength()
+                        uploadRequest.getContentLength(),
+                        uploadRequest.getMetadata()
                     );
                     completionListener.onResponse(null);
                 } catch (Exception ex) {
@@ -542,8 +568,13 @@ private String buildKey(String blobName) {
     /**
      * Uploads a blob using a single upload request
      */
-    void executeSingleUpload(final S3BlobStore blobStore, final String blobName, final InputStream input, final long blobSize)
-        throws IOException {
+    void executeSingleUpload(
+        final S3BlobStore blobStore,
+        final String blobName,
+        final InputStream input,
+        final long blobSize,
+        final Map<String, String> metadata
+    ) throws IOException {
 
         // Extra safety checks
         if (blobSize > MAX_FILE_SIZE.getBytes()) {
@@ -560,6 +591,10 @@ void executeSingleUpload(final S3BlobStore blobStore, final String blobName, fin
             .storageClass(blobStore.getStorageClass())
             .acl(blobStore.getCannedACL())
             .overrideConfiguration(o -> o.addMetricPublisher(blobStore.getStatsMetricPublisher().putObjectMetricPublisher));
+
+        if (CollectionUtils.isNotEmpty(metadata)) {
+            putObjectRequestBuilder = putObjectRequestBuilder.metadata(metadata);
+        }
         if (blobStore.serverSideEncryption()) {
             putObjectRequestBuilder.serverSideEncryption(ServerSideEncryption.AES256);
         }
@@ -583,8 +618,13 @@ void executeSingleUpload(final S3BlobStore blobStore, final String blobName, fin
     /**
      * Uploads a blob using multipart upload requests.
      */
-    void executeMultipartUpload(final S3BlobStore blobStore, final String blobName, final InputStream input, final long blobSize)
-        throws IOException {
+    void executeMultipartUpload(
+        final S3BlobStore blobStore,
+        final String blobName,
+        final InputStream input,
+        final long blobSize,
+        final Map<String, String> metadata
+    ) throws IOException {
 
         ensureMultiPartUploadSize(blobSize);
         final long partSize = blobStore.bufferSizeInBytes();
@@ -609,6 +649,10 @@ void executeMultipartUpload(final S3BlobStore blobStore, final String blobName,
             .acl(blobStore.getCannedACL())
             .overrideConfiguration(o -> o.addMetricPublisher(blobStore.getStatsMetricPublisher().multipartUploadMetricCollector));
 
+        if (CollectionUtils.isNotEmpty(metadata)) {
+            createMultipartUploadRequestBuilder.metadata(metadata);
+        }
+
         if (blobStore.serverSideEncryption()) {
             createMultipartUploadRequestBuilder.serverSideEncryption(ServerSideEncryption.AES256);
         }

diff --git a/...ins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3RetryingInputStream.java b/...ins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3RetryingInputStream.java
@@ -48,6 +48,7 @@
 import java.nio.file.NoSuchFileException;
 import java.util.ArrayList;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.atomic.AtomicBoolean;
 
 /**
@@ -77,6 +78,7 @@ class S3RetryingInputStream extends InputStream {
     private long currentOffset;
     private boolean closed;
     private boolean eof;
+    private Map<String, String> metadata;
 
     S3RetryingInputStream(S3BlobStore blobStore, String blobKey) throws IOException {
         this(blobStore, blobKey, 0, Long.MAX_VALUE - 1);
@@ -122,6 +124,7 @@ private void openStream() throws IOException {
                 getObjectResponseInputStream.response().contentLength()
             );
             this.currentStream = getObjectResponseInputStream;
+            this.metadata = getObjectResponseInputStream.response().metadata();
             this.isStreamAborted.set(false);
         } catch (final SdkException e) {
             if (e instanceof S3Exception) {
@@ -265,4 +268,8 @@ boolean isEof() {
     boolean isAborted() {
         return isStreamAborted.get();
     }
+
+    Map<String, String> getMetadata() {
+        return this.metadata;
+    }
 }
diff --git a/...epository-s3/src/main/java/org/opensearch/repositories/s3/async/AsyncTransferManager.java b/...epository-s3/src/main/java/org/opensearch/repositories/s3/async/AsyncTransferManager.java
@@ -22,6 +22,7 @@
 import software.amazon.awssdk.services.s3.model.DeleteObjectRequest;
 import software.amazon.awssdk.services.s3.model.PutObjectRequest;
 import software.amazon.awssdk.services.s3.model.S3Exception;
+import software.amazon.awssdk.utils.CollectionUtils;
 import software.amazon.awssdk.utils.CompletableFutureUtils;
 
 import org.apache.logging.log4j.LogManager;
@@ -131,6 +132,10 @@ private void uploadInParts(
             .bucket(uploadRequest.getBucket())
             .key(uploadRequest.getKey())
             .overrideConfiguration(o -> o.addMetricPublisher(statsMetricPublisher.multipartUploadMetricCollector));
+
+        if (CollectionUtils.isNotEmpty(uploadRequest.getMetadata())) {
+            createMultipartUploadRequestBuilder.metadata(uploadRequest.getMetadata());
+        }
         if (uploadRequest.doRemoteDataIntegrityCheck()) {
             createMultipartUploadRequestBuilder.checksumAlgorithm(ChecksumAlgorithm.CRC32);
         }
@@ -327,6 +332,10 @@ private void uploadInOneChunk(
             .key(uploadRequest.getKey())
             .contentLength(uploadRequest.getContentLength())
             .overrideConfiguration(o -> o.addMetricPublisher(statsMetricPublisher.putObjectMetricPublisher));
+
+        if (CollectionUtils.isNotEmpty(uploadRequest.getMetadata())) {
+            putObjectRequestBuilder.metadata(uploadRequest.getMetadata());
+        }
         if (uploadRequest.doRemoteDataIntegrityCheck()) {
             putObjectRequestBuilder.checksumAlgorithm(ChecksumAlgorithm.CRC32);
             putObjectRequestBuilder.checksumCRC32(base64StringFromLong(uploadRequest.getExpectedChecksum()));

diff --git a/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/UploadRequest.java b/plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/UploadRequest.java
@@ -9,9 +9,11 @@
 package org.opensearch.repositories.s3.async;
 
 import org.opensearch.common.CheckedConsumer;
+import org.opensearch.common.Nullable;
 import org.opensearch.common.blobstore.stream.write.WritePriority;
 
 import java.io.IOException;
+import java.util.Map;
 
 /**
  * A model encapsulating all details for an upload to S3
@@ -24,8 +26,8 @@ public class UploadRequest {
     private final CheckedConsumer<Boolean, IOException> uploadFinalizer;
     private final boolean doRemoteDataIntegrityCheck;
     private final Long expectedChecksum;
-
     private boolean uploadRetryEnabled;
+    private final Map<String, String> metadata;
 
     /**
      * Construct a new UploadRequest object
@@ -37,6 +39,7 @@ public class UploadRequest {
      * @param uploadFinalizer            An upload finalizer to call once all parts are uploaded
      * @param doRemoteDataIntegrityCheck A boolean to inform vendor plugins whether remote data integrity checks need to be done
      * @param expectedChecksum           Checksum of the file being uploaded for remote data integrity check
+     * @param metadata                   Metadata of the file being uploaded
      */
     public UploadRequest(
         String bucket,
@@ -46,7 +49,8 @@ public UploadRequest(
         CheckedConsumer<Boolean, IOException> uploadFinalizer,
         boolean doRemoteDataIntegrityCheck,
         Long expectedChecksum,
-        boolean uploadRetryEnabled
+        boolean uploadRetryEnabled,
+        @Nullable Map<String, String> metadata
     ) {
         this.bucket = bucket;
         this.key = key;
@@ -56,6 +60,7 @@ public UploadRequest(
         this.doRemoteDataIntegrityCheck = doRemoteDataIntegrityCheck;
         this.expectedChecksum = expectedChecksum;
         this.uploadRetryEnabled = uploadRetryEnabled;
+        this.metadata = metadata;
     }
 
     public String getBucket() {
@@ -89,4 +94,11 @@ public Long getExpectedChecksum() {
     public boolean isUploadRetryEnabled() {
         return uploadRetryEnabled;
     }
+
+    /**
+     * @return metadata of the blob to be uploaded
+     */
+    public Map<String, String> getMetadata() {
+        return metadata;
+    }
 }
diff --git a/...itory-s3/src/test/java/org/opensearch/repositories/s3/S3BlobContainerMockClientTests.java b/...itory-s3/src/test/java/org/opensearch/repositories/s3/S3BlobContainerMockClientTests.java
@@ -57,6 +57,7 @@
 import java.nio.file.Path;
 import java.util.ArrayList;
 import java.util.Arrays;
+import java.util.HashMap;
 import java.util.List;
 import java.util.Locale;
 import java.util.concurrent.CompletableFuture;
@@ -75,6 +76,7 @@
 
 import static org.opensearch.repositories.s3.S3Repository.BULK_DELETE_SIZE;
 import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.anyMap;
 import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.any;
 import static org.mockito.Mockito.doAnswer;
@@ -659,6 +661,7 @@ private void testLargeFilesRedirectedToSlowSyncClient(boolean expectException) t
             .writePriority(WritePriority.HIGH)
             .uploadFinalizer(Assert::assertTrue)
             .doRemoteDataIntegrityCheck(false)
+            .metadata(new HashMap<>())
             .build();
 
         s3BlobContainer.asyncBlobUpload(writeContext, completionListener);
@@ -668,7 +671,13 @@ private void testLargeFilesRedirectedToSlowSyncClient(boolean expectException) t
         } else {
             assertNull(exceptionRef.get());
         }
-        verify(s3BlobContainer, times(1)).executeMultipartUpload(any(S3BlobStore.class), anyString(), any(InputStream.class), anyLong());
+        verify(s3BlobContainer, times(1)).executeMultipartUpload(
+            any(S3BlobStore.class),
+            anyString(),
+            any(InputStream.class),
+            anyLong(),
+            anyMap()
+        );
 
         if (expectException) {
             verify(client, times(1)).abortMultipartUpload(any(AbortMultipartUploadRequest.class));