Add Streams HeaderEnricher

* Checkstyle

Author: garyrussell

spring-kafka/src/main/java/org/springframework/kafka/streams/HeaderEnricher.java

@@ -0,0 +1,113 @@
+/*
+ * Copyright 2019 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.kafka.streams;
+
+import java.nio.charset.StandardCharsets;
+import java.util.HashMap;
+import java.util.Map;
+
+import org.apache.kafka.common.header.Headers;
+import org.apache.kafka.common.header.internals.RecordHeader;
+import org.apache.kafka.streams.KeyValue;
+import org.apache.kafka.streams.kstream.Transformer;
+import org.apache.kafka.streams.processor.ProcessorContext;
+
+import org.springframework.expression.Expression;
+
+/**
+ * Manipulate the headers.
+ *
+ * @param <K> the key type.
+ * @param <V> the value type.
+ *
+ * @author Gary Russell
+ * @since 2.3
+ *
+ */
+public class HeaderEnricher<K, V> implements Transformer<K, V, KeyValue<K, V>> {
+
+	private final Map<String, Expression> headerExpressions = new HashMap<>();
+
+	private ProcessorContext processorContext;
+
+	public HeaderEnricher(Map<String, Expression> headerExpressions) {
+		this.headerExpressions.putAll(headerExpressions);
+	}
+
+	@Override
+	public void init(ProcessorContext context) {
+		this.processorContext = context;
+	}
+
+	@Override
+	public KeyValue<K, V> transform(K key, V value) {
+		Headers headers = this.processorContext.headers();
+		Container<K, V> container = new Container<>(this.processorContext, key, value);
+		this.headerExpressions.forEach((name, expression) -> {
+			Object headerValue = expression.getValue(container);
+			if (headerValue instanceof String) {
+				headerValue = ((String) headerValue).getBytes(StandardCharsets.UTF_8);
+			}
+			else if (!(headerValue instanceof byte[])) {
+				throw new IllegalStateException("Invalid header value type" + headerValue.getClass());
+			}
+			headers.add(new RecordHeader(name, (byte[]) headerValue));
+		});
+		return new KeyValue<>(key, value);
+	}
+
+	@Override
+	public void close() {
+		// NO-OP
+	}
+
+	/**
+	 * Container object for SpEL evaluation.
+	 *
+	 * @param <K> the key type.
+	 * @param <V> the value type.
+	 *
+	 */
+	public static final class Container<K, V> {
+
+		private final ProcessorContext context;
+
+		private final K key;
+
+		private final V value;
+
+		private Container(ProcessorContext context, K key, V value) {
+			this.context = context;
+			this.key = key;
+			this.value = value;
+		}
+
+		public ProcessorContext getContext() {
+			return this.context;
+		}
+
+		public K getKey() {
+			return this.key;
+		}
+
+		public V getValue() {
+			return this.value;
+		}
+
+	}
+
+}

spring-kafka/src/test/java/org/springframework/kafka/kstream/KafkaStreamsBranchTests.java renamed to spring-kafka/src/test/java/org/springframework/kafka/streams/KafkaStreamsBranchTests.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package org.springframework.kafka.kstream;
+package org.springframework.kafka.streams;
 
 import static org.assertj.core.api.Assertions.assertThat;
 

spring-kafka/src/test/java/org/springframework/kafka/kstream/KafkaStreamsJsonSerializationTests.java renamed to spring-kafka/src/test/java/org/springframework/kafka/streams/KafkaStreamsJsonSerializationTests.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package org.springframework.kafka.kstream;
+package org.springframework.kafka.streams;
 
 import static org.assertj.core.api.Assertions.assertThat;
 

spring-kafka/src/test/java/org/springframework/kafka/kstream/KafkaStreamsTests.java renamed to spring-kafka/src/test/java/org/springframework/kafka/streams/KafkaStreamsTests.java

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 
-package org.springframework.kafka.kstream;
+package org.springframework.kafka.streams;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.mock;
@@ -49,6 +49,9 @@
 import org.springframework.beans.factory.annotation.Value;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.expression.Expression;
