Search in sources :

Example 21 with CompleteMultipartUploadRequest

use of software.amazon.awssdk.services.s3.model.CompleteMultipartUploadRequest in project crate by crate.

the class S3BlobStoreContainerTests method testExecuteMultipartUpload.

@Test
public void testExecuteMultipartUpload() throws IOException {
    final String bucketName = randomAlphaOfLengthBetween(1, 10);
    final String blobName = randomAlphaOfLengthBetween(1, 10);
    final BlobPath blobPath = new BlobPath();
    if (randomBoolean()) {
        IntStream.of(randomIntBetween(1, 5)).forEach(value -> blobPath.add("path_" + value));
    }
    final long blobSize = ByteSizeUnit.GB.toBytes(randomIntBetween(1, 128));
    final long bufferSize = ByteSizeUnit.MB.toBytes(randomIntBetween(5, 1024));
    final S3BlobStore blobStore = mock(S3BlobStore.class);
    when(blobStore.bucket()).thenReturn(bucketName);
    when(blobStore.bufferSizeInBytes()).thenReturn(bufferSize);
    final boolean serverSideEncryption = randomBoolean();
    when(blobStore.serverSideEncryption()).thenReturn(serverSideEncryption);
    final StorageClass storageClass = randomFrom(StorageClass.values());
    when(blobStore.getStorageClass()).thenReturn(storageClass);
    final CannedAccessControlList cannedAccessControlList = randomBoolean() ? randomFrom(CannedAccessControlList.values()) : null;
    if (cannedAccessControlList != null) {
        when(blobStore.getCannedACL()).thenReturn(cannedAccessControlList);
    }
    final AmazonS3 client = mock(AmazonS3.class);
    final AmazonS3Reference clientReference = new AmazonS3Reference(client);
    when(blobStore.clientReference()).thenReturn(clientReference);
    final ArgumentCaptor<InitiateMultipartUploadRequest> initArgCaptor = ArgumentCaptor.forClass(InitiateMultipartUploadRequest.class);
    final InitiateMultipartUploadResult initResult = new InitiateMultipartUploadResult();
    initResult.setUploadId(randomAlphaOfLength(10));
    when(client.initiateMultipartUpload(initArgCaptor.capture())).thenReturn(initResult);
    final ArgumentCaptor<UploadPartRequest> uploadArgCaptor = ArgumentCaptor.forClass(UploadPartRequest.class);
    final List<String> expectedEtags = new ArrayList<>();
    final long partSize = Math.min(bufferSize, blobSize);
    long totalBytes = 0;
    do {
        expectedEtags.add(randomAlphaOfLength(50));
        totalBytes += partSize;
    } while (totalBytes < blobSize);
    when(client.uploadPart(uploadArgCaptor.capture())).thenAnswer(invocationOnMock -> {
        final UploadPartRequest request = (UploadPartRequest) invocationOnMock.getArguments()[0];
        final UploadPartResult response = new UploadPartResult();
        response.setPartNumber(request.getPartNumber());
        response.setETag(expectedEtags.get(request.getPartNumber() - 1));
        return response;
    });
    final ArgumentCaptor<CompleteMultipartUploadRequest> compArgCaptor = ArgumentCaptor.forClass(CompleteMultipartUploadRequest.class);
    when(client.completeMultipartUpload(compArgCaptor.capture())).thenReturn(new CompleteMultipartUploadResult());
    final ByteArrayInputStream inputStream = new ByteArrayInputStream(new byte[0]);
    final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore);
    blobContainer.executeMultipartUpload(blobStore, blobName, inputStream, blobSize);
    final InitiateMultipartUploadRequest initRequest = initArgCaptor.getValue();
    assertEquals(bucketName, initRequest.getBucketName());
    assertEquals(blobPath.buildAsString() + blobName, initRequest.getKey());
    assertEquals(storageClass, initRequest.getStorageClass());
    assertEquals(cannedAccessControlList, initRequest.getCannedACL());
    if (serverSideEncryption) {
        assertEquals(ObjectMetadata.AES_256_SERVER_SIDE_ENCRYPTION, initRequest.getObjectMetadata().getSSEAlgorithm());
    }
    final Tuple<Long, Long> numberOfParts = S3BlobContainer.numberOfMultiparts(blobSize, bufferSize);
    final List<UploadPartRequest> uploadRequests = uploadArgCaptor.getAllValues();
    assertEquals(numberOfParts.v1().intValue(), uploadRequests.size());
    for (int i = 0; i < uploadRequests.size(); i++) {
        final UploadPartRequest uploadRequest = uploadRequests.get(i);
        assertEquals(bucketName, uploadRequest.getBucketName());
        assertEquals(blobPath.buildAsString() + blobName, uploadRequest.getKey());
        assertEquals(initResult.getUploadId(), uploadRequest.getUploadId());
        assertEquals(i + 1, uploadRequest.getPartNumber());
        assertEquals(inputStream, uploadRequest.getInputStream());
        if (i == (uploadRequests.size() - 1)) {
            assertTrue(uploadRequest.isLastPart());
            assertEquals(numberOfParts.v2().longValue(), uploadRequest.getPartSize());
        } else {
            assertFalse(uploadRequest.isLastPart());
            assertEquals(bufferSize, uploadRequest.getPartSize());
        }
    }
    final CompleteMultipartUploadRequest compRequest = compArgCaptor.getValue();
    assertEquals(bucketName, compRequest.getBucketName());
    assertEquals(blobPath.buildAsString() + blobName, compRequest.getKey());
    assertEquals(initResult.getUploadId(), compRequest.getUploadId());
    final List<String> actualETags = compRequest.getPartETags().stream().map(PartETag::getETag).collect(Collectors.toList());
    assertEquals(expectedEtags, actualETags);
}
Also used : AmazonS3(com.amazonaws.services.s3.AmazonS3) InitiateMultipartUploadResult(com.amazonaws.services.s3.model.InitiateMultipartUploadResult) ArrayList(java.util.ArrayList) Matchers.containsString(org.hamcrest.Matchers.containsString) CompleteMultipartUploadResult(com.amazonaws.services.s3.model.CompleteMultipartUploadResult) UploadPartResult(com.amazonaws.services.s3.model.UploadPartResult) BlobPath(org.elasticsearch.common.blobstore.BlobPath) S3BlobStoreTests.randomMockS3BlobStore(org.elasticsearch.repositories.s3.S3BlobStoreTests.randomMockS3BlobStore) StorageClass(com.amazonaws.services.s3.model.StorageClass) InitiateMultipartUploadRequest(com.amazonaws.services.s3.model.InitiateMultipartUploadRequest) UploadPartRequest(com.amazonaws.services.s3.model.UploadPartRequest) CannedAccessControlList(com.amazonaws.services.s3.model.CannedAccessControlList) ByteArrayInputStream(java.io.ByteArrayInputStream) CompleteMultipartUploadRequest(com.amazonaws.services.s3.model.CompleteMultipartUploadRequest) Test(org.junit.Test)

