Json codec changes with specific json input codec config

data-prepper-api/src/main/java/org/opensearch/dataprepper/model/codec/JsonDecoder.java

@@ -17,32 +17,81 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.time.Instant;
+import java.util.Collection;
+import java.util.HashMap;
 import java.util.Map;
 import java.util.Objects;
 import java.util.function.Consumer;
 
 public class JsonDecoder implements ByteDecoder {
     private final ObjectMapper objectMapper = new ObjectMapper();
     private final JsonFactory jsonFactory = new JsonFactory();
+    private String keyName;
+    private Collection<String> includeKeys;
+    private Collection<String> includeKeysMetadata;
+
+    public JsonDecoder(String keyName, Collection<String> includeKeys, Collection<String> includeKeysMetadata) {
+        this.keyName = keyName;
+        this.includeKeys = includeKeys;
+        this.includeKeysMetadata = includeKeysMetadata;
+    }
+
+    public JsonDecoder() {
+        this.keyName = null;
+        this.includeKeys = null;
+        this.includeKeysMetadata = null;
+    }
 
     public void parse(InputStream inputStream, Instant timeReceived, Consumer<Record<Event>> eventConsumer) throws IOException {
         Objects.requireNonNull(inputStream);
         Objects.requireNonNull(eventConsumer);
 
         final JsonParser jsonParser = jsonFactory.createParser(inputStream);
 
+        Map<String, Object> includeKeysMap = new HashMap<>();
+        Map<String, Object> includeMetadataKeysMap = new HashMap<>();
         while (!jsonParser.isClosed() && jsonParser.nextToken() != JsonToken.END_OBJECT) {
+            final String nodeName = jsonParser.currentName();
+
+            if (includeKeys != null && includeKeys.contains(nodeName) ||
+                    (includeKeysMetadata != null && includeKeysMetadata.contains(nodeName))) {
+                jsonParser.nextToken();
+                if (includeKeys != null && includeKeys.contains(nodeName)) {
+                    includeKeysMap.put(nodeName, jsonParser.getValueAsString());
+                }
+                if (includeKeysMetadata != null && includeKeysMetadata.contains(nodeName)) {
+                    includeMetadataKeysMap.put(nodeName, jsonParser.getValueAsString());
+                }
+                continue;
+            }
+
             if (jsonParser.getCurrentToken() == JsonToken.START_ARRAY) {
-                parseRecordsArray(jsonParser, timeReceived, eventConsumer);
+                if (keyName != null && !nodeName.equals(keyName)) {
+                    continue;
+                }
+                parseRecordsArray(jsonParser, timeReceived, eventConsumer, includeKeysMap, includeMetadataKeysMap);
             }
         }
     }
 
-    private void parseRecordsArray(final JsonParser jsonParser, final Instant timeReceived, final Consumer<Record<Event>> eventConsumer) throws IOException {
+    private void parseRecordsArray(final JsonParser jsonParser,
+                                   final Instant timeReceived,
+                                   final Consumer<Record<Event>> eventConsumer,
+                                   final Map<String, Object> includeKeysMap,
+                                   final Map<String, Object> includeMetadataKeysMap
+    ) throws IOException {
         while (jsonParser.nextToken() != JsonToken.END_ARRAY) {
             final Map<String, Object> innerJson = objectMapper.readValue(jsonParser, Map.class);
 
             final Record<Event> record = createRecord(innerJson, timeReceived);
+            for (final Map.Entry<String, Object> entry : includeKeysMap.entrySet()) {
+                record.getData().put(entry.getKey(), entry.getValue());
+            }
+
+            for (final Map.Entry<String, Object> entry : includeMetadataKeysMap.entrySet()) {
+                record.getData().getMetadata().setAttribute(entry.getKey(), entry.getValue());
+            }
+
             eventConsumer.accept(record);
         }
     }

data-prepper-api/src/test/java/org/opensearch/dataprepper/model/codec/JsonDecoderTest.java

@@ -1,19 +1,30 @@
 package org.opensearch.dataprepper.model.codec;
 
+import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.ObjectMapper;
+import org.junit.jupiter.api.Nested;
 import org.junit.jupiter.api.Test;
 import org.opensearch.dataprepper.model.event.Event;
-import org.opensearch.dataprepper.model.event.DefaultEventHandle;
 import org.opensearch.dataprepper.model.record.Record;
 
 import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.io.InputStream;
 import java.time.Instant;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.LinkedHashMap;
+import java.util.List;
 import java.util.Map;
 import java.util.Random;
 import java.util.UUID;
 
 import static org.hamcrest.CoreMatchers.equalTo;
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import org.junit.jupiter.api.BeforeEach;
 
@@ -26,7 +37,7 @@ private JsonDecoder createObjectUnderTest() {
         return new JsonDecoder();
     }
 
-@BeforeEach
+    @BeforeEach
     void setup() {
         jsonDecoder = createObjectUnderTest();
         receivedRecord = null;
@@ -60,15 +71,146 @@ void test_basicJsonDecoder_withTimeReceived() {
         try {
             jsonDecoder.parse(new ByteArrayInputStream(inputString.getBytes()), now, (record) -> {
                 receivedRecord = record;
-                receivedTime = ((DefaultEventHandle)(((Event)record.getData()).getEventHandle())).getInternalOriginationTime();
+                receivedTime = record.getData().getEventHandle().getInternalOriginationTime();
             });
         } catch (Exception e){}
-        
+
         assertNotEquals(receivedRecord, null);
         Map<String, Object> map = receivedRecord.getData().toMap();
         assertThat(map.get("key1"), equalTo(stringValue));
         assertThat(map.get("key2"), equalTo(intValue));
         assertThat(receivedTime, equalTo(now));
     }
 
+    @Nested
+    class JsonDecoderWithInputConfig {
+        private ObjectMapper objectMapper;
+        private final List<String> includeKeys = new ArrayList<>();
+        private final List<String> includeMetadataKeys = new ArrayList<>();
+        private static final int numKeyRecords = 10;
+        private static final int numKeyPerRecord = 3;
+        private Map<String, Object> jsonObject;
+        private final String key_name = "logEvents";
+
+        @BeforeEach
+        void setup() {
+            objectMapper = new ObjectMapper();
+            for (int i=0; i<10; i++) {
+                includeKeys.add(UUID.randomUUID().toString());
+                includeMetadataKeys.add(UUID.randomUUID().toString());
+            }
+            jsonObject = generateJsonWithSpecificKeys(includeKeys, includeMetadataKeys, key_name, numKeyRecords, numKeyPerRecord);
+        }
+        @Test
+        void test_basicJsonDecoder_withInputConfig() throws IOException {
+            final Instant now = Instant.now();
+            List<Record<Event>> records = new ArrayList<>();
+            jsonDecoder = new JsonDecoder(key_name, includeKeys, includeMetadataKeys);
+            jsonDecoder.parse(createInputStream(jsonObject), now, (record) -> {
+                records.add(record);
+                receivedTime = record.getData().getEventHandle().getInternalOriginationTime();
+            });
+
+            assertFalse(records.isEmpty());
+            assertEquals(numKeyRecords, records.size());
+
+            records.forEach(record -> {
+                Map<String, Object> dataMap = record.getData().toMap();
+                Map<String, Object> metadataMap = record.getData().getMetadata().getAttributes();
+
+                for (String includeKey: includeKeys) {
+                    assertThat(dataMap.get(includeKey), equalTo(jsonObject.get(includeKey)));
+                }
+                for (String includeMetadataKey: includeMetadataKeys) {
+                    assertThat(metadataMap.get(includeMetadataKey), equalTo(jsonObject.get(includeMetadataKey)));
+                }
+            });
+
+            assertThat(receivedTime, equalTo(now));
+        }
+
+        @Test
+        void test_basicJsonDecoder_withInputConfig_withoutEvents_empty_metadata_keys() throws IOException {
+            final Instant now = Instant.now();
+            List<Record<Event>> records = new ArrayList<>();
+            jsonDecoder = new JsonDecoder("", includeKeys, Collections.emptyList());
+            jsonDecoder.parse(createInputStream(jsonObject), now, (record) -> {
+                records.add(record);
+                receivedTime = record.getData().getEventHandle().getInternalOriginationTime();
+            });
+            assertTrue(records.isEmpty());
+        }
+
+        @Test
+        void test_basicJsonDecoder_withInputConfig_withoutEvents_null_include_metadata_keys() throws IOException {
+            final Instant now = Instant.now();
+            List<Record<Event>> records = new ArrayList<>();
+            jsonDecoder = new JsonDecoder("", includeKeys, null);
+            jsonDecoder.parse(createInputStream(jsonObject), now, (record) -> {
+                records.add(record);
+                receivedTime = record.getData().getEventHandle().getInternalOriginationTime();
+            });
+
+            assertTrue(records.isEmpty());
+        }
+
+        @Test
+        void test_basicJsonDecoder_withInputConfig_withoutEvents_empty_include_keys() throws IOException {
+            final Instant now = Instant.now();
+            List<Record<Event>> records = new ArrayList<>();
+            jsonDecoder = new JsonDecoder("", Collections.emptyList(), includeMetadataKeys);
+            jsonDecoder.parse(createInputStream(jsonObject), now, (record) -> {
+                records.add(record);
+                receivedTime = record.getData().getEventHandle().getInternalOriginationTime();
+            });
+            assertTrue(records.isEmpty());
+        }
+
+        @Test
+        void test_basicJsonDecoder_withInputConfig_withoutEvents_null_include_keys() throws IOException {
+            final Instant now = Instant.now();
+            List<Record<Event>> records = new ArrayList<>();
+            jsonDecoder = new JsonDecoder("", null, includeMetadataKeys);
+            jsonDecoder.parse(createInputStream(jsonObject), now, (record) -> {
+                records.add(record);
+                receivedTime = record.getData().getEventHandle().getInternalOriginationTime();
+            });
+
+            assertTrue(records.isEmpty());
+        }
+
+        private Map<String, Object> generateJsonWithSpecificKeys(final List<String> includeKeys,
+                                                                 final List<String> includeMetadataKeys,
+                                                                 final String key,
+                                                                 final int numKeyRecords,
+                                                                 final int numKeyPerRecord) {
+            final Map<String, Object> jsonObject = new LinkedHashMap<>();
+            final List<Map<String, Object>> innerObjects = new ArrayList<>();
+
+            for (String includeKey: includeKeys) {
+                jsonObject.put(includeKey, UUID.randomUUID().toString());
+            }
+
+            for (String includeMetadataKey: includeMetadataKeys) {
+                jsonObject.put(includeMetadataKey, UUID.randomUUID().toString());
+            }
+
+            for (int i=0; i<numKeyRecords; i++) {
+                final Map<String, Object> innerJsonMap = new LinkedHashMap<>();
+                for (int j=0; j<numKeyPerRecord; j++) {
+                    innerJsonMap.put(UUID.randomUUID().toString(), UUID.randomUUID().toString());
+                }
+                innerObjects.add(innerJsonMap);
+            }
+            jsonObject.put(key, innerObjects);
+            return jsonObject;
+        }
+
+        private InputStream createInputStream(final Map<String, ?> jsonRoot) throws JsonProcessingException {
+            final byte[] jsonBytes = objectMapper.writeValueAsBytes(jsonRoot);
+
+            return new ByteArrayInputStream(jsonBytes);
+        }
+    }
+
 }

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

@@ -6,20 +6,28 @@
 package org.opensearch.dataprepper.plugins.codec.json;
 
 import org.opensearch.dataprepper.model.annotations.DataPrepperPlugin;
+import org.opensearch.dataprepper.model.annotations.DataPrepperPluginConstructor;
 import org.opensearch.dataprepper.model.codec.InputCodec;
 import org.opensearch.dataprepper.model.codec.JsonDecoder;
 import org.opensearch.dataprepper.model.event.Event;
 import org.opensearch.dataprepper.model.record.Record;
 
 import java.io.IOException;
 import java.io.InputStream;
+import java.util.Objects;
 import java.util.function.Consumer;
 
 /**
  * An implementation of {@link InputCodec} which parses JSON Objects for arrays.
  */
-@DataPrepperPlugin(name = "json", pluginType = InputCodec.class)
+@DataPrepperPlugin(name = "json", pluginType = InputCodec.class, pluginConfigurationType = JsonInputCodecConfig.class)
 public class JsonInputCodec extends JsonDecoder implements InputCodec {
+
+    @DataPrepperPluginConstructor
+    public JsonInputCodec(final JsonInputCodecConfig config) {
+        super(Objects.requireNonNull(config).getKeyName(), config.getIncludeKeys(), config.getIncludeKeysMetadata());
+    }
+
     public void parse(InputStream inputStream, Consumer<Record<Event>> eventConsumer) throws IOException {
         parse(inputStream, null, eventConsumer);
     }

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

@@ -0,0 +1,41 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * The OpenSearch Contributors require contributions made to
+ * this file be licensed under the Apache-2.0 license or a
+ * compatible open source license.
+ *
+ */
+
+package org.opensearch.dataprepper.plugins.codec.json;
+
+import com.fasterxml.jackson.annotation.JsonProperty;
+import jakarta.validation.constraints.Size;
+
+import java.util.List;
+
+public class JsonInputCodecConfig {
+
+    @JsonProperty("key_name")
+    @Size(min = 1, max = 2048)
+    private String keyName;
+
+    public String getKeyName() {
+        return keyName;
+    }
+
+    @JsonProperty("include_keys")
+    private List<String> includeKeys;
+
+    public List<String> getIncludeKeys() {
+        return includeKeys;
+    }
+
+    @JsonProperty("include_keys_metadata")
+    private List<String> includeKeysMetadata;
+
+    public List<String> getIncludeKeysMetadata() {
+        return includeKeysMetadata;
+    }
+}

data-prepper-plugins/parse-json-processor/src/test/java/org/opensearch/dataprepper/plugins/codec/json/JsonCodecsIT.java

@@ -36,21 +36,26 @@
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
 
 public class JsonCodecsIT {
 
     private ObjectMapper objectMapper;
     private Consumer<Record<Event>> eventConsumer;
+    private JsonInputCodecConfig jsonInputCodecConfig;
 
     @BeforeEach
     void setUp() {
         objectMapper = new ObjectMapper();
-
+        jsonInputCodecConfig = mock(JsonInputCodecConfig.class);
+        when(jsonInputCodecConfig.getIncludeKeysMetadata()).thenReturn(Collections.emptyList());
+        when(jsonInputCodecConfig.getIncludeKeys()).thenReturn(Collections.emptyList());
+        when(jsonInputCodecConfig.getKeyName()).thenReturn(null);
         eventConsumer = mock(Consumer.class);
     }
 
     private JsonInputCodec createJsonInputCodecObjectUnderTest() {
-        return new JsonInputCodec();
+        return new JsonInputCodec(jsonInputCodecConfig);
     }
 
     private JsonOutputCodec createJsonOutputCodecObjectUnderTest() {