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

libs/common/src/main/java/org/opensearch/common/io/InputStreamContainer.java

@@ -11,6 +11,7 @@
 import org.opensearch.common.annotation.ExperimentalApi;
 
 import java.io.InputStream;
+import java.util.Map;
 
 /**
  * Model composed of an input stream and the total content length of the stream
@@ -23,17 +24,20 @@
     private final InputStream inputStream;
     private final long contentLength;
     private final long offset;
+    private final Map<String, String> metadata;
 
     /**
      * Construct a new stream object
      *
      * @param inputStream   The input stream that is to be encapsulated
      * @param contentLength The total content length that is to be read from the stream
+     * @param metadata      The metadata of the blob. This will be same for each part download.
      */
-    public InputStreamContainer(InputStream inputStream, long contentLength, long offset) {
+    public InputStreamContainer(InputStream inputStream, long contentLength, long offset, Map<String, String> metadata) {
         this.inputStream = inputStream;
         this.contentLength = contentLength;
         this.offset = offset;
+        this.metadata = metadata;
     }
 
     /**
@@ -56,4 +60,11 @@
     public long getOffset() {
         return offset;
     }
+
+    /**
+     * @return metadata of the source content.
+     */
+    public Map<String, String> getMetadata() {
+        return metadata;
+    }
 }

plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/S3BlobContainer.java

@@ -73,6 +73,7 @@
 import org.opensearch.common.annotation.ExperimentalApi;
 import org.opensearch.common.blobstore.AsyncMultiStreamBlobContainer;
 import org.opensearch.common.blobstore.BlobContainer;
+import org.opensearch.common.blobstore.BlobDownloadResponse;
 import org.opensearch.common.blobstore.BlobMetadata;
 import org.opensearch.common.blobstore.BlobPath;
 import org.opensearch.common.blobstore.BlobStoreException;
@@ -138,6 +139,13 @@
         }
     }
 
+    @ExperimentalApi
+    @Override
+    public BlobDownloadResponse readBlobWithMetadata(String blobName) throws IOException {
+        S3RetryingInputStream s3RetryingInputStream = new S3RetryingInputStream(blobStore, buildKey(blobName));
+        return new BlobDownloadResponse(s3RetryingInputStream, s3RetryingInputStream.getMetadata());
+    }
+
     @Override
     public InputStream readBlob(String blobName) throws IOException {
         return new S3RetryingInputStream(blobStore, buildKey(blobName));
@@ -169,12 +177,27 @@
      */
     @Override
     public void writeBlob(String blobName, InputStream inputStream, long blobSize, boolean failIfAlreadyExists) throws IOException {
+        writeBlobWithMetadata(blobName, inputStream, null, blobSize, failIfAlreadyExists);
+    }
+
+    /**
+     * Write blob with its object metadata.
+     */
+    @ExperimentalApi
+    @Override
+    public void writeBlobWithMetadata(
+        String blobName,
+        InputStream inputStream,
+        @Nullable Map<String, String> metadata,
+        long blobSize,
+        boolean failIfAlreadyExists
+    ) 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 +213,8 @@
             writeContext.getUploadFinalizer(),
             writeContext.doRemoteDataIntegrityCheck(),
             writeContext.getExpectedChecksum(),
-            blobStore.isUploadRetryEnabled()
+            blobStore.isUploadRetryEnabled(),
+            writeContext.getMetadata()
         );
         try {
             if (uploadRequest.getContentLength() > ByteSizeUnit.GB.toBytes(10) && blobStore.isRedirectLargeUploads()) {
@@ -203,7 +227,8 @@
                         blobStore,
                         uploadRequest.getKey(),
                         inputStream.getInputStream(),
-                        uploadRequest.getContentLength()
+                        uploadRequest.getContentLength(),
+                        uploadRequest.getMetadata()
                     );
                     completionListener.onResponse(null);
                 } catch (Exception ex) {
@@ -309,6 +334,18 @@
         writeBlob(blobName, inputStream, blobSize, failIfAlreadyExists);
     }
 
