Spring/spring-integration
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Issue#67: fixing custom paritioner issue

Browse Source
This commit is contained in:
Abhijith Prabhakar
2015-10-08 14:46:07 -04:00
committed by Artem Bilan
parent fd93f91a78
commit 32715799d9
5 changed files with 78 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -42,7 +42,7 @@ public class KafkaProducerMessageHandler extends AbstractMessageHandler {
private volatile Expression messageKeyExpression;
private volatile Expression partitionExpression;
private volatile Expression partitionIdExpression;
@SuppressWarnings("unchecked")
public KafkaProducerMessageHandler(final KafkaProducerContext kafkaProducerContext) {
@@ -57,8 +57,8 @@ public class KafkaProducerMessageHandler extends AbstractMessageHandler {
this.messageKeyExpression = messageKeyExpression;
}
public void setPartitionExpression(Expression partitionExpression) {
this.partitionExpression = partitionExpression;
public void setPartitionIdExpression(Expression partitionIdExpression) {
this.partitionIdExpression = partitionIdExpression;
}
public KafkaProducerContext getKafkaProducerContext() {
@@ -77,8 +77,8 @@ public class KafkaProducerMessageHandler extends AbstractMessageHandler {
this.topicExpression.getValue(this.evaluationContext, message, String.class)
: message.getHeaders().get(KafkaHeaders.TOPIC, String.class);
Integer partitionId = this.partitionExpression != null ?
this.partitionExpression.getValue(this.evaluationContext, message, Integer.class)
Integer partitionId = this.partitionIdExpression != null ?
this.partitionIdExpression.getValue(this.evaluationContext, message, Integer.class)
: message.getHeaders().get(KafkaHeaders.PARTITION_ID, Integer.class);
Object messageKey = this.messageKeyExpression != null

8
spring-integration-kafka/src/main/java/org/springframework/integration/kafka/support/ProducerConfiguration.java
View File

@@ -80,7 +80,13 @@ public class ProducerConfiguration<K, V> {
public Future<RecordMetadata> send(String topic, Integer partition, K messageKey, V messagePayload) {
String targetTopic = StringUtils.hasText(topic) ? topic : this.producerMetadata.getTopic();
Future<RecordMetadata> future = this.producer.send(new ProducerRecord<>(targetTopic, partition, messageKey, messagePayload));
//If partition is sent by producer context then that takes precedence over custom partitioner.
if (partition == null && this.getProducerMetadata().getPartitioner() != null) {
partition = this.getProducerMetadata().getPartitioner().partition(messageKey,
this.producer.partitionsFor(targetTopic).size());
}
Future<RecordMetadata> future =
this.producer.send(new ProducerRecord<>(targetTopic, partition, messageKey, messagePayload));
if (!producerMetadata.isSync()) {
return future;

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

@@ -19,7 +19,9 @@
channel="inputToKafka"
order="3"
topic="foo"
message-key-expression="'bar'">
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>

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

@@ -59,6 +59,7 @@ public class KafkaOutboundAdapterParserTests {
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"));
KafkaProducerContext producerContext = messageHandler.getKafkaProducerContext();
assertNotNull(producerContext);
assertEquals(producerContext.getProducerConfigurations().size(), 2);

63
spring-integration-kafka/src/test/java/org/springframework/integration/kafka/support/ProducerConfigurationTests.java
View File

@@ -18,9 +18,12 @@ package org.springframework.integration.kafka.support;
import java.io.ByteArrayInputStream;
import java.io.ObjectInputStream;
import java.util.ArrayList;
import kafka.producer.Partitioner;
import org.apache.kafka.clients.producer.Producer;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.serialization.ByteArraySerializer;
import org.apache.kafka.common.serialization.StringSerializer;
import org.junit.Assert;
@@ -39,7 +42,6 @@ import org.springframework.messaging.Message;
import org.springframework.messaging.support.GenericMessage;
import kafka.serializer.DefaultEncoder;
import kafka.serializer.StringEncoder;
/**
* @author Soby Chacko
@@ -242,6 +244,65 @@ public class ProducerConfigurationTests {
Assert.assertEquals(capturedKeyMessage.topic(), "test");
}
@Test
@SuppressWarnings("unchecked")
public void testSendMessageWithNonDefaultKeyAndValueEncodersSpecifyingPartition() throws Exception {
final ProducerMetadata<String, String> producerMetadata = new ProducerMetadata<String, String>("test", String.class, String.class, new StringSerializer(), new StringSerializer());
final Producer<String, String> producer = Mockito.mock(Producer.class);
final ProducerConfiguration<String, String> configuration =
new ProducerConfiguration<String, String>(producerMetadata, producer);
configuration.send("test", 1, "key", "test message");
Mockito.verify(producer, Mockito.times(1)).send(Mockito.any(ProducerRecord.class));
final ArgumentCaptor<ProducerRecord<String, String>> argument =
(ArgumentCaptor<ProducerRecord<String, String>>) (Object)
ArgumentCaptor.forClass(ProducerRecord.class);
Mockito.verify(producer).send(argument.capture());
final ProducerRecord<String, String> capturedKeyMessage = argument.getValue();
Assert.assertEquals(capturedKeyMessage.key(), "key");
Assert.assertEquals(capturedKeyMessage.value(), "test message");
Assert.assertEquals(capturedKeyMessage.topic(), "test");
Assert.assertEquals(capturedKeyMessage.partition().intValue(), 1);
}
@Test
@SuppressWarnings("unchecked")
public void testSendMessageWithNonDefaultKeyAndValueEncodersCustomPartitioner() throws Exception {
Partitioner customPartitioner = Mockito.mock(Partitioner.class);
final ProducerMetadata<String, String> producerMetadata = new ProducerMetadata<String, String>("test", String.class, String.class, new StringSerializer(), new StringSerializer());
producerMetadata.setPartitioner(customPartitioner);
final Producer<String, String> producer = Mockito.mock(Producer.class);
final ProducerConfiguration<String, String> configuration =
new ProducerConfiguration<String, String>(producerMetadata, producer);
ArrayList<PartitionInfo> partitionInfos = new ArrayList<>();
partitionInfos.add(Mockito.mock(PartitionInfo.class));
Mockito.when(producer.partitionsFor("test")).thenReturn(partitionInfos);
Mockito.when(customPartitioner.partition("key", 1)).thenReturn(4);
configuration.send("test", null, "key", "test message");
Mockito.verify(producer, Mockito.times(1)).send(Mockito.any(ProducerRecord.class));
final ArgumentCaptor<ProducerRecord<String, String>> argument =
(ArgumentCaptor<ProducerRecord<String, String>>) (Object)
ArgumentCaptor.forClass(ProducerRecord.class);
Mockito.verify(producer).send(argument.capture());
final ProducerRecord<String, String> capturedKeyMessage = argument.getValue();
Assert.assertEquals(capturedKeyMessage.key(), "key");
Assert.assertEquals(capturedKeyMessage.value(), "test message");
Assert.assertEquals(capturedKeyMessage.topic(), "test");
Assert.assertEquals(capturedKeyMessage.partition().intValue(), 4);
}
/**
* User does not set an explicit key/value encoder, but send non-serializable object for both key/value
*/