Add selective merge to Event API and refactor S3EnrichProcessor

Signed-off-by: Xun Zhang <xunzh@amazon.com>

Author: Zhangxunmt

.gitignore

@@ -13,6 +13,9 @@ data-prepper-main/src/test/resources/logstash-conf/logstash-filter.yaml
 # output folder created when we run test cases
 **/out/
 
+# Eclipse/IDE compiled output
+**/bin/
+
 # Development tools
 .DS_Store
 .idea

data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/Event.java

@@ -146,6 +146,19 @@ public interface Event extends Serializable {
      */
     void merge(Event other);
 
+    /**
+     * Merges only the specified keys from another Event into the current Event.
+     * Only the keys present in {@code keys} will be copied from {@code other} into this Event.
+     * Values from {@code other} will overwrite existing values in this Event for matching keys.
+     *
+     * @param other the other Event to merge from
+     * @param keys  the list of keys to selectively merge
+     * @throws IllegalArgumentException if the input event is not compatible to merge.
+     * @throws UnsupportedOperationException if the current Event does not support merging.
+     * @since 2.11
+     */
+    void merge(Event other, List<String> keys);
+
     /**
      * Generates a serialized Json string of the entire Event
      *

data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/JacksonEvent.java

@@ -446,6 +446,26 @@ public void merge(final Event other) {
         ((ObjectNode) jsonNode).setAll(otherObjectNode);
     }
 
+    @Override
+    public void merge(final Event other, final List<String> keys) {
+        if (keys == null || keys.isEmpty()) {
+            throw new IllegalArgumentException("Keys list must not be null or empty for selective merge.");
+        }
+        if (!(other instanceof JacksonEvent)) {
+            throw new IllegalArgumentException("Unable to merge the Event. The input Event must be a JacksonEvent.");
+        }
+        if (!(jsonNode instanceof ObjectNode)) {
+            throw new UnsupportedOperationException("Unable to merge the Event. The current Event must have object data.");
+        }
+
+        for (final String key : keys) {
+            final Object value = other.get(key, Object.class);
+            if (value != null) {
+                put(key, value);
+            }
+        }
+    }
+
     @Override
     public String toJsonString() {
         return jsonNode.toString();

data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/JacksonEventTest.java

@@ -682,6 +682,71 @@ void merge_overrides_existing_values() {
         assertThat(event.get("info/ids/id", String.class), equalTo("idx"));
     }
 
+    @Test
+    void merge_with_keys_only_copies_specified_keys() {
+        final String jsonString = "{\"a\": \"alpha\", \"b\": \"beta\", \"c\": \"gamma\"}";
+        Event otherEvent = JacksonEvent.builder().withEventType(EventType.DOCUMENT.toString()).withData(jsonString).build();
+        event.merge(otherEvent, List.of("a", "b"));
+
+        assertThat(event.get("a", Object.class), equalTo("alpha"));
+        assertThat(event.get("b", Object.class), equalTo("beta"));
+        assertThat(event.containsKey("c"), equalTo(false));
+    }
+
+    @Test
+    void merge_with_keys_overwrites_existing_values() {
+        event.put("a", "original");
+        final String jsonString = "{\"a\": \"updated\", \"b\": \"beta\"}";
+        Event otherEvent = JacksonEvent.builder().withEventType(EventType.DOCUMENT.toString()).withData(jsonString).build();
+        event.merge(otherEvent, List.of("a"));
+
+        assertThat(event.get("a", Object.class), equalTo("updated"));
+        assertThat(event.containsKey("b"), equalTo(false));
+    }
+
+    @Test
+    void merge_with_keys_skips_missing_keys_in_other() {
+        event.put("existing", "value");
+        final String jsonString = "{\"a\": \"alpha\"}";
+        Event otherEvent = JacksonEvent.builder().withEventType(EventType.DOCUMENT.toString()).withData(jsonString).build();
+        event.merge(otherEvent, List.of("a", "nonexistent"));
+
+        assertThat(event.get("a", Object.class), equalTo("alpha"));
+        assertThat(event.get("existing", Object.class), equalTo("value"));
+        assertThat(event.containsKey("nonexistent"), equalTo(false));
+    }
+
+    @Test
+    void merge_with_keys_throws_when_keys_list_is_null() {
+        Event otherEvent = JacksonEvent.builder().withEventType(EventType.DOCUMENT.toString()).withData("{}").build();
+        assertThrows(IllegalArgumentException.class, () -> event.merge(otherEvent, null));
+    }
+
+    @Test
+    void merge_with_keys_throws_when_keys_list_is_empty() {
+        Event otherEvent = JacksonEvent.builder().withEventType(EventType.DOCUMENT.toString()).withData("{}").build();
+        assertThrows(IllegalArgumentException.class, () -> event.merge(otherEvent, List.of()));
+    }
+
+    @Test
+    void merge_with_keys_throws_when_other_is_not_JacksonEvent() {
+        final Event otherEvent = mock(Event.class);
+        assertThrows(IllegalArgumentException.class, () -> event.merge(otherEvent, List.of("a")));
+    }
+
+    @Test
+    void merge_with_keys_throws_when_current_event_has_array_data() {
+        final JacksonEvent arrayEvent = JacksonEvent.builder()
+                .withEventType(EventType.DOCUMENT.toString())
+                .withData("[1, 2, 3]")
+                .build();
+        final Event otherEvent = JacksonEvent.builder()
+                .withEventType(EventType.DOCUMENT.toString())
+                .withData("{\"a\": \"alpha\"}")
+                .build();
+        assertThrows(UnsupportedOperationException.class, () -> arrayEvent.merge(otherEvent, List.of("a")));
+    }
+
     @ParameterizedTest
     @ValueSource(strings = {"/", "foo", "/foo", "/foo/bar", "foo/bar", "foo/bar/", "/foo/bar/leaf/key"})
     public void testDelete_withNonexistentKey(final String key) {

data-prepper-core/src/test/java/org/opensearch/dataprepper/core/peerforwarder/discovery/AwsCloudMapPeerListProviderTest.java

@@ -347,7 +347,9 @@ void listener_gets_list_after_several_failed_attempts() {
 
             waitUntilPeerListPopulated(objectUnderTest);
 
-            assertThat(listenerEndpoints.size(), equalTo(knownIpPeers.size()));
+            await().atMost(5, TimeUnit.SECONDS)
+                    .pollDelay(100, TimeUnit.MILLISECONDS)
+                    .untilAsserted(() -> assertThat(listenerEndpoints.size(), equalTo(knownIpPeers.size())));
 
             final Set<String> observedIps = listenerEndpoints.stream()
                     .map(Endpoint::ipAddr)

data-prepper-plugins/parse-json-processor/src/main/java/org/opensearch/dataprepper/plugins/codec/json/NdjsonInputConfig.java

@@ -18,7 +18,18 @@ public class NdjsonInputConfig {
     @JsonProperty("include_empty_objects")
     private boolean includeEmptyObjects = false;
 
+    /**
+     * Optional file extension used to identify enrichment source files.
+     * Defaults to "jsonl".
+     */
+    @JsonProperty("extension")
+    private String extension = "jsonl";
+
     public boolean isIncludeEmptyObjects() {
         return includeEmptyObjects;
     }
+
+    public String getExtension() {
+        return extension;
+    }
 }

