This is an automated email from the ASF dual-hosted git repository.
exceptionfactory 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 5cad5838c5 NIFI-11840 Add Kafka Partition and Topic Attributes before
handoff to RecordReader
5cad5838c5 is described below
commit 5cad5838c5b069e21596404a8f4d80c1e51a2e2b
Author: Paul Grey <[email protected]>
AuthorDate: Thu Jul 20 15:57:26 2023 -0400
NIFI-11840 Add Kafka Partition and Topic Attributes before handoff to
RecordReader
This closes #7510
Signed-off-by: David Handermann <[email protected]>
---
.../processors/kafka/pubsub/ConsumerLease.java | 7 ++--
.../kafka/pubsub/TestConsumeKafkaMock.java | 47 ++++++++++++++++++++++
2 files changed, 50 insertions(+), 4 deletions(-)
diff --git
a/nifi-nar-bundles/nifi-kafka-bundle/nifi-kafka-2-6-processors/src/main/java/org/apache/nifi/processors/kafka/pubsub/ConsumerLease.java
b/nifi-nar-bundles/nifi-kafka-bundle/nifi-kafka-2-6-processors/src/main/java/org/apache/nifi/processors/kafka/pubsub/ConsumerLease.java
index 644ca07048..698ecc36a5 100644
---
a/nifi-nar-bundles/nifi-kafka-bundle/nifi-kafka-2-6-processors/src/main/java/org/apache/nifi/processors/kafka/pubsub/ConsumerLease.java
+++
b/nifi-nar-bundles/nifi-kafka-bundle/nifi-kafka-2-6-processors/src/main/java/org/apache/nifi/processors/kafka/pubsub/ConsumerLease.java
@@ -545,8 +545,6 @@ public abstract class ConsumerLease implements Closeable,
ConsumerRebalanceListe
final Map<String, String> attributes = getAttributes(consumerRecord);
attributes.put(KAFKA_OFFSET, String.valueOf(consumerRecord.offset()));
attributes.put(KAFKA_TIMESTAMP,
String.valueOf(consumerRecord.timestamp()));
- attributes.put(KAFKA_PARTITION,
String.valueOf(consumerRecord.partition()));
- attributes.put(KAFKA_TOPIC, consumerRecord.topic());
FlowFile failureFlowFile = session.create();
@@ -572,6 +570,9 @@ public abstract class ConsumerLease implements Closeable,
ConsumerRebalanceListe
protected Map<String, String> getAttributes(final ConsumerRecord<?, ?>
consumerRecord) {
final Map<String, String> attributes = new HashMap<>();
+ attributes.put(KAFKA_PARTITION,
String.valueOf(consumerRecord.partition()));
+ attributes.put(KAFKA_TOPIC, consumerRecord.topic());
+
if (headerNamePattern == null) {
return attributes;
}
@@ -826,8 +827,6 @@ public abstract class ConsumerLease implements Closeable,
ConsumerRebalanceListe
}
}
- kafkaAttrs.put(KAFKA_PARTITION, String.valueOf(tracker.partition));
- kafkaAttrs.put(KAFKA_TOPIC, tracker.topic);
if (tracker.totalRecords > 1) {
// Add a record.count attribute to remain consistent with other
record-oriented processors. If not
// reading/writing records, then use "kafka.count" attribute.
diff --git
a/nifi-nar-bundles/nifi-kafka-bundle/nifi-kafka-2-6-processors/src/test/java/org/apache/nifi/processors/kafka/pubsub/TestConsumeKafkaMock.java
b/nifi-nar-bundles/nifi-kafka-bundle/nifi-kafka-2-6-processors/src/test/java/org/apache/nifi/processors/kafka/pubsub/TestConsumeKafkaMock.java
index ea43ab5afc..3259b3eac1 100644
---
a/nifi-nar-bundles/nifi-kafka-bundle/nifi-kafka-2-6-processors/src/test/java/org/apache/nifi/processors/kafka/pubsub/TestConsumeKafkaMock.java
+++
b/nifi-nar-bundles/nifi-kafka-bundle/nifi-kafka-2-6-processors/src/test/java/org/apache/nifi/processors/kafka/pubsub/TestConsumeKafkaMock.java
@@ -33,6 +33,9 @@ import org.apache.kafka.common.header.Headers;
import org.apache.kafka.common.header.internals.RecordHeader;
import org.apache.kafka.common.header.internals.RecordHeaders;
import org.apache.kafka.common.record.TimestampType;
+import org.apache.nifi.csv.CSVReader;
+import org.apache.nifi.csv.CSVRecordSetWriter;
+import org.apache.nifi.csv.CSVUtils;
import org.apache.nifi.json.JsonRecordSetWriter;
import org.apache.nifi.json.JsonTreeReader;
import org.apache.nifi.kafka.shared.property.OutputStrategy;
@@ -94,6 +97,50 @@ public class TestConsumeKafkaMock {
*/
private static final String TEST_GROUP = "nifi-group-" + TIMESTAMP;
+ /**
+ * Kafka topic attribute should be available to RecordReader, so that the
'ValueSeparator' character can be
+ * correctly configured.
+ */
+ @Test
+ public void testConsumeRecordDynamicReader() throws Exception {
+ final String value = "ALPHA;BETA\na1;a2\nb1;b2\n";
+ final ConsumerRecord<byte[], byte[]> record = new ConsumerRecord<>(
+ TEST_TOPIC, 0, 0, null, value.getBytes(UTF_8));
+ final ConsumerRecords<byte[], byte[]> consumerRecords =
getConsumerRecords(record);
+
+ final ConsumeKafkaRecord_2_6 processor = new ConsumeKafkaRecord_2_6() {
+ @Override
+ protected ConsumerPool createConsumerPool(final ProcessContext
context, final ComponentLog log) {
+ return getConsumerPool(consumerRecords, context, log);
+ }
+ };
+ final TestRunner runner = TestRunners.newTestRunner(processor);
+ runner.setValidateExpressionUsage(false);
+ runner.setProperty(BOOTSTRAP_SERVERS, TEST_BOOTSTRAP_SERVER);
+ runner.setProperty("topic", TEST_TOPIC);
+ runner.setProperty("topic_type", "names");
+ runner.setProperty(ConsumerConfig.GROUP_ID_CONFIG, TEST_GROUP);
+ runner.setProperty("auto.offset.reset", "earliest");
+ final String readerId = "record-reader";
+ final RecordReaderFactory readerService = new CSVReader();
+ final String writerId = "record-writer";
+ final RecordSetWriterFactory writerService = new CSVRecordSetWriter();
+ runner.addControllerService(readerId, readerService);
+ runner.setProperty(readerService, CSVUtils.VALUE_SEPARATOR,
+ "${kafka.topic:startsWith('nifi-consume'):ifElse(';', ',')}");
+ runner.enableControllerService(readerService);
+ runner.setProperty(readerId, readerId);
+ runner.addControllerService(writerId, writerService);
+ runner.enableControllerService(writerService);
+ runner.setProperty(writerId, writerId);
+ runner.run(1);
+
runner.assertAllFlowFilesTransferred(ConsumeKafkaRecord_2_6.REL_SUCCESS, 1);
+ final List<MockFlowFile> flowFiles =
runner.getFlowFilesForRelationship(ConsumeKafkaRecord_2_6.REL_SUCCESS);
+ assertEquals(1, flowFiles.size());
+ final MockFlowFile flowFile = flowFiles.iterator().next();
+ assertEquals("ALPHA,BETA\na1,a2\nb1,b2\n", flowFile.getContent());
+ }
+
@Test
public void testConsumeRecordNullKey() throws JsonProcessingException,
InitializationException {
final ObjectNode node = mapper.createObjectNode().put("a", 1).put("b",
"2");