Examples with UploadPartResult com.amazonaws.services.s3.model.UploadPartResult used on opensource projects

Example 26 with UploadPartResult

use of com.amazonaws.services.s3.model.UploadPartResult in project nifi by apache.

the class PutS3Object method onTrigger.

@Override
public void onTrigger(final ProcessContext context, final ProcessSession session) {
    FlowFile flowFile = session.get();
    if (flowFile == null) {
        return;
    }
    final long startNanos = System.nanoTime();
    final String bucket = context.getProperty(BUCKET).evaluateAttributeExpressions(flowFile).getValue();
    final String key = context.getProperty(KEY).evaluateAttributeExpressions(flowFile).getValue();
    final String cacheKey = getIdentifier() + "/" + bucket + "/" + key;
    final AmazonS3Client s3 = getClient();
    final FlowFile ff = flowFile;
    final Map<String, String> attributes = new HashMap<>();
    final String ffFilename = ff.getAttributes().get(CoreAttributes.FILENAME.key());
    attributes.put(S3_BUCKET_KEY, bucket);
    attributes.put(S3_OBJECT_KEY, key);
    final Long multipartThreshold = context.getProperty(MULTIPART_THRESHOLD).asDataSize(DataUnit.B).longValue();
    final Long multipartPartSize = context.getProperty(MULTIPART_PART_SIZE).asDataSize(DataUnit.B).longValue();
    final long now = System.currentTimeMillis();
    /*
         * If necessary, run age off for existing uploads in AWS S3 and local state
         */
    ageoffS3Uploads(context, s3, now);
    /*
         * Then
         */
    try {
        session.read(flowFile, new InputStreamCallback() {

            @Override
            public void process(final InputStream rawIn) throws IOException {
                try (final InputStream in = new BufferedInputStream(rawIn)) {
                    final ObjectMetadata objectMetadata = new ObjectMetadata();
                    objectMetadata.setContentDisposition(ff.getAttribute(CoreAttributes.FILENAME.key()));
                    objectMetadata.setContentLength(ff.getSize());
                    final String contentType = context.getProperty(CONTENT_TYPE).evaluateAttributeExpressions(ff).getValue();
                    if (contentType != null) {
                        objectMetadata.setContentType(contentType);
                        attributes.put(S3_CONTENT_TYPE, contentType);
                    }
                    final String expirationRule = context.getProperty(EXPIRATION_RULE_ID).evaluateAttributeExpressions(ff).getValue();
                    if (expirationRule != null) {
                        objectMetadata.setExpirationTimeRuleId(expirationRule);
                    }
                    final Map<String, String> userMetadata = new HashMap<>();
                    for (final Map.Entry<PropertyDescriptor, String> entry : context.getProperties().entrySet()) {
                        if (entry.getKey().isDynamic()) {
                            final String value = context.getProperty(entry.getKey()).evaluateAttributeExpressions(ff).getValue();
                            userMetadata.put(entry.getKey().getName(), value);
                        }
                    }
                    final String serverSideEncryption = context.getProperty(SERVER_SIDE_ENCRYPTION).getValue();
                    if (!serverSideEncryption.equals(NO_SERVER_SIDE_ENCRYPTION)) {
                        objectMetadata.setSSEAlgorithm(serverSideEncryption);
                        attributes.put(S3_SSE_ALGORITHM, serverSideEncryption);
                    }
                    if (!userMetadata.isEmpty()) {
                        objectMetadata.setUserMetadata(userMetadata);
                    }
                    if (ff.getSize() <= multipartThreshold) {
                        // ----------------------------------------
                        // single part upload
                        // ----------------------------------------
                        final PutObjectRequest request = new PutObjectRequest(bucket, key, in, objectMetadata);
                        request.setStorageClass(StorageClass.valueOf(context.getProperty(STORAGE_CLASS).getValue()));
                        final AccessControlList acl = createACL(context, ff);
                        if (acl != null) {
                            request.setAccessControlList(acl);
                        }
                        final CannedAccessControlList cannedAcl = createCannedACL(context, ff);
                        if (cannedAcl != null) {
                            request.withCannedAcl(cannedAcl);
                        }
                        try {
                            final PutObjectResult result = s3.putObject(request);
                            if (result.getVersionId() != null) {
                                attributes.put(S3_VERSION_ATTR_KEY, result.getVersionId());
                            }
                            if (result.getETag() != null) {
                                attributes.put(S3_ETAG_ATTR_KEY, result.getETag());
                            }
                            if (result.getExpirationTime() != null) {
                                attributes.put(S3_EXPIRATION_ATTR_KEY, result.getExpirationTime().toString());
                            }
                            if (result.getMetadata().getRawMetadata().keySet().contains(S3_STORAGECLASS_META_KEY)) {
                                attributes.put(S3_STORAGECLASS_ATTR_KEY, result.getMetadata().getRawMetadataValue(S3_STORAGECLASS_META_KEY).toString());
                            }
                            if (userMetadata.size() > 0) {
                                StringBuilder userMetaBldr = new StringBuilder();
                                for (String userKey : userMetadata.keySet()) {
                                    userMetaBldr.append(userKey).append("=").append(userMetadata.get(userKey));
                                }
                                attributes.put(S3_USERMETA_ATTR_KEY, userMetaBldr.toString());
                            }
                            attributes.put(S3_API_METHOD_ATTR_KEY, S3_API_METHOD_PUTOBJECT);
                        } catch (AmazonClientException e) {
                            getLogger().info("Failure completing upload flowfile={} bucket={} key={} reason={}", new Object[] { ffFilename, bucket, key, e.getMessage() });
                            throw (e);
                        }
                    } else {
                        // ----------------------------------------
                        // multipart upload
                        // ----------------------------------------
                        // load or create persistent state
                        // ------------------------------------------------------------
                        MultipartState currentState;
                        try {
                            currentState = getLocalStateIfInS3(s3, bucket, cacheKey);
                            if (currentState != null) {
                                if (currentState.getPartETags().size() > 0) {
                                    final PartETag lastETag = currentState.getPartETags().get(currentState.getPartETags().size() - 1);
                                    getLogger().info("Resuming upload for flowfile='{}' bucket='{}' key='{}' " + "uploadID='{}' filePosition='{}' partSize='{}' storageClass='{}' " + "contentLength='{}' partsLoaded={} lastPart={}/{}", new Object[] { ffFilename, bucket, key, currentState.getUploadId(), currentState.getFilePosition(), currentState.getPartSize(), currentState.getStorageClass().toString(), currentState.getContentLength(), currentState.getPartETags().size(), Integer.toString(lastETag.getPartNumber()), lastETag.getETag() });
                                } else {
                                    getLogger().info("Resuming upload for flowfile='{}' bucket='{}' key='{}' " + "uploadID='{}' filePosition='{}' partSize='{}' storageClass='{}' " + "contentLength='{}' no partsLoaded", new Object[] { ffFilename, bucket, key, currentState.getUploadId(), currentState.getFilePosition(), currentState.getPartSize(), currentState.getStorageClass().toString(), currentState.getContentLength() });
                                }
                            } else {
                                currentState = new MultipartState();
                                currentState.setPartSize(multipartPartSize);
                                currentState.setStorageClass(StorageClass.valueOf(context.getProperty(STORAGE_CLASS).getValue()));
                                currentState.setContentLength(ff.getSize());
                                persistLocalState(cacheKey, currentState);
                                getLogger().info("Starting new upload for flowfile='{}' bucket='{}' key='{}'", new Object[] { ffFilename, bucket, key });
                            }
                        } catch (IOException e) {
                            getLogger().error("IOException initiating cache state while processing flow files: " + e.getMessage());
                            throw (e);
                        }
                        // ------------------------------------------------------------
                        if (currentState.getUploadId().isEmpty()) {
                            final InitiateMultipartUploadRequest initiateRequest = new InitiateMultipartUploadRequest(bucket, key, objectMetadata);
                            initiateRequest.setStorageClass(currentState.getStorageClass());
                            final AccessControlList acl = createACL(context, ff);
                            if (acl != null) {
                                initiateRequest.setAccessControlList(acl);
                            }
                            final CannedAccessControlList cannedAcl = createCannedACL(context, ff);
                            if (cannedAcl != null) {
                                initiateRequest.withCannedACL(cannedAcl);
                            }
                            try {
                                final InitiateMultipartUploadResult initiateResult = s3.initiateMultipartUpload(initiateRequest);
                                currentState.setUploadId(initiateResult.getUploadId());
                                currentState.getPartETags().clear();
                                try {
                                    persistLocalState(cacheKey, currentState);
                                } catch (Exception e) {
                                    getLogger().info("Exception saving cache state while processing flow file: " + e.getMessage());
                                    throw (new ProcessException("Exception saving cache state", e));
                                }
                                getLogger().info("Success initiating upload flowfile={} available={} position={} " + "length={} bucket={} key={} uploadId={}", new Object[] { ffFilename, in.available(), currentState.getFilePosition(), currentState.getContentLength(), bucket, key, currentState.getUploadId() });
                                if (initiateResult.getUploadId() != null) {
                                    attributes.put(S3_UPLOAD_ID_ATTR_KEY, initiateResult.getUploadId());
                                }
                            } catch (AmazonClientException e) {
                                getLogger().info("Failure initiating upload flowfile={} bucket={} key={} reason={}", new Object[] { ffFilename, bucket, key, e.getMessage() });
                                throw (e);
                            }
                        } else {
                            if (currentState.getFilePosition() > 0) {
                                try {
                                    final long skipped = in.skip(currentState.getFilePosition());
                                    if (skipped != currentState.getFilePosition()) {
                                        getLogger().info("Failure skipping to resume upload flowfile={} " + "bucket={} key={} position={} skipped={}", new Object[] { ffFilename, bucket, key, currentState.getFilePosition(), skipped });
                                    }
                                } catch (Exception e) {
                                    getLogger().info("Failure skipping to resume upload flowfile={} bucket={} " + "key={} position={} reason={}", new Object[] { ffFilename, bucket, key, currentState.getFilePosition(), e.getMessage() });
                                    throw (new ProcessException(e));
                                }
                            }
                        }
                        // upload parts
                        // ------------------------------------------------------------
                        long thisPartSize;
                        for (int part = currentState.getPartETags().size() + 1; currentState.getFilePosition() < currentState.getContentLength(); part++) {
                            if (!PutS3Object.this.isScheduled()) {
                                throw new IOException(S3_PROCESS_UNSCHEDULED_MESSAGE + " flowfile=" + ffFilename + " part=" + part + " uploadId=" + currentState.getUploadId());
                            }
                            thisPartSize = Math.min(currentState.getPartSize(), (currentState.getContentLength() - currentState.getFilePosition()));
                            UploadPartRequest uploadRequest = new UploadPartRequest().withBucketName(bucket).withKey(key).withUploadId(currentState.getUploadId()).withInputStream(in).withPartNumber(part).withPartSize(thisPartSize);
                            try {
                                UploadPartResult uploadPartResult = s3.uploadPart(uploadRequest);
                                currentState.addPartETag(uploadPartResult.getPartETag());
                                currentState.setFilePosition(currentState.getFilePosition() + thisPartSize);
                                try {
                                    persistLocalState(cacheKey, currentState);
                                } catch (Exception e) {
                                    getLogger().info("Exception saving cache state processing flow file: " + e.getMessage());
                                }
                                getLogger().info("Success uploading part flowfile={} part={} available={} " + "etag={} uploadId={}", new Object[] { ffFilename, part, in.available(), uploadPartResult.getETag(), currentState.getUploadId() });
                            } catch (AmazonClientException e) {
                                getLogger().info("Failure uploading part flowfile={} part={} bucket={} key={} " + "reason={}", new Object[] { ffFilename, part, bucket, key, e.getMessage() });
                                throw (e);
                            }
                        }
                        // complete multipart upload
                        // ------------------------------------------------------------
                        CompleteMultipartUploadRequest completeRequest = new CompleteMultipartUploadRequest(bucket, key, currentState.getUploadId(), currentState.getPartETags());
                        try {
                            CompleteMultipartUploadResult completeResult = s3.completeMultipartUpload(completeRequest);
                            getLogger().info("Success completing upload flowfile={} etag={} uploadId={}", new Object[] { ffFilename, completeResult.getETag(), currentState.getUploadId() });
                            if (completeResult.getVersionId() != null) {
                                attributes.put(S3_VERSION_ATTR_KEY, completeResult.getVersionId());
                            }
                            if (completeResult.getETag() != null) {
                                attributes.put(S3_ETAG_ATTR_KEY, completeResult.getETag());
                            }
                            if (completeResult.getExpirationTime() != null) {
                                attributes.put(S3_EXPIRATION_ATTR_KEY, completeResult.getExpirationTime().toString());
                            }
                            if (currentState.getStorageClass() != null) {
                                attributes.put(S3_STORAGECLASS_ATTR_KEY, currentState.getStorageClass().toString());
                            }
                            if (userMetadata.size() > 0) {
                                StringBuilder userMetaBldr = new StringBuilder();
                                for (String userKey : userMetadata.keySet()) {
                                    userMetaBldr.append(userKey).append("=").append(userMetadata.get(userKey));
                                }
                                attributes.put(S3_USERMETA_ATTR_KEY, userMetaBldr.toString());
                            }
                            attributes.put(S3_API_METHOD_ATTR_KEY, S3_API_METHOD_MULTIPARTUPLOAD);
                        } catch (AmazonClientException e) {
                            getLogger().info("Failure completing upload flowfile={} bucket={} key={} reason={}", new Object[] { ffFilename, bucket, key, e.getMessage() });
                            throw (e);
                        }
                    }
                }
            }
        });
        if (!attributes.isEmpty()) {
            flowFile = session.putAllAttributes(flowFile, attributes);
        }
        session.transfer(flowFile, REL_SUCCESS);
        final String url = s3.getResourceUrl(bucket, key);
        final long millis = TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - startNanos);
        session.getProvenanceReporter().send(flowFile, url, millis);
        getLogger().info("Successfully put {} to Amazon S3 in {} milliseconds", new Object[] { ff, millis });
        try {
            removeLocalState(cacheKey);
        } catch (IOException e) {
            getLogger().info("Error trying to delete key {} from cache: {}", new Object[] { cacheKey, e.getMessage() });
        }
    } catch (final ProcessException | AmazonClientException pe) {
        if (pe.getMessage().contains(S3_PROCESS_UNSCHEDULED_MESSAGE)) {
            getLogger().info(pe.getMessage());
            session.rollback();
        } else {
            getLogger().error("Failed to put {} to Amazon S3 due to {}", new Object[] { flowFile, pe });
            flowFile = session.penalize(flowFile);
            session.transfer(flowFile, REL_FAILURE);
        }
    }
}