+    @ExperimentalApi
+    @Override
+    public void writeBlobAtomicWithMetadata(
+        String blobName,
+        InputStream inputStream,
+        Map<String, String> metadata,
+        long blobSize,
+        boolean failIfAlreadyExists
+    ) throws IOException {
+        writeBlobWithMetadata(blobName, inputStream, metadata, blobSize, failIfAlreadyExists);
+    }
+
     @Override
     public DeleteResult delete() throws IOException {
         final AtomicLong deletedBlobs = new AtomicLong();
@@ -542,8 +579,13 @@
     /**
      * 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 +602,10 @@
             .storageClass(blobStore.getStorageClass())
             .acl(blobStore.getCannedACL())
             .overrideConfiguration(o -> o.addMetricPublisher(blobStore.getStatsMetricPublisher().putObjectMetricPublisher));
+
+        if (metadata != null && !metadata.isEmpty()) {
+            putObjectRequestBuilder = putObjectRequestBuilder.metadata(metadata);
+        }
         if (blobStore.serverSideEncryption()) {
             putObjectRequestBuilder.serverSideEncryption(ServerSideEncryption.AES256);
         }
@@ -583,8 +629,13 @@
     /**
      * 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 +660,10 @@
             .acl(blobStore.getCannedACL())
             .overrideConfiguration(o -> o.addMetricPublisher(blobStore.getStatsMetricPublisher().multipartUploadMetricCollector));
 
+        if (metadata != null && !metadata.isEmpty()) {
+            createMultipartUploadRequestBuilder.metadata(metadata);
+        }
+
         if (blobStore.serverSideEncryption()) {
             createMultipartUploadRequestBuilder.serverSideEncryption(ServerSideEncryption.AES256);
         }
@@ -767,11 +822,12 @@
         final GetObjectResponse getObjectResponse = streamResponse.response();
         final String contentRange = getObjectResponse.contentRange();
         final Long contentLength = getObjectResponse.contentLength();
+        final Map<String, String> metadata = getObjectResponse.metadata();
         if ((isMultipartObject && contentRange == null) || contentLength == null) {
             throw SdkException.builder().message("Failed to fetch required metadata for blob part").build();
         }
         final long offset = isMultipartObject ? HttpRangeUtils.getStartOffsetFromRangeHeader(getObjectResponse.contentRange()) : 0L;
-        return new InputStreamContainer(streamResponse, getObjectResponse.contentLength(), offset);
+        return new InputStreamContainer(streamResponse, getObjectResponse.contentLength(), offset, metadata);
     }
 
     /**

plugins/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 @@
     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 @@
                 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 isAborted() {
         return isStreamAborted.get();
     }
+
+    Map<String, String> getMetadata() {
+        return this.metadata;
+    }
 }

plugins/repository-s3/src/main/java/org/opensearch/repositories/s3/async/AsyncTransferManager.java

@@ -131,6 +131,10 @@
             .bucket(uploadRequest.getBucket())
             .key(uploadRequest.getKey())
             .overrideConfiguration(o -> o.addMetricPublisher(statsMetricPublisher.multipartUploadMetricCollector));
+
+        if (uploadRequest.getMetadata() != null && !uploadRequest.getMetadata().isEmpty()) {
+            createMultipartUploadRequestBuilder.metadata(uploadRequest.getMetadata());
+        }
         if (uploadRequest.doRemoteDataIntegrityCheck()) {
             createMultipartUploadRequestBuilder.checksumAlgorithm(ChecksumAlgorithm.CRC32);
         }
@@ -327,6 +331,10 @@
             .key(uploadRequest.getKey())
             .contentLength(uploadRequest.getContentLength())
             .overrideConfiguration(o -> o.addMetricPublisher(statsMetricPublisher.putObjectMetricPublisher));
+
+        if (uploadRequest.getMetadata() != null && !uploadRequest.getMetadata().isEmpty()) {
+            putObjectRequestBuilder.metadata(uploadRequest.getMetadata());
+        }
         if (uploadRequest.doRemoteDataIntegrityCheck()) {
             putObjectRequestBuilder.checksumAlgorithm(ChecksumAlgorithm.CRC32);
             putObjectRequestBuilder.checksumCRC32(base64StringFromLong(uploadRequest.getExpectedChecksum()));

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

plugins/repository-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;
@@ -471,7 +473,7 @@ private void testWriteBlobByStreams(boolean expectException, boolean throwExcept
         StreamContextSupplier streamContextSupplier = partSize -> new StreamContext((partNo, size, position) -> {
             InputStream inputStream = new OffsetRangeIndexInputStream(new ByteArrayIndexInput("desc", bytes), size, position);
             openInputStreams.add(inputStream);
-            return new InputStreamContainer(inputStream, size, position);
+            return new InputStreamContainer(inputStream, size, position, null);
         }, partSize, calculateLastPartSize(bytes.length, partSize), calculateNumberOfParts(bytes.length, partSize));
 
         CheckedConsumer<Boolean, IOException> uploadFinalizer = uploadSuccess -> {
@@ -527,7 +529,7 @@ private void testWriteBlobByStreamsLargeBlob(boolean expectException, boolean th
         StreamContextSupplier streamContextSupplier = partSize1 -> new StreamContext((partNo, size, position) -> {
             InputStream inputStream = new OffsetRangeIndexInputStream(new ZeroIndexInput("desc", blobSize), size, position);
             openInputStreams.add(inputStream);
-            return new InputStreamContainer(inputStream, size, position);
+            return new InputStreamContainer(inputStream, size, position, null);
         }, partSize1, calculateLastPartSize(blobSize, partSize1), calculateNumberOfParts(blobSize, partSize1));
 
         CheckedConsumer<Boolean, IOException> uploadFinalizer = uploadSuccess -> {
@@ -649,7 +651,7 @@ private void testLargeFilesRedirectedToSlowSyncClient(boolean expectException) t
         StreamContextSupplier streamContextSupplier = partSize1 -> new StreamContext((partNo, size, position) -> {
             InputStream inputStream = new OffsetRangeIndexInputStream(new ZeroIndexInput("desc", blobSize), size, position);
             openInputStreams.add(inputStream);
-            return new InputStreamContainer(inputStream, size, position);
+            return new InputStreamContainer(inputStream, size, position, null);
         }, partSize1, calculateLastPartSize(blobSize, partSize1), calculateNumberOfParts(blobSize, partSize1));
 
         WriteContext writeContext = new WriteContext.Builder().fileName("write_large_blob")
@@ -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));