Tagging Events in Data Prepper. Issue opensearch-project#629

Signed-off-by: Krishna Kondaka <krishkdk@amazon.com>
kkondaka · May 3, 2023 · 5ee91c4 · 5ee91c4
1 parent 1f90f38
commit 5ee91c4
Show file tree

Hide file tree

Showing 6 changed files with 110 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 setTag(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/Event.java b/data-prepper-api/src/main/java/org/opensearch/dataprepper/model/event/Event.java
@@ -67,6 +67,14 @@ public interface Event extends Serializable {
      */
     String toJsonString();
 
+    /**
+     * Generates a serialized Json string of the entire Event
+     * @param tagsKey key name for tags
+     * @return Json string of the event with Tags under key
+     * @since 2.3
+     */
+    String toJsonStringWithTags(final String tagsKey);
+
     /**
      * Gets a serialized Json string of the specific key in the Event
      * @param key the field to be returned

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);
+
+    /**
+     * Returns the attributes
+     * @param tag to be set
+     * @since 2.3
+     */
+    void setTag(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
@@ -276,6 +276,12 @@ public String toJsonString() {
         return jsonNode.toString();
     }
 
+    @Override
+    public String toJsonStringWithTags(final String tagsKey) {
+        String result = jsonNode.toString();
+        return result.substring(0, result.length()-1) + ",\""+tagsKey+"\":"+ getMetadata().getTags()+"}";
+    }
+
     @Override
     public String getAsJsonString(final String key) {
         final String trimmedKey = checkAndTrimKey(key);

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.setTag("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,21 @@ void testEventHandleGetAndSet() {
         assertThat(event.getEventHandle(), equalTo(testEventHandle));
     }
 
+    @Test
+    void testEventToJsonStringWithTags() {
+        final String jsonString = "{\"foo\": \"bar\"}";
+
+        final JacksonEvent event = JacksonEvent.builder()
+                .withEventType(eventType)
+                .withData(jsonString)
+                .build();
+        final EventMetadata eventMetadata = event.getMetadata();
+        eventMetadata.setTag("tag1");
+        eventMetadata.setTag("tag2");
+        final String expectedJsonString = "{\"foo\":\"bar\",\"tags\":[tag1, tag2]}";
+        assertThat(event.toJsonStringWithTags("tags"), equalTo(expectedJsonString));
+    }
+
     private static Map<String, Object> createComplexDataMap() {
         final Map<String, Object> dataObject = new HashMap<>();
         final int fullDepth = 6;