Example 27 with UploadPartResult

use of com.amazonaws.services.s3.model.UploadPartResult in project crate by crate.

the class S3BlobStoreContainerTests method testExecuteMultipartUploadAborted.

@Test
public void testExecuteMultipartUploadAborted() {
    final String bucketName = randomAlphaOfLengthBetween(1, 10);
    final String blobName = randomAlphaOfLengthBetween(1, 10);
    final BlobPath blobPath = new BlobPath();
    final long blobSize = ByteSizeUnit.MB.toBytes(765);
    final long bufferSize = ByteSizeUnit.MB.toBytes(150);
    final S3BlobStore blobStore = mock(S3BlobStore.class);
    when(blobStore.bucket()).thenReturn(bucketName);
    when(blobStore.bufferSizeInBytes()).thenReturn(bufferSize);
    when(blobStore.getStorageClass()).thenReturn(randomFrom(StorageClass.values()));
    final AmazonS3 client = mock(AmazonS3.class);
    final AmazonS3Reference clientReference = new AmazonS3Reference(client);
    doAnswer(invocation -> {
        clientReference.incRef();
        return clientReference;
    }).when(blobStore).clientReference();
    final String uploadId = randomAlphaOfLength(25);
    final int stage = randomInt(2);
    final List<AmazonClientException> exceptions = Arrays.asList(new AmazonClientException("Expected initialization request to fail"), new AmazonClientException("Expected upload part request to fail"), new AmazonClientException("Expected completion request to fail"));
    if (stage == 0) {
        // Fail the initialization request
        when(client.initiateMultipartUpload(any(InitiateMultipartUploadRequest.class))).thenThrow(exceptions.get(stage));
    } else if (stage == 1) {
        final InitiateMultipartUploadResult initResult = new InitiateMultipartUploadResult();
        initResult.setUploadId(uploadId);
        when(client.initiateMultipartUpload(any(InitiateMultipartUploadRequest.class))).thenReturn(initResult);
        // Fail the upload part request
        when(client.uploadPart(any(UploadPartRequest.class))).thenThrow(exceptions.get(stage));
    } else {
        final InitiateMultipartUploadResult initResult = new InitiateMultipartUploadResult();
        initResult.setUploadId(uploadId);
        when(client.initiateMultipartUpload(any(InitiateMultipartUploadRequest.class))).thenReturn(initResult);
        when(client.uploadPart(any(UploadPartRequest.class))).thenAnswer(invocationOnMock -> {
            final UploadPartRequest request = (UploadPartRequest) invocationOnMock.getArguments()[0];
            final UploadPartResult response = new UploadPartResult();
            response.setPartNumber(request.getPartNumber());
            response.setETag(randomAlphaOfLength(20));
            return response;
        });
        // Fail the completion request
        when(client.completeMultipartUpload(any(CompleteMultipartUploadRequest.class))).thenThrow(exceptions.get(stage));
    }
    final ArgumentCaptor<AbortMultipartUploadRequest> argumentCaptor = ArgumentCaptor.forClass(AbortMultipartUploadRequest.class);
    doNothing().when(client).abortMultipartUpload(argumentCaptor.capture());
    final IOException e = expectThrows(IOException.class, () -> {
        final S3BlobContainer blobContainer = new S3BlobContainer(blobPath, blobStore);
        blobContainer.executeMultipartUpload(blobStore, blobName, new ByteArrayInputStream(new byte[0]), blobSize);
    });
    assertEquals("Unable to upload object [" + blobName + "] using multipart upload", e.getMessage());
    assertThat(e.getCause(), instanceOf(AmazonClientException.class));
    assertEquals(exceptions.get(stage).getMessage(), e.getCause().getMessage());
    if (stage == 0) {
        verify(client, times(1)).initiateMultipartUpload(any(InitiateMultipartUploadRequest.class));
        verify(client, times(0)).uploadPart(any(UploadPartRequest.class));
        verify(client, times(0)).completeMultipartUpload(any(CompleteMultipartUploadRequest.class));
        verify(client, times(0)).abortMultipartUpload(any(AbortMultipartUploadRequest.class));
    } else {
        verify(client, times(1)).initiateMultipartUpload(any(InitiateMultipartUploadRequest.class));
        if (stage == 1) {
            verify(client, times(1)).uploadPart(any(UploadPartRequest.class));
            verify(client, times(0)).completeMultipartUpload(any(CompleteMultipartUploadRequest.class));
        } else {
            verify(client, times(6)).uploadPart(any(UploadPartRequest.class));
            verify(client, times(1)).completeMultipartUpload(any(CompleteMultipartUploadRequest.class));
        }
        verify(client, times(1)).abortMultipartUpload(any(AbortMultipartUploadRequest.class));
        final AbortMultipartUploadRequest abortRequest = argumentCaptor.getValue();
        assertEquals(bucketName, abortRequest.getBucketName());
        assertEquals(blobName, abortRequest.getKey());
        assertEquals(uploadId, abortRequest.getUploadId());
    }
}

