diff --git a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/KafkaItemReader.java b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/KafkaItemReader.java index 3d4bb49eb..ec761bc3b 100644 --- a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/KafkaItemReader.java +++ b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/KafkaItemReader.java @@ -140,12 +140,30 @@ public class KafkaItemReader extends AbstractItemStreamItemReader { return this.saveState; } + /** + * 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. + * + *

In case of a restart, offsets stored in the execution context + * will take precedence.

+ * + * @param partitionOffsets mapping of starting offset in each partition + */ + public void setPartitionOffsets(Map partitionOffsets) { + Assert.notNull(partitionOffsets, "partitionOffsets must not be null"); + this.partitionOffsets = partitionOffsets; + } + @Override public void open(ExecutionContext executionContext) { this.kafkaConsumer = new KafkaConsumer<>(this.consumerProperties); - this.partitionOffsets = new HashMap<>(); - for (TopicPartition topicPartition : this.topicPartitions) { - this.partitionOffsets.put(topicPartition, 0L); + if (this.partitionOffsets == null) { + this.partitionOffsets = new HashMap<>(); + for (TopicPartition topicPartition : this.topicPartitions) { + this.partitionOffsets.put(topicPartition, 0L); + } } if (this.saveState && executionContext.containsKey(TOPIC_PARTITION_OFFSETS)) { Map offsets = (Map) executionContext.get(TOPIC_PARTITION_OFFSETS); diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/kafka/KafkaItemReaderTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/kafka/KafkaItemReaderTests.java index 75e256353..72458b9f1 100644 --- a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/kafka/KafkaItemReaderTests.java +++ b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/kafka/KafkaItemReaderTests.java @@ -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. @@ -26,6 +26,7 @@ import java.util.concurrent.ExecutionException; import org.apache.kafka.clients.admin.NewTopic; import org.apache.kafka.clients.consumer.ConsumerConfig; +import org.apache.kafka.clients.consumer.OffsetAndMetadata; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.serialization.StringDeserializer; import org.junit.Before; @@ -67,7 +68,9 @@ public class KafkaItemReaderTests { new NewTopic("topic1", 1, (short) 1), new NewTopic("topic2", 2, (short) 1), new NewTopic("topic3", 1, (short) 1), - new NewTopic("topic4", 2, (short) 1) + new NewTopic("topic4", 2, (short) 1), + new NewTopic("topic5", 1, (short) 1), + new NewTopic("topic6", 1, (short) 1) ); } @@ -212,6 +215,89 @@ public class KafkaItemReaderTests { this.reader.close(); } + @Test + public void testReadFromSinglePartitionFromCustomOffset() { + this.template.setDefaultTopic("topic5"); + this.template.sendDefault("val0"); // <-- offset 0 + this.template.sendDefault("val1"); // <-- offset 1 + this.template.sendDefault("val2"); // <-- offset 2 + this.template.sendDefault("val3"); // <-- offset 3 + + this.reader = new KafkaItemReader<>(this.consumerProperties, "topic5", 0); + + // specify which offset to start from + Map partitionOffsets = new HashMap<>(); + partitionOffsets.put(new TopicPartition("topic5", 0), 2L); + this.reader.setPartitionOffsets(partitionOffsets); + + this.reader.setPollTimeout(Duration.ofSeconds(1)); + this.reader.open(new ExecutionContext()); + + String item = this.reader.read(); + assertThat(item, is("val2")); + + item = this.reader.read(); + assertThat(item, is("val3")); + + item = this.reader.read(); + assertNull(item); + + this.reader.close(); + } + + @Test + public void testReadFromSinglePartitionFromTheOffsetStoredInKafka() throws Exception { + // first run: read a topic from the beginning + + this.template.setDefaultTopic("topic6"); + this.template.sendDefault("val0"); // <-- offset 0 + this.template.sendDefault("val1"); // <-- offset 1 + + this.reader = new KafkaItemReader<>(this.consumerProperties, "topic6", 0); + this.reader.setPollTimeout(Duration.ofSeconds(1)); + this.reader.open(new ExecutionContext()); + + String item = this.reader.read(); + assertThat(item, is("val0")); + + item = this.reader.read(); + assertThat(item, is("val1")); + + item = this.reader.read(); + assertNull(item); + + this.reader.close(); + + // The offset stored in Kafka should be equal to 2 at this point + OffsetAndMetadata currentOffset = KafkaTestUtils.getCurrentOffset( + embeddedKafka.getEmbeddedKafka().getBrokersAsString(), + "1", "topic6", + 0); + assertEquals(2, currentOffset.offset()); + + // second run (with same consumer group ID): new messages arrived since the last run. + + this.template.sendDefault("val2"); // <-- offset 2 + this.template.sendDefault("val3"); // <-- offset 3 + + this.reader = new KafkaItemReader<>(this.consumerProperties, "topic6", 0); + // Passing an empty map means the reader should start from the offset stored in Kafka (offset 2 in this case) + this.reader.setPartitionOffsets(new HashMap<>()); + this.reader.setPollTimeout(Duration.ofSeconds(1)); + this.reader.open(new ExecutionContext()); + + item = this.reader.read(); + assertThat(item, is("val2")); + + item = this.reader.read(); + assertThat(item, is("val3")); + + item = this.reader.read(); + assertNull(item); + + this.reader.close(); + } + @Test public void testReadFromMultiplePartitions() { this.template.setDefaultTopic("topic2");