Add header routing setting to KafkaOutboundChannelAdapter

Addressing PR comments

Simple polishing

spring-integration-kafka/src/main/java/org/springframework/integration/kafka/config/xml/KafkaOutboundChannelAdapterParser.java

@@ -67,6 +67,9 @@ public class KafkaOutboundChannelAdapterParser extends AbstractOutboundChannelAd
 			kafkaProducerMessageHandlerBuilder.addPropertyValue("partitionIdExpression", partitionIdExpressionDef);
 		}
 
+		IntegrationNamespaceUtils.setValueIfAttributeDefined(kafkaProducerMessageHandlerBuilder, element,
+				"enable-header-routing");
+
 		return kafkaProducerMessageHandlerBuilder.getBeanDefinition();
 	}
 

spring-integration-kafka/src/main/java/org/springframework/integration/kafka/outbound/KafkaProducerMessageHandler.java

@@ -28,6 +28,7 @@ import org.springframework.messaging.Message;
  * @author Soby Chacko
  * @author Artem Bilan
  * @author Gary Russell
+ * @author Marius Bogoevici
  * @since 0.5
  */
 public class KafkaProducerMessageHandler extends AbstractMessageHandler {
@@ -36,17 +37,31 @@ public class KafkaProducerMessageHandler extends AbstractMessageHandler {
 
 	private EvaluationContext evaluationContext;
 
+	private boolean enableHeaderRouting = true;
+
 	private volatile Expression topicExpression;
 
 	private volatile Expression messageKeyExpression;
 
 	private volatile Expression partitionIdExpression;
 
-	@SuppressWarnings("unchecked")
 	public KafkaProducerMessageHandler(final KafkaProducerContext kafkaProducerContext) {
 		this.kafkaProducerContext = kafkaProducerContext;
 	}
 
+	/**
+	 * Enable the use of headers for determining the target topic and partition of outbound messages. By default it is
+	 * set to true, but it can be disabled when those values are produced by upstream components that read messages
+	 * from Kafka sources themselves.
+	 * @param enableHeaderRouting whether the topic and destination headers should be considered
+	 * @since 1.3
+	 * @see KafkaHeaders#TOPIC
+	 * @see KafkaHeaders#PARTITION_ID
+	 */
+	public void setEnableHeaderRouting(boolean enableHeaderRouting) {
+		this.enableHeaderRouting = enableHeaderRouting;
+	}
+
 	public void setTopicExpression(Expression topicExpression) {
 		this.topicExpression = topicExpression;
 	}
@@ -82,11 +97,12 @@ public class KafkaProducerMessageHandler extends AbstractMessageHandler {
 	protected void handleMessageInternal(final Message<?> message) throws Exception {
 		String topic = this.topicExpression != null ?
 				this.topicExpression.getValue(this.evaluationContext, message, String.class)
-				: message.getHeaders().get(KafkaHeaders.TOPIC, String.class);
+				//TODO revise the headers fallback behavior in favor of just expression
+				: (this.enableHeaderRouting ? message.getHeaders().get(KafkaHeaders.TOPIC, String.class) : null);
 
 		Integer partitionId = this.partitionIdExpression != null ?
 				this.partitionIdExpression.getValue(this.evaluationContext, message, Integer.class)
-				: message.getHeaders().get(KafkaHeaders.PARTITION_ID, Integer.class);
+				: (this.enableHeaderRouting ? message.getHeaders().get(KafkaHeaders.PARTITION_ID, Integer.class) : null);
 
 		Object messageKey = this.messageKeyExpression != null
 				? this.messageKeyExpression.getValue(this.evaluationContext, message)

spring-integration-kafka/src/main/resources/org/springframework/integration/kafka/config/spring-integration-kafka-1.3.xsd

@@ -558,6 +558,14 @@
 							]]></xsd:documentation>
 				</xsd:annotation>
 			</xsd:attribute>
+			<xsd:attribute name="enable-header-routing" type="xsd:string">
+				<xsd:annotation>
+					<xsd:documentation><![CDATA[
+								Enables the use of message headers for routing messages to specific topics and
+								partitions.
+							]]></xsd:documentation>
+				</xsd:annotation>
+			</xsd:attribute>
 			<xsd:attribute name="order">
 				<xsd:annotation>
 					<xsd:documentation>

spring-integration-kafka/src/test/java/org/springframework/integration/kafka/config/xml/KafkaOutboundAdapterParserTests-context.xml

@@ -14,6 +14,17 @@
 	</int:channel>
 
 	<int-kafka:outbound-channel-adapter id="kafkaOutboundChannelAdapter"
+										kafka-producer-context-ref="kafkaProducerContext"
+										auto-startup="false"
+										channel="inputToKafka"
+										order="3"
+										topic="foo"
+										message-key-expression="'bar'"
+										partition-id-expression="2">
+		<int:poller fixed-delay="1000" time-unit="MILLISECONDS" receive-timeout="0" task-executor="taskExecutor"/>
+	</int-kafka:outbound-channel-adapter>
+
+	<int-kafka:outbound-channel-adapter id="kafkaOutboundChannelAdapterWithHeaderRoutingDisabled"
 										kafka-producer-context-ref="kafkaProducerContext"
 										auto-startup="false"
 										channel="inputToKafka"
@@ -21,7 +32,7 @@
 										topic="foo"
 										message-key-expression="'bar'"
 										partition-id-expression="2"
-										>
+										enable-header-routing="false">
 		<int:poller fixed-delay="1000" time-unit="MILLISECONDS" receive-timeout="0" task-executor="taskExecutor"/>
 	</int-kafka:outbound-channel-adapter>
 

spring-integration-kafka/src/test/java/org/springframework/integration/kafka/config/xml/KafkaOutboundAdapterParserTests.java

@@ -27,7 +27,8 @@ import org.springframework.beans.factory.annotation.Autowired;
 import org.springframework.context.ApplicationContext;
 import org.springframework.integration.endpoint.PollingConsumer;
 import org.springframework.integration.kafka.outbound.KafkaProducerMessageHandler;
-import org.springframework.integration.kafka.rule.KafkaRunning;
+import org.springframework.integration.kafka.rule.KafkaEmbedded;
+import org.springframework.integration.kafka.rule.KafkaRule;
 import org.springframework.integration.kafka.support.KafkaProducerContext;
 import org.springframework.integration.test.util.TestUtils;
 import org.springframework.test.context.ContextConfiguration;
@@ -44,25 +45,32 @@ import org.springframework.test.context.junit4.SpringJUnit4ClassRunner;
 public class KafkaOutboundAdapterParserTests {
 
 	@ClassRule
-	public static KafkaRunning kafkaRunning = KafkaRunning.isRunning();
+	public static KafkaRule kafkaRunning = new KafkaEmbedded(1);
 
 	@Autowired
 	private ApplicationContext appContext;
 
 	@Test
-	@SuppressWarnings("unchecked")
 	public void testOutboundAdapterConfiguration() {
 		PollingConsumer pollingConsumer = this.appContext.getBean("kafkaOutboundChannelAdapter", PollingConsumer.class);
-		KafkaProducerMessageHandler messageHandler = this.appContext.getBean(KafkaProducerMessageHandler.class);
+		KafkaProducerMessageHandler messageHandler
+				= this.appContext.getBean("org.springframework.integration.kafka.outbound.KafkaProducerMessageHandler#0",
+				KafkaProducerMessageHandler.class);
 		assertNotNull(pollingConsumer);
 		assertNotNull(messageHandler);
 		assertEquals(messageHandler.getOrder(), 3);
 		assertEquals("foo", TestUtils.getPropertyValue(messageHandler, "topicExpression.literalValue"));
 		assertEquals("'bar'", TestUtils.getPropertyValue(messageHandler, "messageKeyExpression.expression"));
 		assertEquals("2", TestUtils.getPropertyValue(messageHandler, "partitionIdExpression.expression"));
+		assertEquals(true, TestUtils.getPropertyValue(messageHandler, "enableHeaderRouting"));
 		KafkaProducerContext producerContext = messageHandler.getKafkaProducerContext();
 		assertNotNull(producerContext);
 		assertEquals(producerContext.getProducerConfigurations().size(), 2);
+
+		KafkaProducerMessageHandler messageHandler2
+				= this.appContext.getBean("org.springframework.integration.kafka.outbound.KafkaProducerMessageHandler#1",
+				KafkaProducerMessageHandler.class);
+		assertEquals(false, TestUtils.getPropertyValue(messageHandler2, "enableHeaderRouting"));
 	}
 
 }

spring-integration-kafka/src/test/java/org/springframework/integration/kafka/outbound/OutboundTests.java

@@ -16,15 +16,6 @@
 
 package org.springframework.integration.kafka.outbound;
 
-import static org.hamcrest.Matchers.contains;
-import static org.hamcrest.Matchers.endsWith;
-import static org.hamcrest.Matchers.equalTo;
-import static org.hamcrest.Matchers.hasItem;
-import static org.hamcrest.Matchers.hasSize;
-import static org.junit.Assert.assertThat;
-import static org.junit.Assert.assertTrue;
-import static org.mockito.Mockito.mock;
-
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.List;
@@ -34,6 +25,13 @@ import java.util.UUID;
 import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.TimeUnit;
 
+import com.gs.collections.api.multimap.MutableMultimap;
+import com.gs.collections.impl.factory.Multimaps;
+import kafka.admin.AdminUtils;
+import kafka.api.OffsetRequest;
+import kafka.common.TopicExistsException;
+import kafka.serializer.Decoder;
+import kafka.serializer.Encoder;
 import org.apache.kafka.clients.producer.RecordMetadata;
 import org.apache.kafka.common.serialization.ByteArraySerializer;
 import org.apache.kafka.common.serialization.StringSerializer;
@@ -65,14 +63,14 @@ import org.springframework.integration.kafka.util.MessageUtils;
 import org.springframework.integration.kafka.util.TopicUtils;
 import org.springframework.messaging.support.MessageBuilder;
 
-import com.gs.collections.api.multimap.MutableMultimap;
-import com.gs.collections.impl.factory.Multimaps;
-
-import kafka.admin.AdminUtils;
-import kafka.api.OffsetRequest;
-import kafka.common.TopicExistsException;
-import kafka.serializer.Decoder;
-import kafka.serializer.Encoder;
+import static org.hamcrest.Matchers.contains;
+import static org.hamcrest.Matchers.endsWith;
+import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.hasItem;
+import static org.hamcrest.Matchers.hasSize;
+import static org.junit.Assert.assertThat;
+import static org.junit.Assert.assertTrue;
+import static org.mockito.Mockito.mock;
 
 /**
  * @author Gary Russell
@@ -265,6 +263,97 @@ public class OutboundTests {
 		kafkaMessageListenerContainer.stop();
 	}
 
+	@Test
+	public void testHeaderRoutingDisabled() throws Exception {
+
+		// create the topic
+
+		try {
+			TopicUtils.ensureTopicCreated(kafkaRule.getZookeeperConnectionString(), TOPIC, 1, 1);
+		}
+		catch (TopicExistsException e) {
+			// do nothing
+		}
+
+		try {
+			TopicUtils.ensureTopicCreated(kafkaRule.getZookeeperConnectionString(), TOPIC2, 1, 1);
+		}
+		catch (TopicExistsException e) {
+			// do nothing
+		}
+
+		final String suffix = UUID.randomUUID().toString();
+
+		KafkaMessageListenerContainer kafkaMessageListenerContainer = createMessageListenerContainer(TOPIC, TOPIC2);
+
+		final Decoder<String> decoder = new StringDecoder();
+
+		int expectedMessageCount = 2;
+		final MutableMultimap<String, String> payloadsByTopic = Multimaps.mutable.list.with();
+		final CountDownLatch latch = new CountDownLatch(expectedMessageCount);
+		kafkaMessageListenerContainer.setMessageListener(new MessageListener() {
+
+			@Override
+			public void onMessage(KafkaMessage message) {
+				payloadsByTopic.put(message.getMetadata().getPartition().getTopic(),
+						MessageUtils.decodePayload(message, decoder));
+				latch.countDown();
+			}
+
+		});
+
+		kafkaMessageListenerContainer.start();
+
+		int expectedDeliveryConfirmations = 2;
+		final List<RecordMetadata> results = new ArrayList<RecordMetadata>();
+		final CountDownLatch sendResultLatch = new CountDownLatch(expectedDeliveryConfirmations);
+		ProducerListener listener = new ProducerListener() {
+
+			@Override
+			public void onSuccess(String topic, Integer partition, Object key, Object value, RecordMetadata recordMetadata) {
+				results.add(recordMetadata);
+				sendResultLatch.countDown();
+			}
+
+			@Override
+			public void onError(String topic, Integer partition, Object key, Object value, Exception exception) {
+				sendResultLatch.countDown();
+			}
+		};
+
+		KafkaProducerContext producerContext = createProducerContext(listener);
+
+		KafkaProducerMessageHandler handler
+				= new KafkaProducerMessageHandler(producerContext);
+
+		handler.setEnableHeaderRouting(false);
+
+		handler.handleMessage(MessageBuilder.withPayload("fooTopic1Header" + suffix)
+				.setHeader(KafkaHeaders.MESSAGE_KEY, "3")
+				.setHeader(KafkaHeaders.TOPIC, TOPIC)
+				.build());
+
+		// even if the header is set to TOPIC2, it should be ignored
+		handler.handleMessage(MessageBuilder.withPayload("fooTopic2Header" + suffix)
+				.setHeader(KafkaHeaders.MESSAGE_KEY, "3")
+				.setHeader(KafkaHeaders.TOPIC, TOPIC2)
+				.build());
+
+		assertTrue(sendResultLatch.await(10, TimeUnit.SECONDS));
+		assertThat(results.size(), equalTo(expectedDeliveryConfirmations));
+
+		producerContext.stop();
+
+		latch.await(10000, TimeUnit.MILLISECONDS);
+		assertThat(latch.getCount(), equalTo(0L));
+		// messages are routed to both topics
+		assertThat(payloadsByTopic.keysView().size(), equalTo(1));
+		assertThat(payloadsByTopic.keysView(), hasItem(TOPIC));
+		assertThat(payloadsByTopic.toMap().get(TOPIC),
+				contains("fooTopic1Header" + suffix, "fooTopic2Header" + suffix));
+		kafkaMessageListenerContainer.stop();
+	}
+
 
 	@Test
 	public void testNoHeader() throws Exception {
@@ -272,7 +361,8 @@ public class OutboundTests {
 		// create the topic
 
 		try {
-			TopicUtils.ensureTopicCreated(kafkaRule.getZookeeperConnectionString(), TOPIC, 1, 1);
+			TopicUtils.ensureTopicCreated(kafkaRule.getZookeeperConnectionString(),
+					TOPIC, 1, 1);
 		}
 		catch (TopicExistsException e) {
 			// do nothing