Example 28 with UploadPartResult

use of com.amazonaws.services.s3.model.UploadPartResult 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);
}

Example 29 with UploadPartResult

use of com.amazonaws.services.s3.model.UploadPartResult in project kafka-connect-storage-cloud by confluentinc.

the class S3SinkTaskTest method testWriteRecordsSpanningMultiplePartsWithRetry.

@Test
public void testWriteRecordsSpanningMultiplePartsWithRetry() throws Exception {
    localProps.put(S3SinkConnectorConfig.FLUSH_SIZE_CONFIG, "10000");
    localProps.put(S3SinkConnectorConfig.S3_PART_RETRIES_CONFIG, "3");
    setUp();
    List<SinkRecord> sinkRecords = createRecords(11000);
    int totalBytes = calcByteSize(sinkRecords);
    final int parts = totalBytes / connectorConfig.getPartSize();
    // From time to time fail S3 upload part method
    final AtomicInteger count = new AtomicInteger();
    PowerMockito.doAnswer(new Answer<UploadPartResult>() {

        @Override
        public UploadPartResult answer(InvocationOnMock invocationOnMock) throws Throwable {
            if (count.getAndIncrement() % parts == 0) {
                throw new SdkClientException("Boom!");
            } else {
                return (UploadPartResult) invocationOnMock.callRealMethod();
            }
        }
    }).when(s3).uploadPart(Mockito.isA(UploadPartRequest.class));
    replayAll();
    task = new S3SinkTask();
    task.initialize(context);
    task.start(properties);
    verifyAll();
    task.put(sinkRecords);
    task.close(context.assignment());
    task.stop();
    long[] validOffsets = { 0, 10000 };
    verify(sinkRecords, validOffsets);
}