data-prepper-plugins/s3-enrich-processor/README.md

@@ -8,7 +8,7 @@ This plugin enables you to merge data from a S3 file with source data from your
 ml_merge-pipeline:
 ...
   processor:
-    - s3_enricher:
+    - s3_enrich:
         # =============================================================================
         # S3 SOURCE BUCKET CONFIGURATION
         # Defines where to fetch the original/source data for enrichment
@@ -45,7 +45,7 @@ ml_merge-pipeline:
         # =============================================================================
         # Maximum size (in MB) of S3 source files to process
         # Files exceeding this limit will be skipped
-        s3_object_size_limit_mb: 100
+        s3_object_size_limit: 100mb
     
         # JSON path in the incoming pipeline event that contains the S3 object key
         # Used to determine which source file to fetch for enrichment

data-prepper-plugins/s3-enrich-processor/build.gradle

@@ -19,9 +19,9 @@ dependencies {
     implementation 'com.fasterxml.jackson.datatype:jackson-datatype-jsr310'
     implementation 'org.projectlombok:lombok:1.18.22'
     implementation libs.parquet.common
-    implementation 'dev.failsafe:failsafe:3.3.2'
     implementation 'org.apache.httpcomponents:httpcore:4.4.16'
     implementation libs.caffeine
+    implementation 'org.hibernate.validator:hibernate-validator:8.0.1.Final'
     annotationProcessor 'org.projectlombok:lombok:1.18.20'
     implementation 'software.amazon.awssdk:s3'
     testImplementation project(':data-prepper-test:test-event')

data-prepper-plugins/s3-enrich-processor/src/main/java/org/opensearch/dataprepper/plugins/s3_enrich/processor/S3EnrichProcessor.java

@@ -43,7 +43,6 @@
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
-import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 
 import static org.opensearch.dataprepper.logging.DataPrepperMarkers.NOISY;
@@ -61,7 +60,6 @@ public class S3EnrichProcessor extends AbstractProcessor<Record<Event>, Record<E
     private final AwsCredentialsSupplier awsCredentialsSupplier;
     private final PluginSetting codecPluginSettings;
     private final PluginFactory pluginFactory;
-    private final Pattern baseNamePattern;
     private final InputCodec codec;
     private final S3ObjectWorker s3ObjectWorker;
     protected final List<String> tagsOnFailure;
@@ -91,7 +89,6 @@ public S3EnrichProcessor(final S3EnrichProcessorConfig s3EnrichProcessorConfig,
         this.pluginFactory = pluginFactory;
         CacheFactory factory = new CacheFactory(s3EnrichProcessorConfig);
         this.cacheService = new S3EnricherCacheService(factory);
-        this.baseNamePattern = Pattern.compile(s3EnrichProcessorConfig.getEnricherNamePattern());
         this.tagsOnFailure = s3EnrichProcessorConfig.getTagsOnFailure();
 
         final AwsAuthenticationAdapter awsAuthenticationAdapter = new AwsAuthenticationAdapter(awsCredentialsSupplier, s3EnrichProcessorConfig);
@@ -224,52 +221,11 @@ private void mergeData(Event targetEvent, Event sourceEvent) {
             throw new IllegalArgumentException("No merge keys configured");
         }
 
-        List<String> failedKeys = new ArrayList<>();
-        // Merge only specified keys
-        for (EventKey eventKey : mergeKeys) {
-            try {
-                mergeKey(targetEvent, sourceEvent, eventKey);
-            } catch (Exception e) {
-                LOG.error("Failed to merge key '{}': {}", eventKey, e.getMessage(), e);
-                failedKeys.add(eventKey.getKey());
-            }
-        }
-        // Handle failures based on configuration or policy
-        if (!failedKeys.isEmpty()) {
-            if (failedKeys.size() == mergeKeys.size()) {
-                // All failed - always throw
-                throw new RuntimeException("All merge keys failed: " + failedKeys);
-            } else {
-                // Partial failure - log warning
-                LOG.warn("Failed to merge {}/{} keys: {}",
-                    failedKeys.size(), mergeKeys.size(), failedKeys);
-            }
-        }
-    }
-
-    /**
-     * Merges a single key from source to target event.
-     */
-    private void mergeKey(Event targetEvent, Event sourceEvent, EventKey eventKey) {
-        String keyPath = eventKey.getKey();
-
-        // Check if source has this key
-        if (!sourceEvent.containsKey(keyPath)) {
-            LOG.debug("Source event does not contain key: {}", keyPath);
-            throw new IllegalArgumentException("Source event does not contain key: " + keyPath);
-        }
-
-        // Get value from source
-        Object sourceValue = sourceEvent.get(keyPath, Object.class);
-
-        if (sourceValue == null) {
-            LOG.debug("Source value is null for key: {}", keyPath);
-            throw new IllegalArgumentException("Source value is null for key: " + keyPath);
-        }
+        final List<String> keyPaths = mergeKeys.stream()
+                .map(EventKey::getKey)
+                .collect(Collectors.toList());
 
-        // Put into target (will overwrite if exists)
-        targetEvent.put(eventKey, sourceValue);
-        LOG.trace("Merged key '{}' with value: {}", keyPath, sourceValue);
+        targetEvent.merge(sourceEvent, keyPaths);
     }
 
     /**

data-prepper-plugins/s3-enrich-processor/src/main/java/org/opensearch/dataprepper/plugins/s3_enrich/processor/S3EnrichProcessorConfig.java

@@ -20,6 +20,8 @@
 import jakarta.validation.constraints.NotNull;
 import jakarta.validation.constraints.Size;
 import lombok.Getter;
+import org.hibernate.validator.constraints.time.DurationMax;
+import org.hibernate.validator.constraints.time.DurationMin;
 import org.opensearch.dataprepper.aws.validator.AwsAccountId;
 import org.opensearch.dataprepper.model.annotations.ExampleValues;
 import org.opensearch.dataprepper.model.configuration.PluginModel;
@@ -77,12 +79,14 @@ public class S3EnrichProcessorConfig {
     @ByteCountMax("300mb")
     private ByteCount enricherSizeLimit = ByteCount.parse(DEFAULT_ENRICHER_SIZE_LIMIT);
 
-    @JsonProperty(value = "cache_max_size", defaultValue="100000")
+    @JsonProperty(value = "cache_max_count", defaultValue="200000")
     @Min(0)
-    @Max(300000)
-    private int cacheSizeLimit = DEFAULT_CACHE_SIZE_LIMIT;
+    @Max(1000000)
+    private int cacheCountLimit = DEFAULT_CACHE_SIZE_LIMIT;
 
     @JsonProperty(value = "cache_ttl", defaultValue = "PT10M")
+    @DurationMin(minutes = 1)
+    @DurationMax(minutes = 120)
     @JsonPropertyDescription("The TTL for cache entries. Accepts ISO-8601 duration format (e.g., PT10M for 10 minutes, PT1H for 1 hour).")
     private Duration cacheTtl = DEFAULT_CACHE_TTL;
 
@@ -118,6 +122,18 @@ public class S3EnrichProcessorConfig {
                     "other parts of the configuration.")
     private List<String> tagsOnFailure = Collections.emptyList();
 
+    /**
+     * Returns the file extension configured on the codec (e.g. "jsonl").
+     * Reads the {@code extension} key from the codec plugin settings, falling back to {@code "jsonl"}.
+     */
+    public String getCodecExtension() {
+        if (codec == null || codec.getPluginSettings() == null) {
+            return "jsonl";
+        }
+        final Object ext = codec.getPluginSettings().get("extension");
+        return (ext instanceof String && !((String) ext).isBlank()) ? (String) ext : "jsonl";
+    }
+
     /**
      * Safely retrieves the S3 scan include prefix from the configuration chain.
      *