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

Add partitionOffsets setter in KafkaItemReaderBuilder

Browse Source 
Resolves #3761
This commit is contained in:
Mahmoud Ben Hassine
2020-08-14 21:29:50 +02:00
parent 0b30554ded
commit ce22cddbbc
3 changed files with 34 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/KafkaItemReader.java
View File

@@ -152,7 +152,6 @@ public class KafkaItemReader<K, V> extends AbstractItemStreamItemReader<V> {
* @param partitionOffsets mapping of starting offset in each partition
*/
public void setPartitionOffsets(Map<TopicPartition, Long> partitionOffsets) {
Assert.notNull(partitionOffsets, "partitionOffsets must not be null");
this.partitionOffsets = partitionOffsets;
}

24
spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/builder/KafkaItemReaderBuilder.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2019 the original author or authors.
* Copyright 2019-2020 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.
@@ -20,9 +20,11 @@ import java.time.Duration;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.common.TopicPartition;
import org.springframework.batch.item.kafka.KafkaItemReader;
import org.springframework.util.Assert;
@@ -43,6 +45,8 @@ public class KafkaItemReaderBuilder<K, V> {
private List<Integer> partitions = new ArrayList<>();
private Map<TopicPartition, Long> partitionOffsets;
private Duration pollTimeout = Duration.ofSeconds(30L);
private boolean saveState = true;
@@ -104,6 +108,23 @@ public class KafkaItemReaderBuilder<K, V> {
return this;
}
/**
* Setter for partition offsets. This mapping tells the reader the offset to start
* reading from in each partition. This is optional, defaults to starting from
* offset 0 in each partition. Passing an empty map makes the reader start
* from the offset stored in Kafka for the consumer group ID.
*
* <p><strong>In case of a restart, offsets stored in the execution context
* will take precedence.</strong></p>
*
* @param partitionOffsets mapping of starting offset in each partition
* @return The current instance of the builder.
*/
public KafkaItemReaderBuilder<K, V> partitionOffsets(Map<TopicPartition, Long> partitionOffsets) {
this.partitionOffsets = partitionOffsets;
return this;
}
/**
* A topic name to manually assign to the consumer.
* @param topic name to assign to the consumer
@@ -148,6 +169,7 @@ public class KafkaItemReaderBuilder<K, V> {
reader.setPollTimeout(this.pollTimeout);
reader.setSaveState(this.saveState);
reader.setName(this.name);
reader.setPartitionOffsets(this.partitionOffsets);
return reader;
}
}

12
spring-batch-infrastructure/src/test/java/org/springframework/batch/item/kafka/builder/KafkaItemReaderBuilderTests.java
View File

@@ -1,5 +1,5 @@
/*
* Copyright 2019 the original author or authors.
* Copyright 2019-2020 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.
@@ -18,7 +18,9 @@ package org.springframework.batch.item.kafka.builder;
import java.time.Duration;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import org.apache.kafka.clients.consumer.ConsumerConfig;
@@ -213,6 +215,9 @@ public class KafkaItemReaderBuilderTests {
Duration pollTimeout = Duration.ofSeconds(100);
String topic = "test";
List<Integer> partitions = Arrays.asList(0, 1);
Map<TopicPartition, Long> partitionOffsets = new HashMap<>();
partitionOffsets.put(new TopicPartition(topic, partitions.get(0)), 10L);
partitionOffsets.put(new TopicPartition(topic, partitions.get(1)), 15L);
// when
KafkaItemReader<String, String> reader = new KafkaItemReaderBuilder<String, String>()
@@ -220,6 +225,7 @@ public class KafkaItemReaderBuilderTests {
.consumerProperties(this.consumerProperties)
.topic(topic)
.partitions(partitions)
.partitionOffsets(partitionOffsets)
.pollTimeout(pollTimeout)
.saveState(saveState)
.build();
@@ -234,5 +240,9 @@ public class KafkaItemReaderBuilderTests {
assertEquals(partitions.get(0).intValue(), topicPartitions.get(0).partition());
assertEquals(topic, topicPartitions.get(1).topic());
assertEquals(partitions.get(1).intValue(), topicPartitions.get(1).partition());
Map<TopicPartition, Long> partitionOffsetsMap = (Map<TopicPartition, Long>) ReflectionTestUtils.getField(reader, "partitionOffsets");
assertEquals(2, partitionOffsetsMap.size());
assertEquals(new Long(10), partitionOffsetsMap.get(new TopicPartition(topic, partitions.get(0))));
assertEquals(new Long(15), partitionOffsetsMap.get(new TopicPartition(topic, partitions.get(1))));
}
}