Example 30 with UploadPartResult

use of com.amazonaws.services.s3.model.UploadPartResult in project bender by Nextdoor.

the class S3Transport method sendStream.

protected void sendStream(InputStream input, String key, long streamSize) throws TransportException {
    /*
     * Create metadata
     */
    ObjectMetadata metadata = new ObjectMetadata();
    /*
     * Find if a multipart upload has already begun or start a new one.
     */
    MultiPartUpload upload;
    synchronized (multiPartUploads) {
        if (!multiPartUploads.containsKey(key)) {
            InitiateMultipartUploadRequest uploadRequest = new InitiateMultipartUploadRequest(bucketName, key);
            uploadRequest.setObjectMetadata(metadata);
            InitiateMultipartUploadResult res = client.initiateMultipartUpload(uploadRequest);
            upload = new MultiPartUpload(bucketName, key, res.getUploadId());
            multiPartUploads.put(key, upload);
        } else {
            upload = multiPartUploads.get(key);
        }
    }
    /*
     * Write out to S3. Note that the S3 client auto closes the input stream.
     */
    UploadPartRequest req = upload.getUploadPartRequest().withInputStream(input).withPartSize(streamSize);
    req.getRequestClientOptions().setReadLimit((int) (streamSize + 1));
    try {
        UploadPartResult res = client.uploadPart(req);
        upload.addPartETag(res.getPartETag());
    } catch (AmazonClientException e) {
        client.abortMultipartUpload(upload.getAbortMultipartUploadRequest());
        multiPartUploads.remove(key);
        String debuggingMessage = String.format("unable to put file with key %s and uploadId %s when trying to upload part number %s, leading to exception: ", key, upload.getUploadId(), upload.getPartCount());
        throw new TransportException(debuggingMessage, e);
    } finally {
        try {
            input.close();
        } catch (IOException e) {
            logger.warn("error encountered while closing input stream", e);
        }
    }
}