Example 22 with CompleteMultipartUploadRequest

use of software.amazon.awssdk.services.s3.model.CompleteMultipartUploadRequest in project Singularity by HubSpot.

the class SingularityS3Uploader method multipartUpload.

private void multipartUpload(AmazonS3 s3Client, String key, File file, ObjectMetadata objectMetadata, Optional<StorageClass> maybeStorageClass) throws Exception {
    List<PartETag> partETags = new ArrayList<>();
    InitiateMultipartUploadRequest initRequest = new InitiateMultipartUploadRequest(bucketName, key, objectMetadata);
    maybeStorageClass.ifPresent(initRequest::setStorageClass);
    InitiateMultipartUploadResult initResponse = s3Client.initiateMultipartUpload(initRequest);
    long contentLength = file.length();
    long partSize = configuration.getUploadPartSize();
    try {
        long filePosition = 0;
        for (int i = 1; filePosition < contentLength; i++) {
            partSize = Math.min(partSize, (contentLength - filePosition));
            UploadPartRequest uploadRequest = new UploadPartRequest().withBucketName(bucketName).withKey(key).withUploadId(initResponse.getUploadId()).withPartNumber(i).withFileOffset(filePosition).withFile(file).withPartSize(partSize);
            partETags.add(s3Client.uploadPart(uploadRequest).getPartETag());
            filePosition += partSize;
        }
        CompleteMultipartUploadRequest completeRequest = new CompleteMultipartUploadRequest(bucketName, key, initResponse.getUploadId(), partETags);
        s3Client.completeMultipartUpload(completeRequest);
    } catch (Exception e) {
        s3Client.abortMultipartUpload(new AbortMultipartUploadRequest(bucketName, key, initResponse.getUploadId()));
        throw new RuntimeException(e);
    }
}
Also used : InitiateMultipartUploadResult(com.amazonaws.services.s3.model.InitiateMultipartUploadResult) ArrayList(java.util.ArrayList) InitiateMultipartUploadRequest(com.amazonaws.services.s3.model.InitiateMultipartUploadRequest) UploadPartRequest(com.amazonaws.services.s3.model.UploadPartRequest) AbortMultipartUploadRequest(com.amazonaws.services.s3.model.AbortMultipartUploadRequest) PartETag(com.amazonaws.services.s3.model.PartETag) NoSuchFileException(java.nio.file.NoSuchFileException) AmazonS3Exception(com.amazonaws.services.s3.model.AmazonS3Exception) CompleteMultipartUploadRequest(com.amazonaws.services.s3.model.CompleteMultipartUploadRequest)

Aggregations

CompleteMultipartUploadRequest (com.amazonaws.services.s3.model.CompleteMultipartUploadRequest)17 InitiateMultipartUploadRequest (com.amazonaws.services.s3.model.InitiateMultipartUploadRequest)12 PartETag (com.amazonaws.services.s3.model.PartETag)12 InitiateMultipartUploadResult (com.amazonaws.services.s3.model.InitiateMultipartUploadResult)11 ArrayList (java.util.ArrayList)10 UploadPartRequest (com.amazonaws.services.s3.model.UploadPartRequest)9 AbortMultipartUploadRequest (com.amazonaws.services.s3.model.AbortMultipartUploadRequest)8 IOException (java.io.IOException)7 ObjectMetadata (com.amazonaws.services.s3.model.ObjectMetadata)6 AmazonClientException (com.amazonaws.AmazonClientException)5 CompleteMultipartUploadResult (com.amazonaws.services.s3.model.CompleteMultipartUploadResult)5 UploadPartResult (com.amazonaws.services.s3.model.UploadPartResult)5 CompleteMultipartUploadRequest (software.amazon.awssdk.services.s3.model.CompleteMultipartUploadRequest)5 AmazonS3Exception (com.amazonaws.services.s3.model.AmazonS3Exception)4 CreateMultipartUploadRequest (software.amazon.awssdk.services.s3.model.CreateMultipartUploadRequest)4 CreateMultipartUploadResponse (software.amazon.awssdk.services.s3.model.CreateMultipartUploadResponse)4 CannedAccessControlList (com.amazonaws.services.s3.model.CannedAccessControlList)3 InputStream (java.io.InputStream)3 CompletedMultipartUpload (software.amazon.awssdk.services.s3.model.CompletedMultipartUpload)3 UploadPartRequest (software.amazon.awssdk.services.s3.model.UploadPartRequest)3