+import org.springframework.expression.common.LiteralExpression;
+import org.springframework.expression.spel.standard.SpelExpressionParser;
 import org.springframework.kafka.annotation.EnableKafka;
 import org.springframework.kafka.annotation.EnableKafkaStreams;
 import org.springframework.kafka.annotation.KafkaListener;
@@ -147,6 +150,9 @@ public void testKStreams() throws Exception {
 
 		assertThat(result.topic()).isEqualTo(streamingTopic2);
 		assertThat(result.value()).isEqualTo(payload.toUpperCase() + payload2.toUpperCase());
+		assertThat(result.headers().lastHeader("foo")).isNotNull();
+		assertThat(result.headers().lastHeader("foo").value()).isEqualTo("bar".getBytes());
+		assertThat(result.headers().lastHeader("spel")).isNotNull();
 
 		assertThat(stateLatch.await(10, TimeUnit.SECONDS)).isTrue();
 
@@ -214,6 +220,11 @@ public StreamsBuilderFactoryBeanCustomizer customizer() {
 		@Bean
 		public KStream<Integer, String> kStream(StreamsBuilder kStreamBuilder) {
 			KStream<Integer, String> stream = kStreamBuilder.stream(STREAMING_TOPIC1);
+			Map<String, Expression> headers = new HashMap<>();
+			headers.put("foo", new LiteralExpression("bar"));
+			SpelExpressionParser parser = new SpelExpressionParser();
+			headers.put("spel", parser.parseExpression("context.timestamp() + key + value"));
+			HeaderEnricher<Integer, String> enricher = new HeaderEnricher<>(headers);
 			stream.mapValues((ValueMapper<String, String>) String::toUpperCase)
 					.mapValues(Foo::new)
 					.through(FOOS, Produced.with(Serdes.Integer(), new JsonSerde<Foo>() {
@@ -225,7 +236,9 @@ public KStream<Integer, String> kStream(StreamsBuilder kStreamBuilder) {
 					.reduce((value1, value2) -> value1 + value2, Materialized.as("windowStore"))
 					.toStream()
 					.map((windowedId, value) -> new KeyValue<>(windowedId.key(), value))
-					.filter((i, s) -> s.length() > 40).to(streamingTopic2);
+					.filter((i, s) -> s.length() > 40)
+					.transform(() -> enricher)
+					.to(streamingTopic2);
 
 			stream.print(Printed.toSysOut());
 

src/reference/asciidoc/streams.adoc

@@ -189,6 +189,28 @@ There must only be one such bean, or one must be marked `@Primary`.
 By default, when the factory bean is stopped, the `KafkaStreams.cleanUp()` method is called.
 Starting with version 2.1.2, the factory bean has additional constructors, taking a `CleanupConfig` object that has properties to let you control whether the `cleanUp()` method is called during `start()` or `stop()` or neither.
 
+==== Header Enricher
+
+Version 2.3 added the `HeaderEnricher` implementation of `Transformer`.
+This can be used to add headers within the stream processing; the header values are SpEL expressions; the root object of the expression evaluation has 3 properties:
+
+* `context` - the `ProcessorContext`, allowing access to the current record metadata
+* `key` - the key of the current record
+* `value` - the value of the current record
+
+The expressions must return a `byte[]` or a `String` (which will be converted to `byte[]` using `UTF-8`).
+
+To use the enricher within a stream:
+
+====
+[source, java]
+----
+.transform(() -> enricher)
+----
+====
+
+The transformer does not change the `key` or `value`; it simply adds headers.
+
 ==== Kafka Streams Example
 
 The following example combines all the topics we have covered in this chapter:

src/reference/asciidoc/whats-new.adoc

@@ -61,3 +61,8 @@ When a reply times out, the future is completed exceptionally with a `KafkaReply
 
 Extends the `ReplyingKafkaTemplate` by aggregating replies from multiple receivers.
 See <<aggregating-request-reply>> for more information.
+
+==== Kafka Streams
+
+The `HeaderEnricher` transformer has been provided, using SpEL to generate the header values.
+See <<streams-header-enricher>> for more information.