Add ability to start reading from a custom offset in KafkaItemReader

Issue #737

Author: fmbenhassine

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

@@ -140,12 +140,30 @@ public boolean isSaveState() {
 		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.
+	 * 
+	 * <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
+	 */
+	public void setPartitionOffsets(Map<TopicPartition, Long> 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<TopicPartition, Long> offsets = (Map<TopicPartition, Long>) executionContext.get(TOPIC_PARTITION_OFFSETS);

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 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 static void setUpTopics() {
 				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 void testReadFromSinglePartition() {
 		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<TopicPartition, Long> 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");