Merge pull request #179 from BillFarber/feature/storeKafkaHeaders

Can now include Kafka headers in the document metadata

Author: BillFarber

build.gradle

@@ -5,7 +5,7 @@ plugins {
   id "com.github.jk1.dependency-license-report" version "1.19"
 
   // Only used for testing
-  id 'com.marklogic.ml-gradle' version '4.7.0'
+  id 'com.marklogic.ml-gradle' version '4.8.0'
   id 'jacoco'
   id "org.sonarqube" version "4.4.1.3373"
 
@@ -42,9 +42,9 @@ dependencies {
   compileOnly "org.apache.kafka:connect-runtime:${kafkaVersion}"
   compileOnly "org.slf4j:slf4j-api:2.0.13"
 
-  implementation 'com.marklogic:ml-javaclient-util:4.7.0'
+  implementation 'com.marklogic:ml-javaclient-util:4.8.0'
   // Force DHF to use the latest version of ml-app-deployer, which minimizes security vulnerabilities
-  implementation "com.marklogic:ml-app-deployer:4.7.0"
+  implementation "com.marklogic:ml-app-deployer:4.8.0"
 
   implementation "com.fasterxml.jackson.dataformat:jackson-dataformat-csv:2.15.3"
 

docs/writing-data.md

@@ -113,6 +113,24 @@ This will result in the following pieces of Kafka record metadata being in each
 - `kafka.partition` = the partition of the Kafka record
 - `kafka.timestamp` = the timestamp of the Kafka record
 
+### Including Kafka headers
+
+Each Kafka record passed to the MarkLogic connector also has headers that may contain useful information which can be
+included in the metadata written with documents. This includes the headers that are included in Kafka records by
+default as well as any custom headers. Kafka headers can be included in each document by configuring the following
+property:
+
+- `ml.dmsdk.includeKafkaHeaders` = `true` to include Kafka headers
+
+When the headers are added to the document metadata, they are simply given the same name as the key for the header.
+However, you may also specify a prefix that will be prepended to each header key. To set that prefix, use the following
+property:
+
+- `ml.dmsdk.includeKafkaHeaders.prefix` = `<prefix>` to be prepended to header keys in the metadata.
+
+The headers that are on the Kafka records will depend on the Kafka distribution you are using and the message producer
+configuration.
+
 ### Configuring DMSDK performance
 
 The performance of how data is written to MarkLogic can be configured via the following properties:

src/main/java/com/marklogic/kafka/connect/sink/DefaultSinkRecordConverter.java

@@ -52,13 +52,19 @@ public class DefaultSinkRecordConverter implements SinkRecordConverter {
     private DocumentWriteOperationBuilder documentWriteOperationBuilder;
     private Format format;
     private String mimeType;
-    private boolean addTopicToCollections = false;
-    private boolean includeKafkaMetadata = false;
-    private IdStrategy idStrategy;
+    private final boolean addTopicToCollections;
+    private final boolean includeKafkaMetadata;
+    private final boolean includeKafkaHeaders;
+    private final String kafkaHeadersPrefix;
+
+    private final IdStrategy idStrategy;
 
     public DefaultSinkRecordConverter(Map<String, Object> parsedConfig) {
         this.addTopicToCollections = ConfigUtil.getBoolean(MarkLogicSinkConfig.DOCUMENT_COLLECTIONS_ADD_TOPIC, parsedConfig);
         this.includeKafkaMetadata = ConfigUtil.getBoolean(MarkLogicSinkConfig.DMSDK_INCLUDE_KAFKA_METADATA, parsedConfig);
+        this.includeKafkaHeaders = ConfigUtil.getBoolean(MarkLogicSinkConfig.DMSDK_INCLUDE_KAFKA_HEADERS, parsedConfig);
+        String configPrefixValue = (String) parsedConfig.get(MarkLogicSinkConfig.DMSDK_INCLUDE_KAFKA_HEADERS_PREFIX);
+        this.kafkaHeadersPrefix = (configPrefixValue == null) ? "" : configPrefixValue;
 
         documentWriteOperationBuilder = new DocumentWriteOperationBuilder()
             .withCollections((String) parsedConfig.get(MarkLogicSinkConfig.DOCUMENT_COLLECTIONS))
@@ -94,8 +100,20 @@ private DocumentMetadataHandle buildAdditionalMetadata(SinkRecord sinkRecord) {
         if (this.addTopicToCollections) {
             metadata.getCollections().add(sinkRecord.topic());
         }
-        if (this.includeKafkaMetadata) {
+        if (this.includeKafkaMetadata || this.includeKafkaHeaders) {
             DocumentMetadataHandle.DocumentMetadataValues values = metadata.getMetadataValues();
+            addKafkaMetadataToDocumentMetadata(sinkRecord, values);
+            if (this.includeKafkaHeaders) {
+                sinkRecord.headers().forEach(
+                    header -> values.add(kafkaHeadersPrefix + header.key(), header.value().toString())
+                );
+            }
+        }
+        return metadata;
+    }
+
+    private void addKafkaMetadataToDocumentMetadata(SinkRecord sinkRecord, DocumentMetadataHandle.DocumentMetadataValues values) {
+        if (this.includeKafkaMetadata) {
             Object key = sinkRecord.key();
             if (key != null) {
                 values.add("kafka-key", key.toString());
@@ -111,7 +129,6 @@ private DocumentMetadataHandle buildAdditionalMetadata(SinkRecord sinkRecord) {
             }
             values.add("kafka-topic", sinkRecord.topic());
         }
-        return metadata;
     }
 
     /**

src/main/java/com/marklogic/kafka/connect/sink/MarkLogicSinkConfig.java

@@ -37,6 +37,8 @@ public class MarkLogicSinkConfig extends MarkLogicConfig {
     public static final String DMSDK_TRANSFORM_PARAMS = "ml.dmsdk.transformParams";
     public static final String DMSDK_TRANSFORM_PARAMS_DELIMITER = "ml.dmsdk.transformParamsDelimiter";
     public static final String DMSDK_INCLUDE_KAFKA_METADATA = "ml.dmsdk.includeKafkaMetadata";
+    public static final String DMSDK_INCLUDE_KAFKA_HEADERS = "ml.dmsdk.includeKafkaHeaders";
+    public static final String DMSDK_INCLUDE_KAFKA_HEADERS_PREFIX = "ml.dmsdk.includeKafkaHeaders.prefix";
 
     public static final String BULK_DS_ENDPOINT_URI = "ml.sink.bulkds.endpointUri";
     public static final String BULK_DS_BATCH_SIZE = "ml.sink.bulkds.batchSize";
@@ -118,6 +120,12 @@ private static ConfigDef getConfigDef() {
             .define(DMSDK_INCLUDE_KAFKA_METADATA, Type.BOOLEAN, null, Importance.LOW,
                 "Set to true so that Kafka record metadata is added to document metadata before it is written. If the document fails to be written, the Kafka record metadata will be logged as well.",
                 GROUP, -1, ConfigDef.Width.MEDIUM, "DMSDK Include Kafka Metadata")
+            .define(DMSDK_INCLUDE_KAFKA_HEADERS, Type.BOOLEAN, null, Importance.LOW,
+                "Set to true so that Kafka record headers are added to document metadata before it is written. If the document fails to be written, the Kafka record headers will be logged as well.",
+                GROUP, -1, ConfigDef.Width.MEDIUM, "DMSDK Include Kafka Headers")
+            .define(DMSDK_INCLUDE_KAFKA_HEADERS_PREFIX, Type.STRING, null, Importance.LOW,
+                "Optional value of the prefix to be added to the Kafka headers when included in document metadata.",
+                GROUP, -1, ConfigDef.Width.MEDIUM, "DMSDK Kafka Header Prefix")
 
             .define(BULK_DS_ENDPOINT_URI, Type.STRING, null, Importance.LOW,
                 "Defines the URI of a Bulk Data Services endpoint for writing data. " +

src/test/java/com/marklogic/kafka/connect/sink/ConvertSinkRecordTest.java

@@ -24,13 +24,13 @@
 import com.marklogic.client.io.Format;
 import com.marklogic.client.io.StringHandle;
 import org.apache.kafka.common.record.TimestampType;
+import org.apache.kafka.connect.data.Schema;
+import org.apache.kafka.connect.header.Header;
 import org.apache.kafka.connect.sink.SinkRecord;
 import org.junit.jupiter.api.Test;
 
 import java.io.IOException;
-import java.util.HashMap;
-import java.util.Iterator;
-import java.util.Map;
+import java.util.*;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNotNull;
@@ -78,7 +78,7 @@ void allPropertiesSet() {
 
     @Test
     void noPropertiesSet() {
-        Map<String, Object> kafkaConfig = new HashMap<String, Object>();
+        Map<String, Object> kafkaConfig = new HashMap<>();
         converter = new DefaultSinkRecordConverter(kafkaConfig);
 
         DocumentWriteOperation op = converter.convert(newSinkRecord("doesn't matter"));
@@ -221,7 +221,7 @@ void binaryContent() {
 
     @Test
     void includeKafkaMetadata() {
-        Map<String, Object> kafkaConfig = new HashMap<String, Object>();
+        Map<String, Object> kafkaConfig = new HashMap<>();
         kafkaConfig.put(MarkLogicSinkConfig.DMSDK_INCLUDE_KAFKA_METADATA, true);
         converter = new DefaultSinkRecordConverter(kafkaConfig);
 
@@ -260,6 +260,89 @@ void dontIncludeKafkaMetadata() {
         assertNull(values.get("kafka-topic"));
     }
 
+    @Test
+    void includeKafkaHeadersWithPrefix() {
+        Map<String, Object> kafkaConfig = new HashMap<>();
+        kafkaConfig.put(MarkLogicSinkConfig.DMSDK_INCLUDE_KAFKA_HEADERS, true);
+        kafkaConfig.put(MarkLogicSinkConfig.DMSDK_INCLUDE_KAFKA_HEADERS_PREFIX, "kafkaHeader_");
+        converter = new DefaultSinkRecordConverter(kafkaConfig);
+
+        final int partition = 5;
+        final long offset = 2;
+        final String key = "some-key";
+        final Long timestamp = System.currentTimeMillis();
+        List<Header> headers = new ArrayList<Header>() {{
+            add(new TestHeaders("A", "1"));
+            add(new TestHeaders("B", "2"));
+        }};
+        DocumentWriteOperation op = converter.convert(new SinkRecord("topic1", partition, null, key,
+            null, "some-value", offset, timestamp, TimestampType.CREATE_TIME, headers));
+
+        DocumentMetadataHandle metadata = (DocumentMetadataHandle) op.getMetadata();
+        DocumentMetadataHandle.DocumentMetadataValues values = metadata.getMetadataValues();
+        assertEquals("1", values.get("kafkaHeader_A"));
+        assertEquals("2", values.get("kafkaHeader_B"));
+        assertNull(values.get("kafka-offset"));
+    }
+
+    @Test
+    void includeKafkaHeadersWithoutPrefix() {
+        Map<String, Object> kafkaConfig = new HashMap<>();
+        kafkaConfig.put(MarkLogicSinkConfig.DMSDK_INCLUDE_KAFKA_HEADERS, true);
+        converter = new DefaultSinkRecordConverter(kafkaConfig);
+
+        final int partition = 5;
+        final long offset = 2;
+        final String key = "some-key";
+        final Long timestamp = System.currentTimeMillis();
+        List<Header> headers = new ArrayList<Header>() {{
+            add(new TestHeaders("A", "1"));
+            add(new TestHeaders("B", "2"));
+        }};
+        DocumentWriteOperation op = converter.convert(new SinkRecord("topic1", partition, null, key,
+            null, "some-value", offset, timestamp, TimestampType.CREATE_TIME, headers));
+
+        DocumentMetadataHandle metadata = (DocumentMetadataHandle) op.getMetadata();
+        DocumentMetadataHandle.DocumentMetadataValues values = metadata.getMetadataValues();
+        assertEquals("1", values.get("A"));
+        assertEquals("2", values.get("B"));
+        assertNull(values.get("kafka-offset"));
+    }
+
+    static class TestHeaders implements Header {
+        private final String key;
+        private final String value;
+        TestHeaders(String key, String value) {
+            this.key = key;
+            this.value = value;
+        }
+
+        @Override
+        public String key() {
+            return key;
+        }
+
+        @Override
+        public Schema schema() {
+            return null;
+        }
+
+        @Override
+        public Object value() {
+            return value;
+        }
+
+        @Override
+        public Header with(Schema schema, Object o) {
+            return null;
+        }
+
+        @Override
+        public Header rename(String s) {
+            return null;
+        }
+    }
+
     private SinkRecord newSinkRecord(Object value) {
         return new SinkRecord("test-topic", 1, null, null, null, value, 0);
     }