From ce22cddbbccc5a753541821bd5eb878032f2fda5 Mon Sep 17 00:00:00 2001 From: Mahmoud Ben Hassine Date: Fri, 14 Aug 2020 21:29:50 +0200 Subject: [PATCH] Add partitionOffsets setter in KafkaItemReaderBuilder Resolves #3761 --- .../batch/item/kafka/KafkaItemReader.java | 1 - .../kafka/builder/KafkaItemReaderBuilder.java | 24 ++++++++++++++++++- .../builder/KafkaItemReaderBuilderTests.java | 12 +++++++++- 3 files changed, 34 insertions(+), 3 deletions(-) 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 ec761bc3b..e5c2a477c 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 @@ -152,7 +152,6 @@ public class KafkaItemReader extends AbstractItemStreamItemReader { * @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; } diff --git a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/builder/KafkaItemReaderBuilder.java b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/builder/KafkaItemReaderBuilder.java index 901ce2bbd..24c37c9ee 100644 --- a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/builder/KafkaItemReaderBuilder.java +++ b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/kafka/builder/KafkaItemReaderBuilder.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. @@ -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 { private List partitions = new ArrayList<>(); + private Map partitionOffsets; + private Duration pollTimeout = Duration.ofSeconds(30L); private boolean saveState = true; @@ -104,6 +108,23 @@ public class KafkaItemReaderBuilder { 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. + * + *

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

+ * + * @param partitionOffsets mapping of starting offset in each partition + * @return The current instance of the builder. + */ + public KafkaItemReaderBuilder partitionOffsets(Map 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 { reader.setPollTimeout(this.pollTimeout); reader.setSaveState(this.saveState); reader.setName(this.name); + reader.setPartitionOffsets(this.partitionOffsets); return reader; } } diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/kafka/builder/KafkaItemReaderBuilderTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/kafka/builder/KafkaItemReaderBuilderTests.java index 4eed03850..4d725f6d7 100644 --- a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/kafka/builder/KafkaItemReaderBuilderTests.java +++ b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/kafka/builder/KafkaItemReaderBuilderTests.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. @@ -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 partitions = Arrays.asList(0, 1); + Map partitionOffsets = new HashMap<>(); + partitionOffsets.put(new TopicPartition(topic, partitions.get(0)), 10L); + partitionOffsets.put(new TopicPartition(topic, partitions.get(1)), 15L); // when KafkaItemReader reader = new KafkaItemReaderBuilder() @@ -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 partitionOffsetsMap = (Map) 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)))); } }