Skip to content

Commit eaa6345

Browse files
mzurita-amzmzurita-amz
committed
Added data_selection support to S3 SQS source
Signed-off-by: Manuel Mangas Zurita <mzurita@amazon.com>
1 parent a4db6da commit eaa6345

8 files changed

Lines changed: 201 additions & 12 deletions

File tree

‎data-prepper-plugins/s3-source/src/integrationTest/java/org/opensearch/dataprepper/plugins/source/s3/S3DataSelectionIT.java‎

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.dataprepper.plugins.source.s3;
7+
8+
import io.micrometer.core.instrument.Counter;
9+
import io.micrometer.core.instrument.DistributionSummary;
10+
import io.micrometer.core.instrument.Meter;
11+
import io.micrometer.core.instrument.Tags;
12+
import io.micrometer.core.instrument.Timer;
13+
import io.micrometer.core.instrument.noop.NoopTimer;
14+
import org.junit.jupiter.api.BeforeEach;
15+
import org.junit.jupiter.api.Test;
16+
import org.opensearch.dataprepper.model.buffer.Buffer;
17+
import org.opensearch.dataprepper.model.event.Event;
18+
import org.opensearch.dataprepper.model.record.Record;
19+
import org.opensearch.dataprepper.plugins.codec.CompressionOption;
20+
import org.opensearch.dataprepper.plugins.source.s3.configuration.S3DataSelection;
21+
import org.opensearch.dataprepper.plugins.source.s3.ownership.BucketOwnerProvider;
22+
import software.amazon.awssdk.regions.Region;
23+
import software.amazon.awssdk.services.s3.S3Client;
24+
25+
import java.io.IOException;
26+
import java.time.Duration;
27+
import java.time.Instant;
28+
import java.util.Collection;
29+
import java.util.Optional;
30+
import java.util.UUID;
31+
import java.util.function.Consumer;
32+
33+
import static org.hamcrest.CoreMatchers.equalTo;
34+
import static org.hamcrest.CoreMatchers.notNullValue;
35+
import static org.hamcrest.MatcherAssert.assertThat;
36+
import static org.hamcrest.Matchers.greaterThanOrEqualTo;
37+
import static org.mockito.ArgumentMatchers.anyCollection;
38+
import static org.mockito.ArgumentMatchers.anyInt;
39+
import static org.mockito.Mockito.doAnswer;
40+
import static org.mockito.Mockito.mock;
41+
import static org.mockito.Mockito.verify;
42+
import static org.mockito.Mockito.when;
43+
44+
class S3DataSelectionIT {
45+
private static final int TIMEOUT_IN_MILLIS = 200;
46+
47+
private S3Client s3Client;
48+
private S3ObjectGenerator s3ObjectGenerator;
49+
private String bucket;
50+
private Buffer<Record<Event>> buffer;
51+
private S3ObjectPluginMetrics s3ObjectPluginMetrics;
52+
private BucketOwnerProvider bucketOwnerProvider;
53+
private EventMetadataModifier eventMetadataModifier;
54+
55+
@BeforeEach
56+
void setUp() {
57+
s3Client = S3Client.builder()
58+
.region(Region.of(System.getProperty("tests.s3source.region")))
59+
.build();
60+
bucket = System.getProperty("tests.s3source.bucket");
61+
s3ObjectGenerator = new S3ObjectGenerator(s3Client, bucket);
62+
buffer = mock(Buffer.class);
63+
64+
eventMetadataModifier = new EventMetadataModifier(S3SourceConfig.DEFAULT_METADATA_ROOT_KEY, false);
65+
bucketOwnerProvider = b -> Optional.empty();
66+
67+
s3ObjectPluginMetrics = mock(S3ObjectPluginMetrics.class);
68+
final Counter counter = mock(Counter.class);
69+
final DistributionSummary distributionSummary = mock(DistributionSummary.class);
70+
final Timer timer = new NoopTimer(new Meter.Id("test", Tags.empty(), null, null, Meter.Type.TIMER));
71+
when(s3ObjectPluginMetrics.getS3ObjectsFailedCounter()).thenReturn(counter);
72+
when(s3ObjectPluginMetrics.getS3ObjectsSucceededCounter()).thenReturn(counter);
73+
when(s3ObjectPluginMetrics.getS3ObjectsFailedAccessDeniedCounter()).thenReturn(counter);
74+
when(s3ObjectPluginMetrics.getS3ObjectsFailedNotFoundCounter()).thenReturn(counter);
75+
when(s3ObjectPluginMetrics.getS3ObjectSizeSummary()).thenReturn(distributionSummary);
76+
when(s3ObjectPluginMetrics.getS3ObjectEventsSummary()).thenReturn(distributionSummary);
77+
when(s3ObjectPluginMetrics.getS3ObjectSizeProcessedSummary()).thenReturn(distributionSummary);
78+
when(s3ObjectPluginMetrics.getS3ObjectReadTimer()).thenReturn(timer);
79+
when(s3ObjectPluginMetrics.getS3ObjectNoRecordsFound()).thenReturn(counter);
80+
}
81+
82+
@Test
83+
void data_and_metadata_selection_includes_both_data_and_metadata() throws Exception {
84+
testS3DataSelection(S3DataSelection.DATA_AND_METADATA, event -> {
85+
assertThat(event.toMap().containsKey("message"), equalTo(true)); // data content
86+
assertThat(event.get("s3/bucket", String.class), equalTo(bucket)); // metadata present
87+
});
88+
}
89+
90+
@Test
91+
void data_only_selection_includes_only_data() throws Exception {
92+
testS3DataSelection(S3DataSelection.DATA_ONLY, event -> {
93+
assertThat(event.toMap().containsKey("message"), equalTo(true)); // data content
94+
assertThat(event.get("s3/bucket", String.class) == null, equalTo(true)); // no metadata
95+
assertThat(event.get("s3/key", String.class) == null, equalTo(true)); // no metadata
96+
});
97+
}
98+
99+
@Test
100+
void metadata_only_selection_includes_only_metadata() throws Exception {
101+
testS3DataSelection(S3DataSelection.METADATA_ONLY, event -> {
102+
assertThat(event.toMap().containsKey("message"), equalTo(false)); // no data content
103+
assertThat(event.get("bucket", String.class), equalTo(bucket)); // metadata present as top-level key
104+
});
105+
}
106+
107+
private void testS3DataSelection(S3DataSelection dataSelection, Consumer<Event> eventAssertions) throws Exception {
108+
String key = writeTestObject();
109+
S3ObjectReference s3ObjectReference = S3ObjectReference.bucketAndKey(bucket, key).build();
110+
111+
stubBufferWriter(eventAssertions);
112+
113+
S3ObjectWorker objectWorker = createObjectUnderTest();
114+
objectWorker.processS3Object(s3ObjectReference, dataSelection, null, null, null);
115+
116+
verify(buffer).writeAll(anyCollection(), anyInt());
117+
}
118+
119+
private void stubBufferWriter(final Consumer<Event> eventAssertions) throws Exception {
120+
doAnswer(a -> {
121+
final Collection<Record<Event>> recordsCollection = a.getArgument(0);
122+
assertThat(recordsCollection.size(), greaterThanOrEqualTo(1));
123+
for (Record<Event> eventRecord : recordsCollection) {
124+
assertThat(eventRecord, notNullValue());
125+
assertThat(eventRecord.getData(), notNullValue());
126+
eventAssertions.accept(eventRecord.getData());
127+
}
128+
return null;
129+
}).when(buffer).writeAll(anyCollection(), anyInt());
130+
}
131+
132+
private S3ObjectWorker createObjectUnderTest() {
133+
final S3ObjectRequest request = new S3ObjectRequest.Builder(buffer, 100,
134+
Duration.ofMillis(TIMEOUT_IN_MILLIS), s3ObjectPluginMetrics)
135+
.bucketOwnerProvider(bucketOwnerProvider)
136+
.eventConsumer(eventMetadataModifier)
137+
.codec(new NewlineDelimitedRecordsGenerator().getCodec())
138+
.s3Client(s3Client)
139+
.compressionOption(CompressionOption.NONE)
140+
.build();
141+
return new S3ObjectWorker(request);
142+
}
143+
144+
private String writeTestObject() throws IOException {
145+
final NewlineDelimitedRecordsGenerator generator = new NewlineDelimitedRecordsGenerator();
146+
final String key = "s3 source/data-selection/" + UUID.randomUUID() + "_" + Instant.now().toString() + generator.getFileExtension();
147+
s3ObjectGenerator.write(5, key, generator, false);
148+
return key;
149+
}
150+
}

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

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,10 @@ void addS3Object(final S3ObjectReference s3ObjectReference, AcknowledgementSet a
2020
s3ObjectHandler.processS3Object(s3ObjectReference, S3DataSelection.DATA_AND_METADATA, acknowledgementSet, null, null);
2121
}
2222

