diff --git a/data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/S3DataSelectionIT.java b/data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/S3DataSelectionIT.java new file mode 100644 index 0000000000..772746b164 --- /dev/null +++ b/data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/S3DataSelectionIT.java @@ -0,0 +1,151 @@ +/* + * Copyright OpenSearch Contributors + * SPDX-License-Identifier: Apache-2.0 + */ + +package org.opensearch.dataprepper.plugins.source.s3; + +import io.micrometer.core.instrument.Counter; +import io.micrometer.core.instrument.DistributionSummary; +import io.micrometer.core.instrument.Meter; +import io.micrometer.core.instrument.Tags; +import io.micrometer.core.instrument.Timer; +import io.micrometer.core.instrument.noop.NoopTimer; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.opensearch.dataprepper.model.buffer.Buffer; +import org.opensearch.dataprepper.model.event.Event; +import org.opensearch.dataprepper.model.record.Record; +import org.opensearch.dataprepper.plugins.codec.CompressionOption; +import org.opensearch.dataprepper.plugins.source.s3.configuration.S3DataSelection; +import org.opensearch.dataprepper.plugins.source.s3.ownership.BucketOwnerProvider; +import software.amazon.awssdk.regions.Region; +import software.amazon.awssdk.services.s3.S3Client; + +import java.io.IOException; +import java.time.Duration; +import java.time.Instant; +import java.util.Collection; +import java.util.Optional; +import java.util.UUID; +import java.util.function.Consumer; + +import static org.hamcrest.CoreMatchers.equalTo; +import static org.hamcrest.CoreMatchers.notNullValue; +import static org.hamcrest.CoreMatchers.nullValue; +import static org.hamcrest.MatcherAssert.assertThat; +import static org.hamcrest.Matchers.greaterThanOrEqualTo; +import static org.mockito.ArgumentMatchers.anyCollection; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +class S3DataSelectionIT { + private static final int TIMEOUT_IN_MILLIS = 200; + + private S3Client s3Client; + private S3ObjectGenerator s3ObjectGenerator; + private String bucket; + private Buffer> buffer; + private S3ObjectPluginMetrics s3ObjectPluginMetrics; + private BucketOwnerProvider bucketOwnerProvider; + private EventMetadataModifier eventMetadataModifier; + + @BeforeEach + void setUp() { + s3Client = S3Client.builder() + .region(Region.of(System.getProperty("tests.s3source.region"))) + .build(); + bucket = System.getProperty("tests.s3source.bucket"); + s3ObjectGenerator = new S3ObjectGenerator(s3Client, bucket); + buffer = mock(Buffer.class); + + eventMetadataModifier = new EventMetadataModifier(S3SourceConfig.DEFAULT_METADATA_ROOT_KEY, false); + bucketOwnerProvider = b -> Optional.empty(); + + s3ObjectPluginMetrics = mock(S3ObjectPluginMetrics.class); + final Counter counter = mock(Counter.class); + final DistributionSummary distributionSummary = mock(DistributionSummary.class); + final Timer timer = new NoopTimer(new Meter.Id("test", Tags.empty(), null, null, Meter.Type.TIMER)); + when(s3ObjectPluginMetrics.getS3ObjectsFailedCounter()).thenReturn(counter); + when(s3ObjectPluginMetrics.getS3ObjectsSucceededCounter()).thenReturn(counter); + when(s3ObjectPluginMetrics.getS3ObjectsFailedAccessDeniedCounter()).thenReturn(counter); + when(s3ObjectPluginMetrics.getS3ObjectsFailedNotFoundCounter()).thenReturn(counter); + when(s3ObjectPluginMetrics.getS3ObjectSizeSummary()).thenReturn(distributionSummary); + when(s3ObjectPluginMetrics.getS3ObjectEventsSummary()).thenReturn(distributionSummary); + when(s3ObjectPluginMetrics.getS3ObjectSizeProcessedSummary()).thenReturn(distributionSummary); + when(s3ObjectPluginMetrics.getS3ObjectReadTimer()).thenReturn(timer); + when(s3ObjectPluginMetrics.getS3ObjectNoRecordsFound()).thenReturn(counter); + } + + @Test + void data_and_metadata_selection_includes_both_data_and_metadata() throws Exception { + testS3DataSelection(S3DataSelection.DATA_AND_METADATA, event -> { + assertThat(event.toMap().containsKey("message"), equalTo(true)); // data content + assertThat(event.get("s3/bucket", String.class), equalTo(bucket)); // metadata present + }); + } + + @Test + void data_only_selection_includes_only_data() throws Exception { + testS3DataSelection(S3DataSelection.DATA_ONLY, event -> { + assertThat(event.toMap().containsKey("message"), equalTo(true)); // data content + assertThat(event.get("s3/bucket", String.class), nullValue()); // no metadata + assertThat(event.get("s3/key", String.class), nullValue()); // no metadata + }); + } + + @Test + void metadata_only_selection_includes_only_metadata() throws Exception { + testS3DataSelection(S3DataSelection.METADATA_ONLY, event -> { + assertThat(event.toMap().containsKey("message"), equalTo(false)); // no data content + assertThat(event.get("bucket", String.class), equalTo(bucket)); // metadata present as top-level key + }); + } + + private void testS3DataSelection(S3DataSelection dataSelection, Consumer eventAssertions) throws Exception { + String key = writeTestObject(); + S3ObjectReference s3ObjectReference = S3ObjectReference.bucketAndKey(bucket, key).build(); + + stubBufferWriter(eventAssertions); + + S3ObjectWorker objectWorker = createObjectUnderTest(); + objectWorker.processS3Object(s3ObjectReference, dataSelection, null, null, null); + + verify(buffer).writeAll(anyCollection(), anyInt()); + } + + private void stubBufferWriter(final Consumer eventAssertions) throws Exception { + doAnswer(a -> { + final Collection> recordsCollection = a.getArgument(0); + assertThat(recordsCollection.size(), greaterThanOrEqualTo(1)); + for (Record eventRecord : recordsCollection) { + assertThat(eventRecord, notNullValue()); + assertThat(eventRecord.getData(), notNullValue()); + eventAssertions.accept(eventRecord.getData()); + } + return null; + }).when(buffer).writeAll(anyCollection(), anyInt()); + } + + private S3ObjectWorker createObjectUnderTest() { + final S3ObjectRequest request = new S3ObjectRequest.Builder(buffer, 100, + Duration.ofMillis(TIMEOUT_IN_MILLIS), s3ObjectPluginMetrics) + .bucketOwnerProvider(bucketOwnerProvider) + .eventConsumer(eventMetadataModifier) + .codec(new NewlineDelimitedRecordsGenerator().getCodec()) + .s3Client(s3Client) + .compressionOption(CompressionOption.NONE) + .build(); + return new S3ObjectWorker(request); + } + + private String writeTestObject() throws IOException { + final NewlineDelimitedRecordsGenerator generator = new NewlineDelimitedRecordsGenerator(); + final String key = "s3 source/data-selection/" + UUID.randomUUID() + "_" + Instant.now().toString() + generator.getFileExtension(); + s3ObjectGenerator.write(5, key, generator, false); + return key; + } +} diff --git a/data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/SqsServiceIT.java b/data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/SqsServiceIT.java index 3d1e6343f0..0e4749b125 100644 --- a/data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/SqsServiceIT.java +++ b/data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/SqsServiceIT.java @@ -18,6 +18,7 @@ import org.mockito.junit.jupiter.MockitoExtension; import org.opensearch.dataprepper.metrics.PluginMetrics; import org.opensearch.dataprepper.model.acknowledgements.AcknowledgementSetManager; +import org.opensearch.dataprepper.plugins.source.s3.configuration.S3DataSelection; import org.opensearch.dataprepper.plugins.source.s3.configuration.NotificationSourceOption; import org.opensearch.dataprepper.plugins.source.s3.configuration.OnErrorOption; import org.opensearch.dataprepper.plugins.source.s3.configuration.AwsAuthenticationOptions; @@ -148,7 +149,7 @@ public void test_sqsService(int numWorkers) throws IOException { assertThat(numMessages.get(), equalTo(numberOfObjectsToWrite)); }); final ArgumentCaptor s3ObjectReferenceArgumentCaptor = ArgumentCaptor.forClass(S3ObjectReference.class); - verify(s3Service, atLeastOnce()).addS3Object(s3ObjectReferenceArgumentCaptor.capture(), eq(null)); + verify(s3Service, atLeastOnce()).addS3Object(s3ObjectReferenceArgumentCaptor.capture(), eq(S3DataSelection.DATA_AND_METADATA), eq(null)); assertThat(s3ObjectReferenceArgumentCaptor.getValue().getBucketName(), equalTo(bucket)); assertThat(s3ObjectReferenceArgumentCaptor.getValue().getKey(), startsWith("s3 source/sqs/")); objectUnderTest.stop(); diff --git a/data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorkerIT.java b/data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorkerIT.java index a74e1b35bb..61362129c1 100644 --- a/data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorkerIT.java +++ b/data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorkerIT.java @@ -26,6 +26,7 @@ import org.opensearch.dataprepper.model.event.JacksonEvent; import org.opensearch.dataprepper.plugins.source.s3.configuration.NotificationSourceOption; import org.opensearch.dataprepper.plugins.source.s3.configuration.OnErrorOption; +import org.opensearch.dataprepper.plugins.source.s3.configuration.S3DataSelection; import org.opensearch.dataprepper.plugins.source.s3.configuration.SqsOptions; import org.opensearch.dataprepper.plugins.source.sqs.common.SqsBackoff; import software.amazon.awssdk.regions.Region; @@ -140,7 +141,7 @@ void processSqsMessages_should_return_at_least_one_message(final int numberOfObj final int sqsMessagesProcessed = objectUnderTest.processSqsMessages(); final ArgumentCaptor s3ObjectReferenceArgumentCaptor = ArgumentCaptor.forClass(S3ObjectReference.class); - verify(s3Service, atLeastOnce()).addS3Object(s3ObjectReferenceArgumentCaptor.capture(), eq(null)); + verify(s3Service, atLeastOnce()).addS3Object(s3ObjectReferenceArgumentCaptor.capture(), eq(S3DataSelection.DATA_AND_METADATA), eq(null)); assertThat(s3ObjectReferenceArgumentCaptor.getValue().getBucketName(), equalTo(bucket)); assertThat(s3ObjectReferenceArgumentCaptor.getValue().getKey(), startsWith("s3 source/sqs/")); @@ -183,7 +184,7 @@ void processSqsMessages_should_return_at_least_one_message_with_acks_with_callba event = (Event)JacksonEvent.fromMessage(val.getArgument(0).toString()); ackSet.add(event); return null; - }).when(s3Service).addS3Object(any(S3ObjectReference.class), any(AcknowledgementSet.class)); + }).when(s3Service).addS3Object(any(S3ObjectReference.class), eq(S3DataSelection.DATA_AND_METADATA), any(AcknowledgementSet.class)); ScheduledExecutorService executor = Executors.newScheduledThreadPool(2); acknowledgementSetManager = new DefaultAcknowledgementSetManager(executor); final SqsWorker objectUnderTest = createObjectUnderTest(); @@ -256,7 +257,7 @@ void processSqsMessages_should_return_at_least_one_message_with_acks_with_callba } catch (Exception e){} return null; - }).when(s3Service).addS3Object(any(S3ObjectReference.class), any(AcknowledgementSet.class)); + }).when(s3Service).addS3Object(any(S3ObjectReference.class), eq(S3DataSelection.DATA_AND_METADATA), any(AcknowledgementSet.class)); ScheduledExecutorService executor = Executors.newScheduledThreadPool(2); acknowledgementSetManager = new DefaultAcknowledgementSetManager(executor); final SqsWorker objectUnderTest = createObjectUnderTest(); @@ -333,7 +334,7 @@ void processSqsMessages_with_acks_and_progress_check_callbacks(final int numberO Thread.sleep(2000); } catch (Exception e) {} return null; - }).when(s3Service).addS3Object(any(S3ObjectReference.class), any(AcknowledgementSet.class)); + }).when(s3Service).addS3Object(any(S3ObjectReference.class), eq(S3DataSelection.DATA_AND_METADATA), any(AcknowledgementSet.class)); ScheduledExecutorService executor = Executors.newScheduledThreadPool(2); when(sqsOptions.getVisibilityTimeout()).thenReturn(Duration.ofSeconds(6)); when(sqsOptions.getVisibilityDuplicateProtectionTimeout()).thenReturn(Duration.ofSeconds(60)); @@ -408,7 +409,7 @@ void processSqsMessages_with_acks_and_progress_check_callbacks_expires(final int Thread.sleep(2000); } catch (Exception e) {} return null; - }).when(s3Service).addS3Object(any(S3ObjectReference.class), any(AcknowledgementSet.class)); + }).when(s3Service).addS3Object(any(S3ObjectReference.class), eq(S3DataSelection.DATA_AND_METADATA), any(AcknowledgementSet.class)); ScheduledExecutorService executor = Executors.newScheduledThreadPool(2); when(sqsOptions.getVisibilityTimeout()).thenReturn(Duration.ofSeconds(6)); when(sqsOptions.getVisibilityDuplicateProtectionTimeout()).thenReturn(Duration.ofSeconds(60)); @@ -439,7 +440,7 @@ void processSqsMessages_should_return_at_least_one_message_when_using_eventbridg final int sqsMessagesProcessed = objectUnderTest.processSqsMessages(); final ArgumentCaptor s3ObjectReferenceArgumentCaptor = ArgumentCaptor.forClass(S3ObjectReference.class); - verify(s3Service, atLeastOnce()).addS3Object(s3ObjectReferenceArgumentCaptor.capture(), eq(null)); + verify(s3Service, atLeastOnce()).addS3Object(s3ObjectReferenceArgumentCaptor.capture(), eq(S3DataSelection.DATA_AND_METADATA), eq(null)); assertThat(s3ObjectReferenceArgumentCaptor.getValue().getBucketName(), equalTo(bucket)); assertThat(s3ObjectReferenceArgumentCaptor.getValue().getKey(), startsWith("s3 source/sqs/")); diff --git a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/S3Service.java b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/S3Service.java index a76d993153..8c3eb7d6f1 100644 --- a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/S3Service.java +++ b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/S3Service.java @@ -16,8 +16,8 @@ public class S3Service { this.s3ObjectHandler = s3ObjectHandler; } - void addS3Object(final S3ObjectReference s3ObjectReference, AcknowledgementSet acknowledgementSet) throws IOException { - s3ObjectHandler.processS3Object(s3ObjectReference, S3DataSelection.DATA_AND_METADATA, acknowledgementSet, null, null); + void addS3Object(final S3ObjectReference s3ObjectReference, final S3DataSelection dataSelection, AcknowledgementSet acknowledgementSet) throws IOException { + s3ObjectHandler.processS3Object(s3ObjectReference, dataSelection, acknowledgementSet, null, null); } void deleteS3Object(final S3ObjectReference s3ObjectReference) throws IOException { diff --git a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/S3SourceConfig.java b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/S3SourceConfig.java index 7cf1568866..280024176d 100644 --- a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/S3SourceConfig.java +++ b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/S3SourceConfig.java @@ -21,6 +21,7 @@ import org.opensearch.dataprepper.plugins.source.s3.configuration.S3ScanScanOptions; import org.opensearch.dataprepper.plugins.source.s3.configuration.S3SelectOptions; import org.opensearch.dataprepper.plugins.source.s3.configuration.SqsOptions; +import org.opensearch.dataprepper.plugins.source.s3.configuration.S3DataSelection; import java.time.Duration; import java.util.Map; @@ -100,6 +101,9 @@ public class S3SourceConfig { @JsonProperty("disable_s3_metadata_in_event") private boolean deleteS3MetadataInEvent = false; + @JsonProperty("data_selection") + private S3DataSelection dataSelection = S3DataSelection.DATA_AND_METADATA; + @AssertTrue(message = "A codec is required for reading objects.") boolean isCodecProvidedWhenNeeded() { if(s3SelectOptions == null) @@ -214,4 +218,8 @@ public String getDefaultBucketOwner() { public boolean isDeleteS3MetadataInEvent() { return deleteS3MetadataInEvent; } + + public S3DataSelection getDataSelection() { + return dataSelection; + } } diff --git a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/ScanObjectWorker.java b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/ScanObjectWorker.java index 625e9e0aa6..6ea8841635 100644 --- a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/ScanObjectWorker.java +++ b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/ScanObjectWorker.java @@ -136,7 +136,7 @@ public ScanObjectWorker(final S3Client s3Client, S3DataSelection dataSelection = s3ScanBucketOption.getDataSelection(); if (scanFilter != null && scanFilter.getS3scanIncludePrefixOptions() != null) { for (final String prefix : scanFilter.getS3scanIncludePrefixOptions()) { - prefixMap.put(prefix, dataSelection == null ? S3DataSelection.DATA_AND_METADATA : dataSelection); + prefixMap.put(prefix, dataSelection == null ? s3SourceConfig.getDataSelection() : dataSelection); } bucketDataSelectionMap.put(bucketName, prefixMap); } diff --git a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorker.java b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorker.java index 1b6f76154b..c9eb7c30f8 100644 --- a/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorker.java +++ b/data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorker.java @@ -350,7 +350,7 @@ private Optional processS3Object( final AcknowledgementSet acknowledgementSet) { // SQS messages won't be deleted if we are unable to process S3Objects because of an exception try { - s3Service.addS3Object(s3ObjectReference, acknowledgementSet); + s3Service.addS3Object(s3ObjectReference, s3SourceConfig.getDataSelection(), acknowledgementSet); return Optional.of(buildDeleteMessageBatchRequestEntry(parsedMessage.getMessage())); } catch (final Exception e) { LOG.error("Error processing from S3: {}. Retrying with exponential backoff.", e.getMessage()); diff --git a/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3ScanObjectWorkerTest.java b/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3ScanObjectWorkerTest.java index aec63e040c..3e0cd2131b 100644 --- a/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3ScanObjectWorkerTest.java +++ b/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3ScanObjectWorkerTest.java @@ -194,6 +194,7 @@ void test_bucketOptions_with_different_data_selections() throws Exception { when(s3SourceConfig.getAcknowledgements()).thenReturn(false); when(s3SourceConfig.isDeleteS3ObjectsOnRead()).thenReturn(false); + when(s3SourceConfig.getDataSelection()).thenReturn(S3DataSelection.DATA_AND_METADATA); final String objectKey = UUID.randomUUID().toString(); final String partitionKey1 = bucket + "|" + "prefix1/"+objectKey; diff --git a/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3ServiceTest.java b/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3ServiceTest.java index fec4d90275..69ba677f24 100644 --- a/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3ServiceTest.java +++ b/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3ServiceTest.java @@ -25,17 +25,18 @@ private S3Service createObjectUnderTest() { } @Test - void addS3Object_calls_parseS3Object_on_S3ObjectHandler() throws IOException { + void addS3Object_with_dataSelection_calls_processS3Object_on_S3ObjectHandler() throws IOException { final AcknowledgementSet acknowledgementSet = mock(AcknowledgementSet.class); final S3ObjectReference s3ObjectReference = mock(S3ObjectReference.class); + final S3DataSelection dataSelection = S3DataSelection.DATA_ONLY; - doNothing().when(s3ObjectHandler).processS3Object(eq(s3ObjectReference), eq(S3DataSelection.DATA_AND_METADATA), eq(acknowledgementSet), eq(null), eq(null)); + doNothing().when(s3ObjectHandler).processS3Object(eq(s3ObjectReference), eq(dataSelection), eq(acknowledgementSet), eq(null), eq(null)); final S3Service objectUnderTest = createObjectUnderTest(); - objectUnderTest.addS3Object(s3ObjectReference, acknowledgementSet); + objectUnderTest.addS3Object(s3ObjectReference, dataSelection, acknowledgementSet); - verify(s3ObjectHandler).processS3Object(s3ObjectReference, S3DataSelection.DATA_AND_METADATA, acknowledgementSet, null, null); + verify(s3ObjectHandler).processS3Object(s3ObjectReference, dataSelection, acknowledgementSet, null, null); } @Test diff --git a/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3SourceConfigTest.java b/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3SourceConfigTest.java index b4676410bc..ae42a74478 100644 --- a/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3SourceConfigTest.java +++ b/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3SourceConfigTest.java @@ -13,6 +13,7 @@ import org.opensearch.dataprepper.plugins.source.s3.configuration.FolderPartitioningOptions; import org.opensearch.dataprepper.plugins.source.s3.configuration.S3ScanScanOptions; import org.opensearch.dataprepper.plugins.source.s3.configuration.S3SelectOptions; +import org.opensearch.dataprepper.plugins.source.s3.configuration.S3DataSelection; import org.opensearch.dataprepper.test.helper.ReflectivelySetField; import java.util.HashMap; @@ -125,4 +126,17 @@ void s3Select_is_not_supported_with_deleteS3ObjectsOnRead() throws NoSuchFieldEx assertFalse(s3SourceConfig.isS3SelectNotUsingDeleteS3ObjectsOnRead()); } + + @Test + void getDataSelection_returns_default_value() { + final S3SourceConfig s3SourceConfig = new S3SourceConfig(); + assertThat(s3SourceConfig.getDataSelection(), equalTo(S3DataSelection.DATA_AND_METADATA)); + } + + @Test + void getDataSelection_returns_configured_value() throws NoSuchFieldException, IllegalAccessException { + final S3SourceConfig s3SourceConfig = new S3SourceConfig(); + ReflectivelySetField.setField(S3SourceConfig.class, s3SourceConfig, "dataSelection", S3DataSelection.METADATA_ONLY); + assertThat(s3SourceConfig.getDataSelection(), equalTo(S3DataSelection.METADATA_ONLY)); + } } diff --git a/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorkerTest.java b/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorkerTest.java index f96ace5ca2..e62ec71bc2 100644 --- a/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorkerTest.java +++ b/data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorkerTest.java @@ -29,6 +29,7 @@ import org.opensearch.dataprepper.model.acknowledgements.ProgressCheck; import org.opensearch.dataprepper.plugins.source.s3.configuration.NotificationSourceOption; import org.opensearch.dataprepper.plugins.source.s3.configuration.OnErrorOption; +import org.opensearch.dataprepper.plugins.source.s3.configuration.S3DataSelection; import org.opensearch.dataprepper.plugins.source.s3.configuration.SqsOptions; import org.opensearch.dataprepper.plugins.source.s3.exception.SqsRetriesExhaustedException; import org.opensearch.dataprepper.plugins.source.s3.filter.S3EventFilter; @@ -70,6 +71,7 @@ import static org.junit.jupiter.params.provider.Arguments.arguments; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.inOrder; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; @@ -174,6 +176,7 @@ void setUp() { @ParameterizedTest @ValueSource(strings = {"ObjectCreated:Put", "ObjectCreated:Post", "ObjectCreated:Copy", "ObjectCreated:CompleteMultipartUpload"}) void processSqsMessages_should_return_number_of_messages_processed(final String eventName) throws IOException { + when(s3SourceConfig.getDataSelection()).thenReturn(S3DataSelection.DATA_AND_METADATA); Instant startTime = Instant.now().minus(1, ChronoUnit.HOURS); final Message message = mock(Message.class); when(message.body()).thenReturn(createEventNotification(eventName, startTime)); @@ -200,7 +203,7 @@ void processSqsMessages_should_return_number_of_messages_processed(final String assertThat(actualDeleteMessageBatchRequest.entries().get(0).id(), equalTo(message.messageId())); assertThat(actualDeleteMessageBatchRequest.entries().get(0).receiptHandle(), equalTo(message.receiptHandle())); assertThat(messagesProcessed, equalTo(1)); - verify(s3Service).addS3Object(any(S3ObjectReference.class), any()); + verify(s3Service).addS3Object(any(S3ObjectReference.class), eq(S3DataSelection.DATA_AND_METADATA), any()); verify(sqsClient).deleteMessageBatch(any(DeleteMessageBatchRequest.class)); verify(sqsMessagesReceivedCounter).increment(1); verify(sqsMessagesDeletedCounter).increment(1); @@ -365,6 +368,7 @@ void processSqsMessages_should_invoke_delete_if_input_is_not_valid_JSON_and_dele @Test void processSqsMessages_should_return_number_of_messages_processed_when_using_EventBridge() throws IOException { + when(s3SourceConfig.getDataSelection()).thenReturn(S3DataSelection.DATA_AND_METADATA); when(s3SourceConfig.getNotificationSource()).thenReturn(NotificationSourceOption.EVENTBRIDGE); Instant startTime = Instant.now().minus(1, ChronoUnit.HOURS); final Message message = mock(Message.class); @@ -392,7 +396,7 @@ void processSqsMessages_should_return_number_of_messages_processed_when_using_Ev assertThat(actualDeleteMessageBatchRequest.entries().get(0).id(), equalTo(message.messageId())); assertThat(actualDeleteMessageBatchRequest.entries().get(0).receiptHandle(), equalTo(message.receiptHandle())); assertThat(messagesProcessed, equalTo(1)); - verify(s3Service).addS3Object(any(S3ObjectReference.class), any()); + verify(s3Service).addS3Object(any(S3ObjectReference.class), eq(S3DataSelection.DATA_AND_METADATA), any()); verify(sqsClient).deleteMessageBatch(any(DeleteMessageBatchRequest.class)); verify(sqsMessagesReceivedCounter).increment(1); verify(sqsMessagesDeletedCounter).increment(1); @@ -402,6 +406,7 @@ void processSqsMessages_should_return_number_of_messages_processed_when_using_Ev @Test void processSqsMessages_should_return_number_of_messages_processed_when_using_SecurityLake() throws IOException { + when(s3SourceConfig.getDataSelection()).thenReturn(S3DataSelection.DATA_AND_METADATA); when(s3SourceConfig.getNotificationSource()).thenReturn(NotificationSourceOption.EVENTBRIDGE); Instant startTime = Instant.now().minus(1, ChronoUnit.HOURS); final Message message = mock(Message.class); @@ -429,7 +434,7 @@ void processSqsMessages_should_return_number_of_messages_processed_when_using_Se assertThat(actualDeleteMessageBatchRequest.entries().get(0).id(), equalTo(message.messageId())); assertThat(actualDeleteMessageBatchRequest.entries().get(0).receiptHandle(), equalTo(message.receiptHandle())); assertThat(messagesProcessed, equalTo(1)); - verify(s3Service).addS3Object(any(S3ObjectReference.class), any()); + verify(s3Service).addS3Object(any(S3ObjectReference.class), eq(S3DataSelection.DATA_AND_METADATA), any()); verify(sqsClient).deleteMessageBatch(any(DeleteMessageBatchRequest.class)); verify(sqsMessagesReceivedCounter).increment(1); verify(sqsMessagesDeletedCounter).increment(1); @@ -441,6 +446,7 @@ void processSqsMessages_should_return_number_of_messages_processed_when_using_Se @Test void processSqsMessages_should_report_correct_metrics_for_DeleteMessages_when_some_succeed_and_some_fail() throws IOException { + when(s3SourceConfig.getDataSelection()).thenReturn(S3DataSelection.DATA_AND_METADATA); Instant startTime = Instant.now().minus(1, ChronoUnit.HOURS); final List messages = IntStream.range(0, 6).mapToObj(i -> { final Message message = mock(Message.class); @@ -479,7 +485,7 @@ void processSqsMessages_should_report_correct_metrics_for_DeleteMessages_when_so assertThat(actualDeleteMessageBatchRequest, notNullValue()); assertThat(actualDeleteMessageBatchRequest.entries().size(), equalTo(6)); assertThat(actualDeleteMessageBatchRequest.queueUrl(), equalTo(s3SourceConfig.getSqsOptions().getSqsUrl())); - verify(s3Service, times(6)).addS3Object(any(S3ObjectReference.class), any()); + verify(s3Service, times(6)).addS3Object(any(S3ObjectReference.class), eq(S3DataSelection.DATA_AND_METADATA), any()); verify(sqsClient).deleteMessageBatch(any(DeleteMessageBatchRequest.class)); verify(sqsMessagesReceivedCounter).increment(6); verify(sqsMessagesDeletedCounter).increment(3); @@ -492,6 +498,7 @@ void processSqsMessages_should_report_correct_metrics_for_DeleteMessages_when_so @ArgumentsSource(DeleteExceptionToBackoffCalls.class) void processSqsMessages_should_report_correct_metrics_for_DeleteMessages_when_request_fails( final Class exClass, final int timesToCallBackoff) throws IOException { + when(s3SourceConfig.getDataSelection()).thenReturn(S3DataSelection.DATA_AND_METADATA); Instant startTime = Instant.now().minus(1, ChronoUnit.HOURS); final List messages = IntStream.range(0, 6).mapToObj(i -> { final Message message = mock(Message.class); @@ -519,7 +526,7 @@ void processSqsMessages_should_report_correct_metrics_for_DeleteMessages_when_re assertThat(actualDeleteMessageBatchRequest, notNullValue()); assertThat(actualDeleteMessageBatchRequest.entries().size(), equalTo(6)); assertThat(actualDeleteMessageBatchRequest.queueUrl(), equalTo(s3SourceConfig.getSqsOptions().getSqsUrl())); - verify(s3Service, times(6)).addS3Object(any(S3ObjectReference.class), any()); + verify(s3Service, times(6)).addS3Object(any(S3ObjectReference.class), eq(S3DataSelection.DATA_AND_METADATA), any()); verify(sqsClient).deleteMessageBatch(any(DeleteMessageBatchRequest.class)); verify(sqsMessagesReceivedCounter).increment(6); verifyNoInteractions(sqsMessagesDeletedCounter); @@ -601,6 +608,7 @@ void populateS3Reference_should_interact_with_getUrlDecodedKey() throws NoSuchMe @ParameterizedTest @ValueSource(strings = {"ObjectCreated:Put", "ObjectCreated:Post", "ObjectCreated:Copy", "ObjectCreated:CompleteMultipartUpload"}) void processSqsMessages_should_return_number_of_messages_processed_with_acknowledgements(final String eventName) throws IOException { + when(s3SourceConfig.getDataSelection()).thenReturn(S3DataSelection.DATA_AND_METADATA); when(acknowledgementSetManager.create(any(), any(Duration.class))).thenReturn(acknowledgementSet); when(s3SourceConfig.getAcknowledgements()).thenReturn(true); Instant startTime = Instant.now().minus(1, ChronoUnit.HOURS); @@ -622,7 +630,7 @@ void processSqsMessages_should_return_number_of_messages_processed_with_acknowle Duration actualDelay = durationArgumentCaptor.getValue(); assertThat(messagesProcessed, equalTo(1)); - verify(s3Service).addS3Object(any(S3ObjectReference.class), any()); + verify(s3Service).addS3Object(any(S3ObjectReference.class), eq(S3DataSelection.DATA_AND_METADATA), any()); verify(acknowledgementSetManager).create(any(), any(Duration.class)); verify(sqsMessagesReceivedCounter).increment(1); verifyNoInteractions(sqsMessagesDeletedCounter); @@ -633,6 +641,7 @@ void processSqsMessages_should_return_number_of_messages_processed_with_acknowle @ParameterizedTest @ValueSource(strings = {"ObjectCreated:Put", "ObjectCreated:Post", "ObjectCreated:Copy", "ObjectCreated:CompleteMultipartUpload"}) void processSqsMessages_should_call_s3Service_to_delete_objects_if_deleteS3ObjectsOnRead_is_enabled_on_acknowledgment_callback(final String eventName) throws IOException { + when(s3SourceConfig.getDataSelection()).thenReturn(S3DataSelection.DATA_AND_METADATA); when(s3SourceConfig.getAcknowledgements()).thenReturn(true); when(s3SourceConfig.isDeleteS3ObjectsOnRead()).thenReturn(true); Instant startTime = Instant.now().minus(1, ChronoUnit.HOURS); @@ -658,7 +667,7 @@ void processSqsMessages_should_call_s3Service_to_delete_objects_if_deleteS3Objec Duration actualDelay = durationArgumentCaptor.getValue(); assertThat(messagesProcessed, equalTo(1)); - verify(s3Service).addS3Object(any(S3ObjectReference.class), any()); + verify(s3Service).addS3Object(any(S3ObjectReference.class), eq(S3DataSelection.DATA_AND_METADATA), any()); verify(acknowledgementSetManager).create(ackCallbackArgumentCaptor.capture(), any(Duration.class)); @@ -682,6 +691,7 @@ void processSqsMessages_should_call_s3Service_to_delete_objects_if_deleteS3Objec @ParameterizedTest @ValueSource(strings = {"ObjectCreated:Put", "ObjectCreated:Post", "ObjectCreated:Copy", "ObjectCreated:CompleteMultipartUpload"}) void processSqsMessages_should_return_number_of_messages_processed_with_acknowledgements_and_progress_check(final String eventName) throws IOException { + when(s3SourceConfig.getDataSelection()).thenReturn(S3DataSelection.DATA_AND_METADATA); when(sqsOptions.getVisibilityDuplicateProtection()).thenReturn(true); when(sqsOptions.getVisibilityTimeout()).thenReturn(Duration.ofSeconds(6)); when(acknowledgementSetManager.create(any(), any(Duration.class))).thenReturn(acknowledgementSet); @@ -705,7 +715,7 @@ void processSqsMessages_should_return_number_of_messages_processed_with_acknowle Duration actualDelay = durationArgumentCaptor.getValue(); assertThat(messagesProcessed, equalTo(1)); - verify(s3Service).addS3Object(any(S3ObjectReference.class), any()); + verify(s3Service).addS3Object(any(S3ObjectReference.class), eq(S3DataSelection.DATA_AND_METADATA), any()); verify(acknowledgementSetManager).create(any(), any(Duration.class)); verify(acknowledgementSet).addProgressCheck(any(), any(Duration.class)); verify(sqsMessagesReceivedCounter).increment(1); @@ -736,6 +746,7 @@ void processSqsMessages_should_not_interact_with_S3Service_if_input_is_not_valid @Test void processSqsMessages_should_update_visibility_timeout_when_progress_changes() throws IOException { + when(s3SourceConfig.getDataSelection()).thenReturn(S3DataSelection.DATA_AND_METADATA); when(sqsOptions.getVisibilityDuplicateProtection()).thenReturn(true); when(sqsOptions.getVisibilityTimeout()).thenReturn(Duration.ofMillis(1)); when(acknowledgementSetManager.create(any(), any(Duration.class))).thenReturn(acknowledgementSet); @@ -754,7 +765,7 @@ void processSqsMessages_should_update_visibility_timeout_when_progress_changes() final int messagesProcessed = createObjectUnderTest().processSqsMessages(); assertThat(messagesProcessed, equalTo(1)); - verify(s3Service).addS3Object(any(S3ObjectReference.class), any()); + verify(s3Service).addS3Object(any(S3ObjectReference.class), eq(S3DataSelection.DATA_AND_METADATA), any()); verify(acknowledgementSetManager).create(any(), any(Duration.class)); ArgumentCaptor> progressConsumerArgumentCaptor = ArgumentCaptor.forClass(Consumer.class); @@ -774,6 +785,7 @@ void processSqsMessages_should_update_visibility_timeout_when_progress_changes() @Test void processSqsMessages_should_stop_updating_visibility_timeout_after_stop() throws IOException { + when(s3SourceConfig.getDataSelection()).thenReturn(S3DataSelection.DATA_AND_METADATA); when(sqsOptions.getVisibilityDuplicateProtection()).thenReturn(true); when(sqsOptions.getVisibilityTimeout()).thenReturn(Duration.ofMillis(1)); when(acknowledgementSetManager.create(any(), any(Duration.class))).thenReturn(acknowledgementSet); @@ -797,7 +809,7 @@ void processSqsMessages_should_stop_updating_visibility_timeout_after_stop() thr final InOrder inOrder = inOrder(s3Service, sqsMessageDelayTimer); inOrder.verify(sqsMessageDelayTimer).record(any(Duration.class)); - inOrder.verify(s3Service).addS3Object(any(S3ObjectReference.class), any()); + inOrder.verify(s3Service).addS3Object(any(S3ObjectReference.class), eq(S3DataSelection.DATA_AND_METADATA), any()); verify(acknowledgementSetManager).create(any(), any(Duration.class)); ArgumentCaptor> progressConsumerArgumentCaptor = ArgumentCaptor.forClass(Consumer.class);