Tagging Events in Data Prepper. Issue #629 (#2629)

* Tagging Events in Data Prepper. Issue #629 Signed-off-by: Krishna Kondaka <krishkdk@amazon.com> * Addressed review comments Signed-off-by: Krishna Kondaka <krishkdk@amazon.com> * Addressed review comments. Introduced JsonStringBuilder in JacksonEvent to return event with additinal info (like tags) as json string Signed-off-by: Krishna Kondaka <krishkdk@amazon.com> --------- Signed-off-by: Krishna Kondaka <krishkdk@amazon.com>
opensearch-project · May 5, 2023 · 38b578c · 38b578c
1 parent b7ef680
commit 38b578c
Show file tree

Hide file tree

Showing 5 changed files with 129 additions and 1 deletion.
diff --git a/...repper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventMetadata.java b/...repper-api/src/main/java/org/opensearch/dataprepper/model/event/DefaultEventMetadata.java
@@ -10,6 +10,8 @@
 import java.time.Instant;
 import java.util.Map;
 import java.util.Objects;
+import java.util.Set;
+import java.util.HashSet;
 
 import static com.google.common.base.Preconditions.checkArgument;
 import static com.google.common.base.Preconditions.checkNotNull;
@@ -27,6 +29,8 @@ public class DefaultEventMetadata implements EventMetadata {
 
     private final ImmutableMap<String, Object> attributes;
 
+    private Set<String> tags;
+
     private DefaultEventMetadata(final Builder builder) {
 
         checkNotNull(builder.eventType, "eventType cannot be null");
@@ -37,12 +41,15 @@ private DefaultEventMetadata(final Builder builder) {
         this.timeReceived = builder.timeReceived == null ? Instant.now() : builder.timeReceived;
 
         this.attributes = builder.attributes == null ? ImmutableMap.of() : ImmutableMap.copyOf(builder.attributes);
+
+        this.tags = builder.tags == null ? new HashSet<>() : new HashSet(builder.tags);
     }
 
     private DefaultEventMetadata(final EventMetadata eventMetadata) {
         this.eventType = eventMetadata.getEventType();
         this.timeReceived = eventMetadata.getTimeReceived();
         this.attributes = ImmutableMap.copyOf(eventMetadata.getAttributes());
+        this.tags = new HashSet<>(eventMetadata.getTags());
     }
 
     @Override
@@ -60,14 +67,30 @@ public Map<String, Object> getAttributes() {
         return attributes;
     }
 
+    @Override
+    public Set<String> getTags() {
+        return tags;
+    }
+
+    @Override
+    public Boolean hasTag(final String tag) {
+        return tags.contains(tag);
+    }
+
+    @Override
+    public void addTag(final String tag) {
+        tags.add(tag);
+    }
+
     @Override
     public boolean equals(final Object o) {
         if (this == o) return true;
         if (o == null || getClass() != o.getClass()) return false;
         final DefaultEventMetadata that = (DefaultEventMetadata) o;
         return Objects.equals(eventType, that.eventType)
                 && Objects.equals(timeReceived, that.timeReceived)
-                && Objects.equals(attributes, that.attributes);
+                && Objects.equals(attributes, that.attributes)
+                && Objects.equals(tags, that.tags);
     }
 
     @Override
@@ -81,6 +104,7 @@ public String toString() {
                 "eventType='" + eventType + '\'' +
                 ", timeReceived=" + timeReceived +
                 ", attributes=" + attributes +
+                ", tags=" + tags +
                 '}';
     }
 
@@ -106,6 +130,7 @@ public static class Builder {
         private String eventType;
         private Instant timeReceived;
         private Map<String, Object> attributes;
+        private Set<String> tags;
 
         /**
          * Sets the event type. A non-null or empty event type is required, otherwise it will cause {@link #build()} to fail.
@@ -140,6 +165,17 @@ public Builder withAttributes(final Map<String, Object> attributes) {
             return this;
         }
 
+        /**
+         * Sets the tags. An empty set is the default value.
+         * @param tags a set of string tags
+         * @return returns the builder
+         * @since 2.3
+         */
+        public Builder withTags(final Set<String> tags) {
+            this.tags = tags;
+            return this;
+        }
+
         /**
          * Returns a newly created {@link DefaultEventMetadata}.
          * @return an {@link DefaultEventMetadata}

diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventMetadata.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/EventMetadata.java
@@ -8,6 +8,7 @@
 import java.io.Serializable;
 import java.time.Instant;
 import java.util.Map;
+import java.util.Set;
 
 /**
  * The event metadata contains internal event fields. These fields are used only within Data Prepper, and are not passed down to Sinks.
@@ -35,4 +36,26 @@ public interface EventMetadata extends Serializable {
      * @since 1.2
      */
     Map<String, Object> getAttributes();
+
+    /**
+     * Returns the tags
+     * @return a set of tags
+     * @since 2.3
+     */
+    Set<String> getTags();
+
+    /**
+     * Indicates if a tag is present
+     * @param tag name of the tag to be looked up
+     * @return true if the tag is present, false otherwise
+     * @since 2.3
+     */
+    Boolean hasTag(final String tag);
+
+    /**
+     * Adds a tag to the Metadata
+     * @param tag to be added
+     * @since 2.3
+     */
+    void addTag(final String tag);
 }
diff --git a/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/JacksonEvent.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/JacksonEvent.java
@@ -416,6 +416,10 @@ public Builder getThis() {
         };
     }
 
+    public static JsonStringBuilder jsonBuilder() {
+        return new JsonStringBuilder();
+    }
+
     public static JacksonEvent fromEvent(final Event event) {
         if (event instanceof JacksonEvent) {
             return new JacksonEvent((JacksonEvent) event);
@@ -512,4 +516,31 @@ public JacksonEvent build() {
             return new JacksonEvent(this);
         }
     }
+
+    public static class JsonStringBuilder {
+        private String tagsKey;
+        private JacksonEvent event;
+
+        public JsonStringBuilder withEvent(final JacksonEvent event) {
+            this.event = event;
+            return this;
+        }
+
+        public JsonStringBuilder includeTags(String key) {
+            tagsKey = key;
+            return this;
+        }
+
+        public String toJsonString() {
+            if (event == null) {
+                return null;
+            }
+            final String jsonString = event.toJsonString().trim();
+            if(tagsKey != null) {
+                final JsonNode tagsNode = mapper.valueToTree(event.getMetadata().getTags());
+                return jsonString.substring(0, jsonString.length()-1) + ",\""+tagsKey+"\":" + tagsNode.toString()+"}";
+            }
+            return jsonString;
+        }
+    }
 }
diff --git a/...er-api/src/test/java/org/opensearch/dataprepper/model/event/DefaultEventMetadataTest.java b/...er-api/src/test/java/org/opensearch/dataprepper/model/event/DefaultEventMetadataTest.java
@@ -13,6 +13,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.Set;
 import java.util.UUID;
 
 import static org.hamcrest.CoreMatchers.allOf;
@@ -28,6 +29,7 @@
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
 import static org.junit.jupiter.api.Assertions.assertThrows;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.when;
 
@@ -175,6 +177,25 @@ void fromEventMetadata_returns_matching_EventMetadata() {
         assertThat(copiedMetadata.getAttributes(), not(sameInstance(attributes)));
     }
 
+    @Test
+    public void testBuild_withTags() {
+        final String testEventType = UUID.randomUUID().toString();
+
+        final Set<String> testTags = Set.of("tag1", "tag2");
+        final EventMetadata result = DefaultEventMetadata.builder()
+                .withEventType(testEventType)
+                .withTags(testTags)
+                .build();
+
+        assertThat(result, notNullValue());
+
+        assertThat(result.getTags(), equalTo(testTags));
+        result.addTag("tag3");
+        assertTrue(result.hasTag("tag1"));
+        assertTrue(result.hasTag("tag2"));
+        assertTrue(result.hasTag("tag3"));
+    }
+
     @Nested
     class EqualsAndHashCodeAndToString {
         private String eventType;

diff --git a/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/JacksonEventTest.java b/data-prepper-api/src/test/java/org/opensearch/dataprepper/model/event/JacksonEventTest.java
@@ -638,6 +638,23 @@ void testEventHandleGetAndSet() {
         assertThat(event.getEventHandle(), equalTo(testEventHandle));
     }
 
+    @Test
+    void testJsonStringBuilder() {
+        final String jsonString = "{\"foo\":\"bar\"}";
+
+        final JacksonEvent event = JacksonEvent.builder()
+                .withEventType(eventType)
+                .withData(jsonString)
+                .build();
+        final EventMetadata eventMetadata = event.getMetadata();
+        eventMetadata.addTag("tag1");
+        eventMetadata.addTag("tag2");
+        final String expectedJsonString = "{\"foo\":\"bar\",\"tags\":[\"tag1\",\"tag2\"]}";
+        assertThat(JacksonEvent.jsonBuilder().withEvent(event).includeTags("tags").toJsonString(), equalTo(expectedJsonString));
+        assertThat(JacksonEvent.jsonBuilder().withEvent(event).toJsonString(), equalTo(jsonString));
+        assertThat(JacksonEvent.jsonBuilder().toJsonString(), equalTo(null));
+    }
+
     private static Map<String, Object> createComplexDataMap() {
         final Map<String, Object> dataObject = new HashMap<>();
         final int fullDepth = 6;