23+
void addS3Object(final S3ObjectReference s3ObjectReference, final S3DataSelection dataSelection, AcknowledgementSet acknowledgementSet) throws IOException {
24+
s3ObjectHandler.processS3Object(s3ObjectReference, dataSelection, acknowledgementSet, null, null);
25+
}
26+
2327
void deleteS3Object(final S3ObjectReference s3ObjectReference) throws IOException {
2428
s3ObjectHandler.deleteS3Object(s3ObjectReference);
2529
}

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

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import org.opensearch.dataprepper.plugins.source.s3.configuration.S3ScanScanOptions;
2222
import org.opensearch.dataprepper.plugins.source.s3.configuration.S3SelectOptions;
2323
import org.opensearch.dataprepper.plugins.source.s3.configuration.SqsOptions;
24+
import org.opensearch.dataprepper.plugins.source.s3.configuration.S3DataSelection;
2425

2526
import java.time.Duration;
2627
import java.util.Map;
@@ -100,6 +101,9 @@ public class S3SourceConfig {
100101
@JsonProperty("disable_s3_metadata_in_event")
101102
private boolean deleteS3MetadataInEvent = false;
102103

104+
@JsonProperty("data_selection")
105+
private S3DataSelection dataSelection = S3DataSelection.DATA_AND_METADATA;
106+
103107
@AssertTrue(message = "A codec is required for reading objects.")
104108
boolean isCodecProvidedWhenNeeded() {
105109
if(s3SelectOptions == null)
@@ -214,4 +218,8 @@ public String getDefaultBucketOwner() {
214218
public boolean isDeleteS3MetadataInEvent() {
215219
return deleteS3MetadataInEvent;
216220
}
221+
222+
public S3DataSelection getDataSelection() {
223+
return dataSelection;
224+
}
217225
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -136,7 +136,7 @@ public ScanObjectWorker(final S3Client s3Client,
136136
S3DataSelection dataSelection = s3ScanBucketOption.getDataSelection();
137137
if (scanFilter != null && scanFilter.getS3scanIncludePrefixOptions() != null) {
138138
for (final String prefix : scanFilter.getS3scanIncludePrefixOptions()) {
139-
prefixMap.put(prefix, dataSelection == null ? S3DataSelection.DATA_AND_METADATA : dataSelection);
139+
prefixMap.put(prefix, dataSelection == null ? s3SourceConfig.getDataSelection() : dataSelection);
140140
}
141141
bucketDataSelectionMap.put(bucketName, prefixMap);
142142
}

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

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -350,7 +350,7 @@ private Optional<DeleteMessageBatchRequestEntry> processS3Object(
350350
final AcknowledgementSet acknowledgementSet) {
351351
// SQS messages won't be deleted if we are unable to process S3Objects because of an exception
352352
try {
353-
s3Service.addS3Object(s3ObjectReference, acknowledgementSet);
353+
s3Service.addS3Object(s3ObjectReference, s3SourceConfig.getDataSelection(), acknowledgementSet);
354354
return Optional.of(buildDeleteMessageBatchRequestEntry(parsedMessage.getMessage()));
355355
} catch (final Exception e) {
356356
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‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -194,6 +194,7 @@ void test_bucketOptions_with_different_data_selections() throws Exception {
194194

195195
when(s3SourceConfig.getAcknowledgements()).thenReturn(false);
196196
when(s3SourceConfig.isDeleteS3ObjectsOnRead()).thenReturn(false);
197+
when(s3SourceConfig.getDataSelection()).thenReturn(S3DataSelection.DATA_AND_METADATA);
197198

198199
final String objectKey = UUID.randomUUID().toString();
199200
final String partitionKey1 = bucket + "|" + "prefix1/"+objectKey;

‎data-prepper-plugins/s3-source/src/test/java/org/opensearch/dataprepper/plugins/source/s3/S3SourceConfigTest.java‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@
1313
import org.opensearch.dataprepper.plugins.source.s3.configuration.FolderPartitioningOptions;
1414
import org.opensearch.dataprepper.plugins.source.s3.configuration.S3ScanScanOptions;
1515
import org.opensearch.dataprepper.plugins.source.s3.configuration.S3SelectOptions;
16+
import org.opensearch.dataprepper.plugins.source.s3.configuration.S3DataSelection;
1617
import org.opensearch.dataprepper.test.helper.ReflectivelySetField;
1718

1819
import java.util.HashMap;
@@ -125,4 +126,17 @@ void s3Select_is_not_supported_with_deleteS3ObjectsOnRead() throws NoSuchFieldEx
125126

126127
assertFalse(s3SourceConfig.isS3SelectNotUsingDeleteS3ObjectsOnRead());
127128
}
129+
130+
@Test
131+
void getDataSelection_returns_default_value() {
132+
final S3SourceConfig s3SourceConfig = new S3SourceConfig();
133+
assertThat(s3SourceConfig.getDataSelection(), equalTo(S3DataSelection.DATA_AND_METADATA));
134+
}
135+
136+
@Test
137+
void getDataSelection_returns_configured_value() throws NoSuchFieldException, IllegalAccessException {
138+
final S3SourceConfig s3SourceConfig = new S3SourceConfig();
139+
ReflectivelySetField.setField(S3SourceConfig.class, s3SourceConfig, "dataSelection", S3DataSelection.METADATA_ONLY);
140+
assertThat(s3SourceConfig.getDataSelection(), equalTo(S3DataSelection.METADATA_ONLY));
141+
}
128142
}

0 commit comments

Comments
 (0)