Added data_selection support to S3 SQS source

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<Record<Event>> 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<Event> 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<Event> eventAssertions) throws Exception {
+        doAnswer(a -> {
+            final Collection<Record<Event>> recordsCollection = a.getArgument(0);
+            assertThat(recordsCollection.size(), greaterThanOrEqualTo(1));
+            for (Record<Event> 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;
+    }
+}

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<S3ObjectReference> 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();

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<S3ObjectReference> 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<S3ObjectReference> 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/"));

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 {

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

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

data-prepper-plugins/s3-source/src/main/java/org/opensearch/dataprepper/plugins/source/s3/SqsWorker.java

@@ -350,7 +350,7 @@ private Optional<DeleteMessageBatchRequestEntry> 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());

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;

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

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