Exclude keys (#3055)

* Add exclude keys Signed-off-by: Kat Shen <katshen@amazon.com> --------- Signed-off-by: Kat Shen <katshen@amazon.com> Co-authored-by: Kat Shen <katshen@amazon.com>
opensearch-project · Jul 27, 2023 · 90178c4 · 90178c4
1 parent 42e274d
commit 90178c4
Show file tree

Hide file tree

Showing 4 changed files with 71 additions and 5 deletions.
diff --git a/data-prepper-plugins/key-value-processor/README.md b/data-prepper-plugins/key-value-processor/README.md
@@ -42,6 +42,9 @@ When run, the processor will parse the message into the following output:
 * `include_keys` - An array specifying the keys which should be added to parse. By default, all keys will be added.
   * Default: `[]`
   * Example: `include_keys` is `["key2"]`. `key1=value1&key2=value2` will parse into `{"key2": "value2"}`
+* `exclude_keys` - An array specifying the parsed keys which should not be added to the event. By default no keys will be excluded.
+  * Default: `[]`
+  * Example: `exclude_keys` is `["key2"]`. `key1=value1&key2=value2` will parse into `{"key1": "value1"}`
 * `key_value_delimiter_regex` - A regex specifying the delimiter between a key and a value. Special regex characters such as `[` and `]` must be escaped using `\\`.
   * There is no default.
   * Note: This cannot be defined at the same time as `value_split_characters`

diff --git a/...rc/main/java/org/opensearch/dataprepper/plugins/processor/keyvalue/KeyValueProcessor.java b/...rc/main/java/org/opensearch/dataprepper/plugins/processor/keyvalue/KeyValueProcessor.java
@@ -36,6 +36,7 @@ public class KeyValueProcessor extends AbstractProcessor<Record<Event>, Record<E
     private final Pattern fieldDelimiterPattern;
     private final Pattern keyValueDelimiterPattern;
     private final Set<String> includeKeysSet = new HashSet<String>();
+    private final Set<String> excludeKeysSet = new HashSet<String>();
     private final String lowercaseKey = "lowercase";
     private final String uppercaseKey = "uppercase";
     private final String capitalizeKey = "capitalize";
@@ -99,9 +100,10 @@ public KeyValueProcessor(final PluginMetrics pluginMetrics, final KeyValueProces
             throw new PatternSyntaxException("delete_value_regex is not a valid regex string", keyValueProcessorConfig.getDeleteValueRegex(), -1);
         }
 
-        if (keyValueProcessorConfig.getIncludeKeys() != null) {
-            includeKeysSet.addAll(keyValueProcessorConfig.getIncludeKeys());
-        }
+        includeKeysSet.addAll(keyValueProcessorConfig.getIncludeKeys());
+        excludeKeysSet.addAll(keyValueProcessorConfig.getExcludeKeys());
+
+        validateKeySets(includeKeysSet, excludeKeysSet);
 
         if (!validTransformOptionSet.contains(keyValueProcessorConfig.getTransformKey())) {
             throw new IllegalArgumentException(String.format("The transform_key value: %s is not a valid option", keyValueProcessorConfig.getTransformKey()));
@@ -153,6 +155,14 @@ private boolean validateRegex(final String pattern)
         return true;
     }
 
+    private void validateKeySets(final Set<String> includeSet, final Set<String> excludeSet) {
+        Set<String> intersectionSet = new HashSet<String>(includeSet);
+        intersectionSet.retainAll(excludeSet);
+        if (!intersectionSet.isEmpty()) {
+            throw new IllegalArgumentException("Include keys and exclude keys set cannot have any overlap", null);
+        }
+    }
+
     @Override
     public Collection<Record<Event>> doExecute(final Collection<Record<Event>> records) {
         for(final Record<Event> record : records) {
@@ -171,6 +181,11 @@ public Collection<Record<Event>> doExecute(final Collection<Record<Event>> recor
                     continue;
                 }
 
+                if (!excludeKeysSet.isEmpty() && excludeKeysSet.contains(key)) {
+                    LOG.debug(String.format("Key is being excluded: '%s'", key));
+                    continue;
+                }
+
                 if(keyValueProcessorConfig.getDeleteKeyRegex() != null && !Objects.equals(keyValueProcessorConfig.getDeleteKeyRegex(), "")) {
                     key = key.replaceAll(keyValueProcessorConfig.getDeleteKeyRegex(), "");
                 }

diff --git a/...n/java/org/opensearch/dataprepper/plugins/processor/keyvalue/KeyValueProcessorConfig.java b/...n/java/org/opensearch/dataprepper/plugins/processor/keyvalue/KeyValueProcessorConfig.java
@@ -17,6 +17,7 @@ public class KeyValueProcessorConfig {
     static final String DEFAULT_DESTINATION = "parsed_message";
     public static final String DEFAULT_FIELD_SPLIT_CHARACTERS = "&";
     static final List<String> DEFAULT_INCLUDE_KEYS = new ArrayList<>();
+    static final List<String> DEFAULT_EXCLUDE_KEYS = new ArrayList<>();
     public static final String DEFAULT_VALUE_SPLIT_CHARACTERS = "=";
     static final Object DEFAULT_NON_MATCH_VALUE = null;
     static final String DEFAULT_PREFIX = "";
@@ -44,6 +45,10 @@ public class KeyValueProcessorConfig {
     @NotNull
     private List<String> includeKeys = DEFAULT_INCLUDE_KEYS;
 
+    @JsonProperty("exclude_keys")
+    @NotNull
+    private List<String> excludeKeys = DEFAULT_EXCLUDE_KEYS;
+
     @JsonProperty("key_value_delimiter_regex")
     private String keyValueDelimiterRegex;
 
@@ -100,6 +105,10 @@ public List<String> getIncludeKeys() {
         return includeKeys;
     }
 
+    public List<String> getExcludeKeys() {
+        return excludeKeys;
+    }
+
     public String getKeyValueDelimiterRegex() {
         return keyValueDelimiterRegex;
     }

diff --git a/...st/java/org/opensearch/dataprepper/plugins/processor/keyvalue/KeyValueProcessorTests.java b/...st/java/org/opensearch/dataprepper/plugins/processor/keyvalue/KeyValueProcessorTests.java
@@ -57,6 +57,7 @@ void setup() {
         lenient().when(mockConfig.getFieldDelimiterRegex()).thenReturn(defaultConfig.getFieldDelimiterRegex());
         lenient().when(mockConfig.getFieldSplitCharacters()).thenReturn(defaultConfig.getFieldSplitCharacters());
         lenient().when(mockConfig.getIncludeKeys()).thenReturn(defaultConfig.getIncludeKeys());
+        lenient().when(mockConfig.getExcludeKeys()).thenReturn(defaultConfig.getExcludeKeys());
         lenient().when(mockConfig.getKeyValueDelimiterRegex()).thenReturn(defaultConfig.getKeyValueDelimiterRegex());
         lenient().when(mockConfig.getValueSplitCharacters()).thenReturn(defaultConfig.getValueSplitCharacters());
         lenient().when(mockConfig.getNonMatchValue()).thenReturn(defaultConfig.getNonMatchValue());
@@ -263,8 +264,8 @@ void testIncludeKeysNoMatchKeyValueProcessor() {
     }
 
     @Test
-    void testIncludeKeysAsNullKeyValueProcessor() {
-        when(mockConfig.getIncludeKeys()).thenReturn(null);
+    void testIncludeKeysAsDefaultKeyValueProcessor() {
+        when(mockConfig.getIncludeKeys()).thenReturn(List.of());
         keyValueProcessor = new KeyValueProcessor(pluginMetrics, mockConfig);
 
         final Record<Event> record = getMessage("key1=value1&key2=value2");
@@ -276,6 +277,44 @@ void testIncludeKeysAsNullKeyValueProcessor() {
         assertThatKeyEquals(parsed_message, "key2", "value2");
     }
 
+    @Test
+    void testExcludeKeysKeyValueProcessor() {
+        final List<String> excludeKeys = List.of("key2");
+        when(mockConfig.getExcludeKeys()).thenReturn(excludeKeys);
+        keyValueProcessor = new KeyValueProcessor(pluginMetrics, mockConfig);
+
+        final Record<Event> record = getMessage("key1=value1&key2=value2");
+        final List<Record<Event>> editedRecords = (List<Record<Event>>) keyValueProcessor.doExecute(Collections.singletonList(record));
+        final LinkedHashMap<String, Object> parsed_message = getLinkedHashMap(editedRecords);
+
+        assertThat(parsed_message.size(), equalTo(1));
+        assertThatKeyEquals(parsed_message, "key1", "value1");
+    }
+
+    @Test
+    void testExcludeKeysAsDefaultKeyValueProcessor() {
+        when(mockConfig.getExcludeKeys()).thenReturn(List.of());
+        keyValueProcessor = new KeyValueProcessor(pluginMetrics, mockConfig);
+
+        final Record<Event> record = getMessage("key1=value1&key2=value2");
+        final List<Record<Event>> editedRecords = (List<Record<Event>>) keyValueProcessor.doExecute(Collections.singletonList(record));
+        final LinkedHashMap<String, Object> parsed_message = getLinkedHashMap(editedRecords);
+
+        assertThat(parsed_message.size(), equalTo(2));
+        assertThatKeyEquals(parsed_message, "key1", "value1");
+        assertThatKeyEquals(parsed_message, "key2", "value2");
+    }
+
+    @Test
+    void testIncludeExcludeKeysOverlapKeyValueProcessor() {
+        final List<String> includeKeys = List.of("key1", "key3");
+        final List<String> excludeKeys = List.of("key3");
+        when(mockConfig.getIncludeKeys()).thenReturn(includeKeys);
+        when(mockConfig.getExcludeKeys()).thenReturn(excludeKeys);
+
+        assertThrows(IllegalArgumentException.class, () -> new KeyValueProcessor(pluginMetrics, mockConfig));
+    }
+
     @Test
     void testCustomPrefixKvProcessor() {
         when(mockConfig.getPrefix()).thenReturn("TEST_");