This is an automated email from the ASF dual-hosted git repository.
markap14 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/nifi.git
The following commit(s) were added to refs/heads/main by this push:
new b1ff500d76f NIFI-16450 Support Hex Encoded Kafka Key in PublishKafka
(#11795)
b1ff500d76f is described below
commit b1ff500d76f5255045eaea0113bb13984407e6c9
Author: Juhwan Kim <[email protected]>
AuthorDate: Mon Nov 2 00:01:32 2026 +0900
NIFI-16450 Support Hex Encoded Kafka Key in PublishKafka (#11795)
---
.../apache/nifi/kafka/processors/PublishKafka.java | 2 +-
.../producer/key/AttributeKeyFactory.java | 22 ++++--
.../nifi/kafka/processors/PublishKafkaTest.java | 86 ++++++++++++++++++++++
.../producer/key/AttributeKeyFactoryTest.java | 84 +++++++++++++++++++--
4 files changed, 179 insertions(+), 15 deletions(-)
diff --git
a/nifi-extension-bundles/nifi-kafka-bundle/nifi-kafka-processors/src/main/java/org/apache/nifi/kafka/processors/PublishKafka.java
b/nifi-extension-bundles/nifi-kafka-bundle/nifi-kafka-processors/src/main/java/org/apache/nifi/kafka/processors/PublishKafka.java
index f68d09836f9..00717c22236 100644
---
a/nifi-extension-bundles/nifi-kafka-bundle/nifi-kafka-processors/src/main/java/org/apache/nifi/kafka/processors/PublishKafka.java
+++
b/nifi-extension-bundles/nifi-kafka-bundle/nifi-kafka-processors/src/main/java/org/apache/nifi/kafka/processors/PublishKafka.java
@@ -567,7 +567,7 @@ public class PublishKafka extends AbstractProcessor
implements VerifiableProcess
final RecordSetWriterFactory keyWriterFactory =
context.getProperty(RECORD_KEY_WRITER).asControllerService(RecordSetWriterFactory.class);
final PropertyValue kafkaKeyAttribute = context.getProperty(KAFKA_KEY);
- final String keyAttributeEncoding =
context.getProperty(KEY_ATTRIBUTE_ENCODING).getValue();
+ final KeyEncoding keyAttributeEncoding =
context.getProperty(KEY_ATTRIBUTE_ENCODING).asAllowableValue(KeyEncoding.class);
final String messageKeyField = publishStrategy ==
PublishStrategy.USE_VALUE
?
context.getProperty(MESSAGE_KEY_FIELD).evaluateAttributeExpressions(flowFile).getValue()
: null;
final KeyFactory keyFactory = messageKeyField != null
diff --git
a/nifi-extension-bundles/nifi-kafka-bundle/nifi-kafka-processors/src/main/java/org/apache/nifi/kafka/processors/producer/key/AttributeKeyFactory.java
b/nifi-extension-bundles/nifi-kafka-bundle/nifi-kafka-processors/src/main/java/org/apache/nifi/kafka/processors/producer/key/AttributeKeyFactory.java
index c3673bd6330..ea1bea36180 100644
---
a/nifi-extension-bundles/nifi-kafka-bundle/nifi-kafka-processors/src/main/java/org/apache/nifi/kafka/processors/producer/key/AttributeKeyFactory.java
+++
b/nifi-extension-bundles/nifi-kafka-bundle/nifi-kafka-processors/src/main/java/org/apache/nifi/kafka/processors/producer/key/AttributeKeyFactory.java
@@ -19,29 +19,37 @@ package org.apache.nifi.kafka.processors.producer.key;
import org.apache.nifi.components.PropertyValue;
import org.apache.nifi.flowfile.FlowFile;
import org.apache.nifi.kafka.shared.attribute.KafkaFlowFileAttribute;
+import org.apache.nifi.kafka.shared.property.KeyEncoding;
import org.apache.nifi.serialization.record.Record;
-import java.io.UnsupportedEncodingException;
import java.nio.charset.StandardCharsets;
+import java.util.HexFormat;
import java.util.Map;
-import java.util.Optional;
+import java.util.Objects;
public class AttributeKeyFactory implements KeyFactory {
private final FlowFile flowFile;
private final PropertyValue keyAttribute;
- private final String keyAttributeEncoding;
+ private final KeyEncoding keyEncoding;
- public AttributeKeyFactory(final FlowFile flowFile, final PropertyValue
keyAttribute, final String keyAttributeEncoding) {
+ public AttributeKeyFactory(final FlowFile flowFile, final PropertyValue
keyAttribute, final KeyEncoding keyEncoding) {
this.flowFile = flowFile;
this.keyAttribute = keyAttribute;
- this.keyAttributeEncoding =
Optional.ofNullable(keyAttributeEncoding).orElse(StandardCharsets.UTF_8.name());
+ this.keyEncoding = Objects.requireNonNull(keyEncoding, "Key Encoding
required");
}
@Override
- public byte[] getKey(final Map<String, String> attributes, final Record
record) throws UnsupportedEncodingException {
+ public byte[] getKey(final Map<String, String> attributes, final Record
record) {
final String keyAttributeValue = keyAttribute.isSet()
?
keyAttribute.evaluateAttributeExpressions(flowFile).getValue()
: attributes.get(KafkaFlowFileAttribute.KAFKA_KEY);
- return (keyAttributeValue == null) ? null :
keyAttributeValue.getBytes(keyAttributeEncoding);
+ if (keyAttributeValue == null) {
+ return null;
+ }
+ return switch (keyEncoding) {
+ case HEX -> HexFormat.of().parseHex(keyAttributeValue);
+ case UTF8 -> keyAttributeValue.getBytes(StandardCharsets.UTF_8);
+ case DO_NOT_ADD -> null;
+ };
}
}
diff --git
a/nifi-extension-bundles/nifi-kafka-bundle/nifi-kafka-processors/src/test/java/org/apache/nifi/kafka/processors/PublishKafkaTest.java
b/nifi-extension-bundles/nifi-kafka-bundle/nifi-kafka-processors/src/test/java/org/apache/nifi/kafka/processors/PublishKafkaTest.java
index 69552cee70a..4c6ef80720f 100644
---
a/nifi-extension-bundles/nifi-kafka-bundle/nifi-kafka-processors/src/test/java/org/apache/nifi/kafka/processors/PublishKafkaTest.java
+++
b/nifi-extension-bundles/nifi-kafka-bundle/nifi-kafka-processors/src/test/java/org/apache/nifi/kafka/processors/PublishKafkaTest.java
@@ -19,7 +19,13 @@ package org.apache.nifi.kafka.processors;
import org.apache.nifi.components.ConfigVerificationResult;
import org.apache.nifi.kafka.service.api.KafkaConnectionService;
import org.apache.nifi.kafka.service.api.common.PartitionState;
+import org.apache.nifi.kafka.service.api.producer.FlowFileResult;
import org.apache.nifi.kafka.service.api.producer.KafkaProducerService;
+import org.apache.nifi.kafka.service.api.producer.PublishContext;
+import org.apache.nifi.kafka.service.api.producer.RecordSummary;
+import org.apache.nifi.kafka.service.api.record.KafkaRecord;
+import org.apache.nifi.kafka.shared.attribute.KafkaFlowFileAttribute;
+import org.apache.nifi.kafka.shared.property.KeyEncoding;
import org.apache.nifi.reporting.InitializationException;
import org.apache.nifi.util.TestRunner;
import org.apache.nifi.util.TestRunners;
@@ -29,15 +35,23 @@ import org.junit.jupiter.api.extension.ExtendWith;
import org.mockito.Mock;
import org.mockito.junit.jupiter.MockitoExtension;
+import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
import java.util.Collections;
+import java.util.Iterator;
+import java.util.LinkedHashMap;
import java.util.List;
+import java.util.Map;
import static org.apache.nifi.kafka.processors.PublishKafka.CONNECTION_SERVICE;
import static org.apache.nifi.kafka.processors.PublishKafka.TOPIC_NAME;
+import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNotNull;
+import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.when;
@ExtendWith(MockitoExtension.class)
@@ -64,6 +78,10 @@ class PublishKafkaTest {
private PublishKafka processor;
+ private final List<KafkaRecord> sentRecords = new ArrayList<>();
+
+ private final Map<PublishContext, Long> sentCounts = new LinkedHashMap<>();
+
@BeforeEach
public void setRunner() {
processor = new PublishKafka();
@@ -124,6 +142,43 @@ class PublishKafkaTest {
runner.enableControllerService(kafkaConnectionService);
}
+ @Test
+ public void testPublishHexEncodedKey() throws InitializationException {
+ setProducerService();
+ runner.setProperty(PublishKafka.KEY_ATTRIBUTE_ENCODING,
KeyEncoding.HEX.getValue());
+
+ runner.enqueue("value".getBytes(StandardCharsets.UTF_8),
Map.of(KafkaFlowFileAttribute.KAFKA_KEY, "0A0b0C"));
+ runner.run();
+
+ runner.assertAllFlowFilesTransferred(PublishKafka.REL_SUCCESS, 1);
+ assertEquals(1, sentRecords.size());
+ assertArrayEquals(new byte[]{0x0a, 0x0b, 0x0c},
sentRecords.getFirst().getKey());
+ }
+
+ @Test
+ public void testPublishUtf8Key() throws InitializationException {
+ setProducerService();
+
+ runner.enqueue("value".getBytes(StandardCharsets.UTF_8),
Map.of(KafkaFlowFileAttribute.KAFKA_KEY, "key-1"));
+ runner.run();
+
+ runner.assertAllFlowFilesTransferred(PublishKafka.REL_SUCCESS, 1);
+ assertEquals(1, sentRecords.size());
+ assertArrayEquals("key-1".getBytes(StandardCharsets.UTF_8),
sentRecords.getFirst().getKey());
+ }
+
+ @Test
+ public void testPublishInvalidHexEncodedKey() throws
InitializationException {
+ setProducerService();
+ runner.setProperty(PublishKafka.KEY_ATTRIBUTE_ENCODING,
KeyEncoding.HEX.getValue());
+
+ runner.enqueue("value".getBytes(StandardCharsets.UTF_8),
Map.of(KafkaFlowFileAttribute.KAFKA_KEY, "0a0"));
+ runner.run();
+
+ runner.assertAllFlowFilesTransferred(PublishKafka.REL_FAILURE, 1);
+ assertTrue(sentRecords.isEmpty());
+ }
+
private void setConnectionService() throws InitializationException {
when(kafkaConnectionService.getIdentifier()).thenReturn(SERVICE_ID);
@@ -132,4 +187,35 @@ class PublishKafkaTest {
runner.setProperty(CONNECTION_SERVICE, SERVICE_ID);
}
+
+ private void setProducerService() throws InitializationException {
+ setConnectionService();
+
when(kafkaConnectionService.getProducerService(any())).thenReturn(kafkaProducerService);
+ runner.setProperty(TOPIC_NAME, TEST_TOPIC_NAME);
+
+ doAnswer(invocation -> {
+ final Iterator<KafkaRecord> records = invocation.getArgument(0);
+ final PublishContext publishContext = invocation.getArgument(1);
+ long sentCount = 0;
+ while (records.hasNext()) {
+ sentRecords.add(records.next());
+ sentCount++;
+ }
+ sentCounts.merge(publishContext, sentCount, Long::sum);
+ return null;
+ }).when(kafkaProducerService).send(any(), any());
+
+ when(kafkaProducerService.complete()).thenAnswer(invocation -> {
+ final RecordSummary recordSummary = new RecordSummary();
+ for (final Map.Entry<PublishContext, Long> entry :
sentCounts.entrySet()) {
+ final PublishContext publishContext = entry.getKey();
+ final long sentCount = entry.getValue();
+ final Exception exception = publishContext.getException();
+ final List<Exception> exceptions = exception == null ?
List.of() : List.of(exception);
+ recordSummary.getFlowFileResults().add(new FlowFileResult(
+ publishContext.getFlowFile(), sentCount,
Map.of(TEST_TOPIC_NAME, sentCount), List.of(), exceptions));
+ }
+ return recordSummary;
+ });
+ }
}
diff --git
a/nifi-extension-bundles/nifi-kafka-bundle/nifi-kafka-processors/src/test/java/org/apache/nifi/kafka/processors/producer/key/AttributeKeyFactoryTest.java
b/nifi-extension-bundles/nifi-kafka-bundle/nifi-kafka-processors/src/test/java/org/apache/nifi/kafka/processors/producer/key/AttributeKeyFactoryTest.java
index 5b14a10c852..c9b2675bef2 100644
---
a/nifi-extension-bundles/nifi-kafka-bundle/nifi-kafka-processors/src/test/java/org/apache/nifi/kafka/processors/producer/key/AttributeKeyFactoryTest.java
+++
b/nifi-extension-bundles/nifi-kafka-bundle/nifi-kafka-processors/src/test/java/org/apache/nifi/kafka/processors/producer/key/AttributeKeyFactoryTest.java
@@ -17,6 +17,8 @@
package org.apache.nifi.kafka.processors.producer.key;
import org.apache.nifi.components.PropertyValue;
+import org.apache.nifi.kafka.shared.attribute.KafkaFlowFileAttribute;
+import org.apache.nifi.kafka.shared.property.KeyEncoding;
import org.apache.nifi.serialization.SimpleRecordSchema;
import org.apache.nifi.serialization.record.MapRecord;
import org.apache.nifi.serialization.record.Record;
@@ -27,7 +29,6 @@ import org.apache.nifi.util.MockFlowFile;
import org.apache.nifi.util.MockPropertyValue;
import org.junit.jupiter.api.Test;
-import java.io.UnsupportedEncodingException;
import java.nio.charset.StandardCharsets;
import java.util.Arrays;
import java.util.HashMap;
@@ -36,33 +37,34 @@ import java.util.Map;
import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
public class AttributeKeyFactoryTest {
@Test
- void testNullKeyAttribute() throws UnsupportedEncodingException {
+ void testNullKeyAttribute() {
final Map<String, String> attributes = new HashMap<>();
final Record record = fabricateRecord();
final PropertyValue propertyValue = new MockPropertyValue(null);
- final AttributeKeyFactory attributeKeyFactory = new
AttributeKeyFactory(null, propertyValue, null);
+ final AttributeKeyFactory attributeKeyFactory = new
AttributeKeyFactory(null, propertyValue, KeyEncoding.UTF8);
assertNull(attributeKeyFactory.getKey(attributes, record));
}
@Test
- void testNullKeyAttributeValue() throws UnsupportedEncodingException {
+ void testNullKeyAttributeValue() {
final Map<String, String> attributes = new HashMap<>();
final Record record = fabricateRecord();
final MockFlowFile flowFile = new MockFlowFile(1L);
flowFile.putAttributes(attributes);
final PropertyValue propertyValue = new MockPropertyValue("${A}");
- final AttributeKeyFactory attributeKeyFactory = new
AttributeKeyFactory(flowFile, propertyValue, null);
+ final AttributeKeyFactory attributeKeyFactory = new
AttributeKeyFactory(flowFile, propertyValue, KeyEncoding.UTF8);
assertEquals(0, attributeKeyFactory.getKey(attributes, record).length);
}
@Test
- void testNonNullKeyAttribute() throws UnsupportedEncodingException {
+ void testNonNullKeyAttribute() {
final Map<String, String> attributes = new HashMap<>();
attributes.put("A", "valueA");
attributes.put("B", "valueB");
@@ -71,7 +73,75 @@ public class AttributeKeyFactoryTest {
flowFile.putAttributes(attributes);
final PropertyValue propertyValue = new MockPropertyValue("${A}");
- final AttributeKeyFactory attributeKeyFactory = new
AttributeKeyFactory(flowFile, propertyValue, null);
+ final AttributeKeyFactory attributeKeyFactory = new
AttributeKeyFactory(flowFile, propertyValue, KeyEncoding.UTF8);
+ assertArrayEquals("valueA".getBytes(StandardCharsets.UTF_8),
attributeKeyFactory.getKey(attributes, record));
+ }
+
+ @Test
+ void testHexEncodedKeyAttribute() {
+ final Map<String, String> attributes = new HashMap<>();
+ attributes.put(KafkaFlowFileAttribute.KAFKA_KEY, "0a0b0c");
+ final Record record = fabricateRecord();
+ final PropertyValue propertyValue = new MockPropertyValue(null);
+
+ final AttributeKeyFactory attributeKeyFactory = new
AttributeKeyFactory(null, propertyValue, KeyEncoding.HEX);
+ assertArrayEquals(new byte[]{0x0a, 0x0b, 0x0c},
attributeKeyFactory.getKey(attributes, record));
+ }
+
+ @Test
+ void testHexEncodedKeyAttributeUpperCase() {
+ final Map<String, String> attributes = new HashMap<>();
+ attributes.put(KafkaFlowFileAttribute.KAFKA_KEY, "0A0B0C");
+ final Record record = fabricateRecord();
+ final PropertyValue propertyValue = new MockPropertyValue(null);
+
+ final AttributeKeyFactory attributeKeyFactory = new
AttributeKeyFactory(null, propertyValue, KeyEncoding.HEX);
+ assertArrayEquals(new byte[]{0x0a, 0x0b, 0x0c},
attributeKeyFactory.getKey(attributes, record));
+ }
+
+ @Test
+ void testHexEncodedKafkaKeyProperty() {
+ final Map<String, String> attributes = new HashMap<>();
+ attributes.put("A", "ff00");
+ final Record record = fabricateRecord();
+ final MockFlowFile flowFile = new MockFlowFile(1L);
+ flowFile.putAttributes(attributes);
+ final PropertyValue propertyValue = new MockPropertyValue("${A}");
+
+ final AttributeKeyFactory attributeKeyFactory = new
AttributeKeyFactory(flowFile, propertyValue, KeyEncoding.HEX);
+ assertArrayEquals(new byte[]{(byte) 0xff, 0x00},
attributeKeyFactory.getKey(attributes, record));
+ }
+
+ @Test
+ void testInvalidHexEncodedKeyAttribute() {
+ final Map<String, String> attributes = new HashMap<>();
+ attributes.put(KafkaFlowFileAttribute.KAFKA_KEY, "0a0");
+ final Record record = fabricateRecord();
+ final PropertyValue propertyValue = new MockPropertyValue(null);
+
+ final AttributeKeyFactory attributeKeyFactory = new
AttributeKeyFactory(null, propertyValue, KeyEncoding.HEX);
+ assertThrows(IllegalArgumentException.class, () ->
attributeKeyFactory.getKey(attributes, record));
+ }
+
+ @Test
+ void testDoNotAddKeyEncoding() {
+ final Map<String, String> attributes = new HashMap<>();
+ attributes.put(KafkaFlowFileAttribute.KAFKA_KEY, "0a0b0c");
+ final Record record = fabricateRecord();
+ final PropertyValue propertyValue = new MockPropertyValue(null);
+
+ final AttributeKeyFactory attributeKeyFactory = new
AttributeKeyFactory(null, propertyValue, KeyEncoding.DO_NOT_ADD);
+ assertNull(attributeKeyFactory.getKey(attributes, record));
+ }
+
+ @Test
+ void testUtf8KeyEncodingExplicit() {
+ final Map<String, String> attributes = new HashMap<>();
+ attributes.put(KafkaFlowFileAttribute.KAFKA_KEY, "valueA");
+ final Record record = fabricateRecord();
+ final PropertyValue propertyValue = new MockPropertyValue(null);
+
+ final AttributeKeyFactory attributeKeyFactory = new
AttributeKeyFactory(null, propertyValue, KeyEncoding.UTF8);
assertArrayEquals("valueA".getBytes(StandardCharsets.UTF_8),
attributeKeyFactory.getKey(attributes, record));
}
