Merge pull request #987 from zalando/ARUHA-2106

ARUHA-2106: Added populating of kafka event key
zalando · Jan 2, 2019 · aabbe99 · aabbe99
2 parents b36b05b + 0bd6cb1
commit aabbe99
Show file tree

Hide file tree

Showing 4 changed files with 91 additions and 8 deletions.
diff --git a/src/main/java/org/zalando/nakadi/partitioning/HashPartitionStrategy.java b/src/main/java/org/zalando/nakadi/partitioning/HashPartitionStrategy.java
@@ -5,22 +5,20 @@
 import org.springframework.stereotype.Component;
 import org.zalando.nakadi.domain.EventCategory;
 import org.zalando.nakadi.domain.EventType;
+import org.zalando.nakadi.exceptions.Try;
 import org.zalando.nakadi.exceptions.runtime.InvalidPartitionKeyFieldsException;
 import org.zalando.nakadi.exceptions.runtime.NakadiRuntimeException;
-import org.zalando.nakadi.exceptions.Try;
 import org.zalando.nakadi.util.JsonPathAccess;
-import org.zalando.nakadi.validation.JsonSchemaEnrichment;
 
 import java.util.List;
 import java.util.stream.Collectors;
 
 import static java.lang.Math.abs;
+import static org.zalando.nakadi.validation.JsonSchemaEnrichment.DATA_PATH_PREFIX;
 
 @Component
 public class HashPartitionStrategy implements PartitionStrategy {
 
-    private static final String DATA_PATH_PREFIX = JsonSchemaEnrichment.DATA_CHANGE_WRAP_FIELD + ".";
-
     private final HashPartitionStrategyCrutch hashPartitioningCrutch;
     private final StringHash stringHash;
 

diff --git a/src/main/java/org/zalando/nakadi/service/EventPublisher.java b/src/main/java/org/zalando/nakadi/service/EventPublisher.java
@@ -10,6 +10,7 @@
 import org.zalando.nakadi.domain.BatchItem;
 import org.zalando.nakadi.domain.BatchItemResponse;
 import org.zalando.nakadi.domain.CleanupPolicy;
+import org.zalando.nakadi.domain.EventCategory;
 import org.zalando.nakadi.domain.EventPublishResult;
 import org.zalando.nakadi.domain.EventPublishingStatus;
 import org.zalando.nakadi.domain.EventPublishingStep;
@@ -26,9 +27,11 @@
 import org.zalando.nakadi.exceptions.runtime.PartitioningException;
 import org.zalando.nakadi.exceptions.runtime.ServiceTemporarilyUnavailableException;
 import org.zalando.nakadi.partitioning.PartitionResolver;
+import org.zalando.nakadi.partitioning.PartitionStrategy;
 import org.zalando.nakadi.repository.db.EventTypeCache;
 import org.zalando.nakadi.service.timeline.TimelineService;
 import org.zalando.nakadi.service.timeline.TimelineSync;
+import org.zalando.nakadi.util.JsonPathAccess;
 import org.zalando.nakadi.validation.EventTypeValidator;
 import org.zalando.nakadi.validation.ValidationError;
 
@@ -39,6 +42,8 @@
 import java.util.concurrent.TimeoutException;
 import java.util.stream.Collectors;
 
+import static org.zalando.nakadi.validation.JsonSchemaEnrichment.DATA_PATH_PREFIX;
+
 @Component
 public class EventPublisher {
 
@@ -99,7 +104,7 @@ EventPublishResult publishInternal(final String events,
 
             validate(batch, eventType);
             partition(batch, eventType);
-            compact(batch, eventType);
+            setEventKey(batch, eventType);
             enrich(batch, eventType);
             submit(batch, eventType);
 
@@ -169,14 +174,29 @@ private void partition(final List<BatchItem> batch, final EventType eventType)
         }
     }
 
-    private void compact(final List<BatchItem> batch, final EventType eventType) {
+    private void setEventKey(final List<BatchItem> batch, final EventType eventType) {
         if (eventType.getCleanupPolicy() == CleanupPolicy.COMPACT) {
             for (final BatchItem item : batch) {
                 final String compactionKey = item.getEvent()
                         .getJSONObject("metadata")
                         .getString("partition_compaction_key");
                 item.setEventKey(compactionKey);
             }
+        } else if (PartitionStrategy.HASH_STRATEGY.equals(eventType.getPartitionStrategy())) {
+            final List<String> partitionKeyFields = eventType.getPartitionKeyFields();
+            // we will set event key only if there is exactly one partition key field,
+            // in other case it's not clear what should be set as event key
+            if (partitionKeyFields.size() == 1) {
+                String partitionKeyField = partitionKeyFields.get(0);
+                if (EventCategory.DATA.equals(eventType.getCategory())) {
+                    partitionKeyField = DATA_PATH_PREFIX + partitionKeyField;
+                }
+                for (final BatchItem item : batch) {
+                    final JsonPathAccess jsonPath = new JsonPathAccess(item.getEvent());
+                    final String eventKey = jsonPath.get(partitionKeyField).toString();
+                    item.setEventKey(eventKey);
+                }
+            }
         }
     }
 

diff --git a/src/main/java/org/zalando/nakadi/validation/JsonSchemaEnrichment.java b/src/main/java/org/zalando/nakadi/validation/JsonSchemaEnrichment.java
@@ -19,6 +19,7 @@
 
 public class JsonSchemaEnrichment {
     public static final String DATA_CHANGE_WRAP_FIELD = "data";
+    public static final String DATA_PATH_PREFIX = JsonSchemaEnrichment.DATA_CHANGE_WRAP_FIELD + ".";
 
     private static final String ADDITIONAL_PROPERTIES = "additionalProperties";
     private static final String ADDITIONAL_ITEMS = "additionalItems";

diff --git a/src/test/java/org/zalando/nakadi/service/EventPublisherTest.java b/src/test/java/org/zalando/nakadi/service/EventPublisherTest.java
@@ -1,9 +1,11 @@
 package org.zalando.nakadi.service;
 
+import com.google.common.collect.ImmutableList;
 import org.json.JSONArray;
 import org.json.JSONObject;
 import org.junit.Assert;
 import org.junit.Test;
+import org.mockito.ArgumentCaptor;
 import org.mockito.Mockito;
 import org.zalando.nakadi.config.NakadiSettings;
 import org.zalando.nakadi.domain.BatchItem;
@@ -15,12 +17,13 @@
 import org.zalando.nakadi.domain.EventTypeBase;
 import org.zalando.nakadi.domain.Timeline;
 import org.zalando.nakadi.enrichment.Enrichment;
-import org.zalando.nakadi.exceptions.runtime.EnrichmentException;
-import org.zalando.nakadi.exceptions.runtime.PartitioningException;
 import org.zalando.nakadi.exceptions.runtime.AccessDeniedException;
+import org.zalando.nakadi.exceptions.runtime.EnrichmentException;
 import org.zalando.nakadi.exceptions.runtime.EventPublishingException;
 import org.zalando.nakadi.exceptions.runtime.EventTypeTimeoutException;
+import org.zalando.nakadi.exceptions.runtime.PartitioningException;
 import org.zalando.nakadi.partitioning.PartitionResolver;
+import org.zalando.nakadi.partitioning.PartitionStrategy;
 import org.zalando.nakadi.repository.TopicRepository;
 import org.zalando.nakadi.repository.db.EventTypeCache;
 import org.zalando.nakadi.service.timeline.TimelineService;
@@ -44,6 +47,7 @@
 import static org.mockito.Matchers.any;
 import static org.mockito.Matchers.anyLong;
 import static org.mockito.Matchers.eq;
+import static org.mockito.Mockito.atLeastOnce;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
@@ -404,6 +408,66 @@ public void whenEnrichmentFailsThenResultIsAborted() throws Exception {
         verify(topicRepository, times(0)).syncPostBatch(any(), any());
     }
 
+    @Test
+    public void whenSinglePartitioningKeyThenEventKeyIsSet() throws Exception {
+        final EventType eventType = EventTypeTestBuilder.builder()
+                .partitionStrategy(PartitionStrategy.HASH_STRATEGY)
+                .partitionKeyFields(ImmutableList.of("my_field"))
+                .build();
+
+        final JSONArray batch = buildDefaultBatch(1);
+        batch.getJSONObject(0).put("my_field", "my_key");
+
+        mockSuccessfulValidation(eventType);
+
+        publisher.publish(batch.toString(), eventType.getName());
+
+        final List<BatchItem> publishedBatch = capturePublishedBatch();
+        assertThat(publishedBatch.get(0).getEventKey(), equalTo("my_key"));
+    }
+
+    @Test
+    public void whenMultiplePartitioningKeyThenEventKeyIsNotSet() throws Exception {
+        final EventType eventType = EventTypeTestBuilder.builder()
+                .partitionStrategy(PartitionStrategy.HASH_STRATEGY)
+                .partitionKeyFields(ImmutableList.of("my_field", "other_field"))
+                .build();
+
+        final JSONArray batch = buildDefaultBatch(1);
+        final JSONObject event = batch.getJSONObject(0);
+        event.put("my_field", "my_key");
+        event.put("other_field", "other_value");
+
+        mockSuccessfulValidation(eventType);
+
+        publisher.publish(batch.toString(), eventType.getName());
+
+        final List<BatchItem> publishedBatch = capturePublishedBatch();
+        assertThat(publishedBatch.get(0).getEventKey(), equalTo(null));
+    }
+
+    @Test
+    public void whenNoneHashPartitioningStrategyThenEventKeyIsNotSet() throws Exception {
+        final EventType eventType = EventTypeTestBuilder.builder()
+                .partitionStrategy(PartitionStrategy.RANDOM_STRATEGY)
+                .build();
+        final JSONArray batch = buildDefaultBatch(1);
+
+        mockSuccessfulValidation(eventType);
+
+        publisher.publish(batch.toString(), eventType.getName());
+
+        final List<BatchItem> publishedBatch = capturePublishedBatch();
+        assertThat(publishedBatch.get(0).getEventKey(), equalTo(null));
+    }
+
+    @SuppressWarnings("unchecked")
+    private List<BatchItem> capturePublishedBatch() {
+        final ArgumentCaptor<List> batchCaptor = ArgumentCaptor.forClass(List.class);
+        verify(topicRepository, atLeastOnce()).syncPostBatch(any(), batchCaptor.capture());
+        return (List<BatchItem>) batchCaptor.getValue();
+    }
+
     @Test
     public void whenEnrichmentFailsThenSubsequentItemsAreAborted() throws Exception {
         final EventType eventType = buildDefaultEventType();