-
Notifications
You must be signed in to change notification settings - Fork 190
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Added SQS interactions for S3 source (#1431)
* Added sqs configuration and basic sqs interactions Signed-off-by: Asif Sohail Mohammed <nsifmoh@amazon.com>
- Loading branch information
1 parent
be35437
commit 8d1454c
Showing
12 changed files
with
543 additions
and
35 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
42 changes: 42 additions & 0 deletions
42
...pper-plugins/s3-source/src/main/java/com/amazon/dataprepper/plugins/source/S3Service.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,42 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package com.amazon.dataprepper.plugins.source; | ||
|
|
||
| import org.slf4j.Logger; | ||
| import org.slf4j.LoggerFactory; | ||
| import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration; | ||
| import software.amazon.awssdk.core.retry.RetryPolicy; | ||
| import software.amazon.awssdk.regions.Region; | ||
| import software.amazon.awssdk.services.s3.S3Client; | ||
| import software.amazon.awssdk.services.sts.StsClient; | ||
|
|
||
| public class S3Service { | ||
| private static final Logger LOG = LoggerFactory.getLogger(S3Service.class); | ||
|
|
||
| private final S3SourceConfig s3SourceConfig; | ||
| private final S3Client s3Client; | ||
|
|
||
| public S3Service(final S3SourceConfig s3SourceConfig) { | ||
| this.s3SourceConfig = s3SourceConfig; | ||
| this.s3Client = createS3Client(StsClient.create()); | ||
| } | ||
|
|
||
| S3ObjectReference addS3Object(final S3ObjectReference s3ObjectReference) { | ||
| // TODO: should return message id and receipt handle if successfully converted to event | ||
| return null; | ||
| } | ||
|
|
||
| S3Client createS3Client(final StsClient stsClient) { | ||
| LOG.info("Creating S3 client"); | ||
| return S3Client.builder() | ||
| .region(Region.of(s3SourceConfig.getAWSAuthenticationOptions().getAwsRegion())) | ||
| .credentialsProvider(s3SourceConfig.getAWSAuthenticationOptions().authenticateAwsConfiguration(stsClient)) | ||
| .overrideConfiguration(ClientOverrideConfiguration.builder() | ||
| .retryPolicy(RetryPolicy.builder().numRetries(5).build()) | ||
| .build()) | ||
| .build(); | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
50 changes: 50 additions & 0 deletions
50
...per-plugins/s3-source/src/main/java/com/amazon/dataprepper/plugins/source/SqsService.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,50 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package com.amazon.dataprepper.plugins.source; | ||
|
|
||
| import org.slf4j.Logger; | ||
| import org.slf4j.LoggerFactory; | ||
| import software.amazon.awssdk.core.client.config.ClientOverrideConfiguration; | ||
| import software.amazon.awssdk.core.retry.RetryPolicy; | ||
| import software.amazon.awssdk.regions.Region; | ||
| import software.amazon.awssdk.services.sqs.SqsClient; | ||
| import software.amazon.awssdk.services.sts.StsClient; | ||
|
|
||
| public class SqsService { | ||
| private static final Logger LOG = LoggerFactory.getLogger(SqsService.class); | ||
|
|
||
| private final S3SourceConfig s3SourceConfig; | ||
| private final S3Service s3Accessor; | ||
| private final SqsClient sqsClient; | ||
|
|
||
| private Thread sqsWorkerThread; | ||
|
|
||
| public SqsService(final S3SourceConfig s3SourceConfig, final S3Service s3Accessor) { | ||
| this.s3SourceConfig = s3SourceConfig; | ||
| this.s3Accessor = s3Accessor; | ||
| this.sqsClient = createSqsClient(StsClient.create()); | ||
| } | ||
|
|
||
| public void start() { | ||
| sqsWorkerThread = new Thread(new SqsWorker(sqsClient, s3Accessor, s3SourceConfig)); | ||
| sqsWorkerThread.start(); | ||
| } | ||
|
|
||
| SqsClient createSqsClient(final StsClient stsClient) { | ||
| LOG.info("Creating SQS client"); | ||
| return SqsClient.builder() | ||
| .region(Region.of(s3SourceConfig.getAWSAuthenticationOptions().getAwsRegion())) | ||
| .credentialsProvider(s3SourceConfig.getAWSAuthenticationOptions().authenticateAwsConfiguration(stsClient)) | ||
| .overrideConfiguration(ClientOverrideConfiguration.builder() | ||
| .retryPolicy(RetryPolicy.builder().numRetries(5).build()) | ||
| .build()) | ||
| .build(); | ||
| } | ||
|
|
||
| public void stop() { | ||
| sqsWorkerThread.interrupt(); | ||
| } | ||
| } |
113 changes: 113 additions & 0 deletions
113
...pper-plugins/s3-source/src/main/java/com/amazon/dataprepper/plugins/source/SqsWorker.java
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,113 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package com.amazon.dataprepper.plugins.source; | ||
|
|
||
| import com.amazon.dataprepper.plugins.source.configuration.SqsOptions; | ||
| import com.amazon.dataprepper.plugins.source.filter.ObjectCreatedFilter; | ||
| import com.amazon.dataprepper.plugins.source.filter.S3EventFilter; | ||
| import com.amazonaws.services.s3.event.S3EventNotification; | ||
| import org.slf4j.Logger; | ||
| import org.slf4j.LoggerFactory; | ||
| import software.amazon.awssdk.services.sqs.SqsClient; | ||
| import software.amazon.awssdk.services.sqs.model.Message; | ||
| import software.amazon.awssdk.services.sqs.model.ReceiveMessageRequest; | ||
| import software.amazon.awssdk.services.sqs.model.SqsException; | ||
|
|
||
| import java.util.ArrayList; | ||
| import java.util.List; | ||
| import java.util.Map; | ||
| import java.util.Optional; | ||
| import java.util.stream.Collectors; | ||
|
|
||
| public class SqsWorker implements Runnable { | ||
| private static final Logger LOG = LoggerFactory.getLogger(SqsWorker.class); | ||
|
|
||
| private final S3SourceConfig s3SourceConfig; | ||
| private final SqsClient sqsClient; | ||
| private final S3Service s3Service; | ||
| private final SqsOptions sqsOptions; | ||
| private final S3EventFilter objectCreatedFilter; | ||
|
|
||
| public SqsWorker(final SqsClient sqsClient, final S3Service s3Service, final S3SourceConfig s3SourceConfig) { | ||
| this.s3SourceConfig = s3SourceConfig; | ||
| this.sqsClient = sqsClient; | ||
| this.s3Service = s3Service; | ||
| sqsOptions = s3SourceConfig.getSqsOptions(); | ||
| objectCreatedFilter = new ObjectCreatedFilter(); | ||
| } | ||
|
|
||
| @Override | ||
| public void run() { | ||
|
|
||
| while(true) { | ||
| // get messages from SQS | ||
| List<Message> sqsMessages = getMessagesFromSqs(); | ||
|
|
||
| // convert each message to S3EventNotificationRecord | ||
| Map<Message, S3EventNotification.S3EventNotificationRecord> s3EventNotificationRecords = | ||
| getS3MessageEventNotificationRecordMap(sqsMessages); | ||
|
|
||
| // build s3ObjectReference from S3EventNotificationRecord if event name starts with ObjectCreated | ||
| processS3ObjectAndDeleteSqsMessages(s3EventNotificationRecords); | ||
|
|
||
| if (sqsMessages.size() < sqsOptions.getMaximumMessages() && s3SourceConfig.getSqsOptions().getPollDelay().toMillis() > 0) { | ||
| try { | ||
| Thread.sleep(s3SourceConfig.getSqsOptions().getPollDelay().toMillis()); | ||
| } catch (InterruptedException e) { | ||
| LOG.error("Thread is interrupted while polling SQS.", e); | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| List<Message> getMessagesFromSqs() { | ||
| List<Message> messages = new ArrayList<>(); | ||
| try { | ||
| ReceiveMessageRequest receiveMessageRequest = createReceiveMessageRequest(); | ||
| messages = sqsClient.receiveMessage(receiveMessageRequest).messages(); | ||
| } catch (SqsException e) { | ||
| LOG.error("Error reading from SQS: {}", e.awsErrorDetails().errorMessage()); | ||
| } | ||
| return messages; | ||
| } | ||
|
|
||
| ReceiveMessageRequest createReceiveMessageRequest() { | ||
| return ReceiveMessageRequest.builder() | ||
| .queueUrl(sqsOptions.getSqsUrl()) | ||
| .maxNumberOfMessages(sqsOptions.getMaximumMessages()) | ||
| .visibilityTimeout((int) sqsOptions.getVisibilityTimeout().getSeconds()) | ||
| .waitTimeSeconds((int) sqsOptions.getWaitTime().getSeconds()) | ||
| .build(); | ||
| } | ||
|
|
||
| private Map<Message, S3EventNotification.S3EventNotificationRecord> getS3MessageEventNotificationRecordMap(final List<Message> sqsMessages) { | ||
| return sqsMessages.stream().collect(Collectors.toMap(message -> message, this::convertS3EventMessages)); | ||
| } | ||
|
|
||
| S3EventNotification.S3EventNotificationRecord convertS3EventMessages(final Message message) { | ||
| return S3EventNotification.parseJson(message.body()).getRecords().get(0); | ||
| } | ||
|
|
||
| private void processS3ObjectAndDeleteSqsMessages(final Map<Message, S3EventNotification.S3EventNotificationRecord> s3EventNotificationRecords) { | ||
| for (Map.Entry<Message, S3EventNotification.S3EventNotificationRecord> entry: s3EventNotificationRecords.entrySet()) { | ||
| if(isEventNameCreated(entry.getValue())) { | ||
| S3ObjectReference s3ObjectReference = populateS3Reference(entry.getValue()); | ||
| s3Service.addS3Object(s3ObjectReference); | ||
| // TODO: delete sqsMessages which are successfully processed | ||
| } | ||
| } | ||
| } | ||
|
|
||
| boolean isEventNameCreated(final S3EventNotification.S3EventNotificationRecord s3EventNotificationRecord) { | ||
| Optional<S3EventNotification.S3EventNotificationRecord> filter = objectCreatedFilter.filter(s3EventNotificationRecord); | ||
| return filter.isPresent(); | ||
| } | ||
|
|
||
| S3ObjectReference populateS3Reference(final S3EventNotification.S3EventNotificationRecord s3EventNotificationRecord) { | ||
| return S3ObjectReference.fromBucketAndKey(s3EventNotificationRecord.getS3().getBucket().getName(), | ||
| s3EventNotificationRecord.getS3().getObject().getKey()); | ||
| } | ||
| } |
Oops, something went wrong.