diff --git a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/redis/RedisItemReader.java b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/redis/RedisItemReader.java index f5142fa39a..8857f05378 100644 --- a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/redis/RedisItemReader.java +++ b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/redis/RedisItemReader.java @@ -1,5 +1,5 @@ /* - * Copyright 2023 the original author or authors. + * Copyright 2023-2025 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,21 +18,27 @@ import org.springframework.batch.item.ExecutionContext; import org.springframework.batch.item.ItemStreamException; import org.springframework.batch.item.ItemStreamReader; -import org.springframework.data.redis.core.Cursor; -import org.springframework.data.redis.core.RedisTemplate; -import org.springframework.data.redis.core.ScanOptions; +import org.springframework.dao.DataAccessException; +import org.springframework.data.redis.core.*; import org.springframework.util.Assert; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Deque; +import java.util.List; + /** * Item reader for Redis based on Spring Data Redis. Uses a {@link RedisTemplate} to query * data. The user should provide a {@link ScanOptions} to specify the set of keys to - * query. + * query. The {@code fetchSize} property controls how many items are fetched from Redis in + * a single pipeline round-trip for efficiency. * *

* The implementation is not thread-safe and not restartable. *

* * @author Mahmoud Ben Hassine + * @author Hyunwoo Jung * @since 5.1 * @param type of keys * @param type of values @@ -43,13 +49,20 @@ public class RedisItemReader implements ItemStreamReader { private final ScanOptions scanOptions; + private final int fetchSize; + + private final Deque buffer; + private Cursor cursor; - public RedisItemReader(RedisTemplate redisTemplate, ScanOptions scanOptions) { + public RedisItemReader(RedisTemplate redisTemplate, ScanOptions scanOptions, int fetchSize) { Assert.notNull(redisTemplate, "redisTemplate must not be null"); Assert.notNull(scanOptions, "scanOptions must no be null"); + Assert.isTrue(fetchSize > 0, "fetchSize must be greater than 0"); this.redisTemplate = redisTemplate; this.scanOptions = scanOptions; + this.fetchSize = fetchSize; + this.buffer = new ArrayDeque<>(); } @Override @@ -59,13 +72,11 @@ public void open(ExecutionContext executionContext) throws ItemStreamException { @Override public V read() throws Exception { - if (this.cursor.hasNext()) { - K nextKey = this.cursor.next(); - return this.redisTemplate.opsForValue().get(nextKey); - } - else { - return null; + if (this.buffer.isEmpty()) { + fetchNext(); } + + return this.buffer.pollFirst(); } @Override @@ -73,4 +84,33 @@ public void close() throws ItemStreamException { this.cursor.close(); } + private void fetchNext() { + List keys = new ArrayList<>(); + while (this.cursor.hasNext() && keys.size() < this.fetchSize) { + keys.add(this.cursor.next()); + } + + if (keys.isEmpty()) { + return; + } + + @SuppressWarnings("unchecked") + List items = (List) this.redisTemplate.executePipelined(sessionCallback(keys)); + + this.buffer.addAll(items); + } + + private SessionCallback sessionCallback(List keys) { + return new SessionCallback<>() { + @Override + public Object execute(RedisOperations operations) throws DataAccessException { + for (K key : keys) { + operations.opsForValue().get(key); + } + + return null; + } + }; + } + } diff --git a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/redis/RedisItemWriter.java b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/redis/RedisItemWriter.java index c9b0ae3ee3..a5ec5ddef5 100644 --- a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/redis/RedisItemWriter.java +++ b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/redis/RedisItemWriter.java @@ -1,5 +1,5 @@ /* - * Copyright 2023 the original author or authors. + * Copyright 2023-2025 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,9 +18,16 @@ import org.springframework.batch.item.ItemWriter; import org.springframework.batch.item.KeyValueItemWriter; +import org.springframework.dao.DataAccessException; +import org.springframework.data.redis.core.RedisOperations; import org.springframework.data.redis.core.RedisTemplate; +import org.springframework.data.redis.core.SessionCallback; +import org.springframework.data.util.Pair; import org.springframework.util.Assert; +import java.util.ArrayList; +import java.util.List; + /** *

* An {@link ItemWriter} implementation for Redis using a {@link RedisTemplate} . @@ -28,20 +35,18 @@ * * @author Santiago Molano * @author Mahmoud Ben Hassine + * @author Hyunwoo Jung * @since 5.1 */ public class RedisItemWriter extends KeyValueItemWriter { private RedisTemplate redisTemplate; + private final List> buffer = new ArrayList<>(); + @Override protected void writeKeyValue(K key, T value) { - if (this.delete) { - this.redisTemplate.delete(key); - } - else { - this.redisTemplate.opsForValue().set(key, value); - } + this.buffer.add(Pair.of(key, value)); } @Override @@ -49,6 +54,17 @@ protected void init() { Assert.notNull(this.redisTemplate, "RedisTemplate must not be null"); } + @Override + protected void flush() throws Exception { + if (this.buffer.isEmpty()) { + return; + } + + this.redisTemplate.executePipelined(sessionCallback()); + + this.buffer.clear(); + } + /** * Set the {@link RedisTemplate} to use. * @param redisTemplate the template to use @@ -57,4 +73,33 @@ public void setRedisTemplate(RedisTemplate redisTemplate) { this.redisTemplate = redisTemplate; } + private SessionCallback sessionCallback() { + return new SessionCallback<>() { + + @SuppressWarnings("unchecked") + @Override + public Object execute(RedisOperations operations) throws DataAccessException { + if (RedisItemWriter.this.delete) { + executeDeleteOperations(operations); + } + else { + executeSetOperations(operations); + } + return null; + } + }; + } + + private void executeDeleteOperations(RedisOperations operations) { + for (Pair item : this.buffer) { + operations.delete(item.getFirst()); + } + } + + private void executeSetOperations(RedisOperations operations) { + for (Pair item : this.buffer) { + operations.opsForValue().set(item.getFirst(), item.getSecond()); + } + } + } diff --git a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/redis/builder/RedisItemReaderBuilder.java b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/redis/builder/RedisItemReaderBuilder.java index 7b00778090..c92525b259 100644 --- a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/redis/builder/RedisItemReaderBuilder.java +++ b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/redis/builder/RedisItemReaderBuilder.java @@ -1,5 +1,5 @@ /* - * Copyright 2023 the original author or authors. + * Copyright 2023-2025 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. @@ -23,6 +23,7 @@ * Builder for {@link RedisItemReader}. * * @author Mahmoud Ben Hassine + * @author Hyunwoo Jung * @since 5.1 * @param type of keys * @param type of values @@ -33,6 +34,8 @@ public class RedisItemReaderBuilder { private ScanOptions scanOptions; + private int fetchSize; + /** * Set the {@link RedisTemplate} to use in the reader. * @param redisTemplate the template to use @@ -53,12 +56,22 @@ public RedisItemReaderBuilder scanOptions(ScanOptions scanOptions) { return this; } + /** + * Set the fetchSize to how many items from Redis in a single round-trip. + * @param fetchSize the number of items to fetch per pipeline execution + * @return the current builder instance for fluent chaining + */ + public RedisItemReaderBuilder fetchSize(int fetchSize) { + this.fetchSize = fetchSize; + return this; + } + /** * Build a new {@link RedisItemReader}. * @return a new item reader */ public RedisItemReader build() { - return new RedisItemReader<>(this.redisTemplate, this.scanOptions); + return new RedisItemReader<>(this.redisTemplate, this.scanOptions, this.fetchSize); } } diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/redis/RedisItemReaderIntegrationTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/redis/RedisItemReaderIntegrationTests.java index 66e733fcfb..c69e5b7ce3 100644 --- a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/redis/RedisItemReaderIntegrationTests.java +++ b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/redis/RedisItemReaderIntegrationTests.java @@ -82,7 +82,7 @@ void testRead(RedisConnectionFactory connectionFactory) throws Exception { RedisTemplate redisTemplate = setUpRedisTemplate(connectionFactory); ScanOptions scanOptions = ScanOptions.scanOptions().match("person:*").count(10).build(); - this.reader = new RedisItemReader<>(redisTemplate, scanOptions); + this.reader = new RedisItemReader<>(redisTemplate, scanOptions, 10); this.reader.open(new ExecutionContext()); diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/redis/RedisItemReaderTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/redis/RedisItemReaderTests.java index 2848a24456..9290fa4b30 100644 --- a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/redis/RedisItemReaderTests.java +++ b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/redis/RedisItemReaderTests.java @@ -1,5 +1,5 @@ /* - * Copyright 2023 the original author or authors. + * Copyright 2023-2025 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. @@ -16,20 +16,30 @@ package org.springframework.batch.item.redis; import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; import org.mockito.Answers; import org.mockito.Mock; -import org.mockito.Mockito; import org.mockito.junit.jupiter.MockitoExtension; - import org.springframework.batch.item.ExecutionContext; import org.springframework.data.redis.core.Cursor; import org.springframework.data.redis.core.RedisTemplate; import org.springframework.data.redis.core.ScanOptions; +import org.springframework.data.redis.core.SessionCallback; + +import java.util.ArrayList; +import java.util.List; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.when; +/** + * @author Mahmoud Ben Hassine + * @author Hyunwoo Jung + */ @ExtendWith(MockitoExtension.class) -public class RedisItemReaderTests { +class RedisItemReaderTests { @Mock(answer = Answers.RETURNS_DEEP_STUBS) private RedisTemplate redisTemplate; @@ -40,15 +50,36 @@ public class RedisItemReaderTests { @Mock private Cursor cursor; + private List results; + + @BeforeEach + void setUp() { + this.results = new ArrayList<>(); + + when(this.redisTemplate.executePipelined(any(SessionCallback.class))).thenAnswer(invocation -> { + SessionCallback sessionCallback = invocation.getArgument(0); + sessionCallback.execute(this.redisTemplate); + return this.results; + }); + } + @Test void testRead() throws Exception { // given - Mockito.when(this.redisTemplate.scan(this.scanOptions)).thenReturn(this.cursor); - Mockito.when(this.cursor.hasNext()).thenReturn(true, true, false); - Mockito.when(this.cursor.next()).thenReturn("person:1", "person:2"); - Mockito.when(this.redisTemplate.opsForValue().get("person:1")).thenReturn("foo"); - Mockito.when(this.redisTemplate.opsForValue().get("person:2")).thenReturn("bar"); - RedisItemReader redisItemReader = new RedisItemReader<>(this.redisTemplate, this.scanOptions); + when(this.redisTemplate.scan(this.scanOptions)).thenReturn(this.cursor); + when(this.cursor.hasNext()).thenReturn(true, true, false); + when(this.cursor.next()).thenReturn("person:1", "person:2"); + when(this.redisTemplate.opsForValue().get("person:1")).thenAnswer(invocation -> { + results.add("foo"); + return null; + }); + when(this.redisTemplate.opsForValue().get("person:2")).thenAnswer(invocation -> { + results.add("bar"); + return null; + }); + + RedisItemReader redisItemReader = new RedisItemReader<>(this.redisTemplate, this.scanOptions, + 10); redisItemReader.open(new ExecutionContext()); // when diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/redis/RedisItemWriterTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/redis/RedisItemWriterTests.java index ebf9ae8f87..51a36d9cb2 100644 --- a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/redis/RedisItemWriterTests.java +++ b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/redis/RedisItemWriterTests.java @@ -1,5 +1,5 @@ /* - * Copyright 2023 the original author or authors. + * Copyright 2023-2025 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. @@ -22,15 +22,21 @@ import org.mockito.Answers; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; - import org.springframework.batch.item.Chunk; import org.springframework.core.convert.converter.Converter; import org.springframework.data.redis.core.RedisTemplate; +import org.springframework.data.redis.core.SessionCallback; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +/** + * @author Mahmoud Ben Hassine + * @author Hyunwoo Jung + */ @ExtendWith(MockitoExtension.class) -public class RedisItemWriterTests { +class RedisItemWriterTests { @Mock(answer = Answers.RETURNS_DEEP_STUBS) private RedisTemplate redisTemplate; @@ -38,16 +44,16 @@ public class RedisItemWriterTests { private RedisItemWriter redisItemWriter; @BeforeEach - public void setup() { + void setup() { this.redisItemWriter = new RedisItemWriter<>(); this.redisItemWriter.setRedisTemplate(this.redisTemplate); this.redisItemWriter.setItemKeyMapper(new RedisItemKeyMapper()); - } - @Test - void shouldWriteToRedisDatabaseUsingKeyValue() { - this.redisItemWriter.writeKeyValue("oneKey", "oneValue"); - verify(this.redisTemplate.opsForValue()).set("oneKey", "oneValue"); + when(this.redisTemplate.executePipelined(any(SessionCallback.class))).thenAnswer(invocation -> { + SessionCallback sessionCallback = invocation.getArgument(0); + sessionCallback.execute(this.redisTemplate); + return null; + }); } @Test @@ -58,6 +64,15 @@ void shouldWriteAllItemsToRedis() throws Exception { verify(this.redisTemplate.opsForValue()).set(items.getItems().get(1), items.getItems().get(1)); } + @Test + void shouldDeleteAllItemsToRedis() throws Exception { + this.redisItemWriter.setDelete(true); + Chunk items = new Chunk<>("val1", "val2"); + this.redisItemWriter.write(items); + verify(this.redisTemplate).delete(items.getItems().get(0)); + verify(this.redisTemplate).delete(items.getItems().get(0)); + } + static class RedisItemKeyMapper implements Converter { @Override diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/redis/builder/RedisItemReaderBuilderTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/redis/builder/RedisItemReaderBuilderTests.java index 3192aa784c..02e58a0b19 100644 --- a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/redis/builder/RedisItemReaderBuilderTests.java +++ b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/redis/builder/RedisItemReaderBuilderTests.java @@ -1,5 +1,5 @@ /* - * Copyright 2023 the original author or authors. + * Copyright 2023-2025 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. @@ -30,8 +30,9 @@ * Test class for {@link RedisItemReaderBuilder}. * * @author Mahmoud Ben Hassine + * @author Hyunwoo Jung */ -public class RedisItemReaderBuilderTests { +class RedisItemReaderBuilderTests { @Test void testRedisItemReaderCreation() { @@ -43,12 +44,14 @@ void testRedisItemReaderCreation() { RedisItemReader reader = new RedisItemReaderBuilder() .redisTemplate(redisTemplate) .scanOptions(scanOptions) + .fetchSize(10) .build(); // then assertNotNull(reader); assertEquals(redisTemplate, ReflectionTestUtils.getField(reader, "redisTemplate")); assertEquals(scanOptions, ReflectionTestUtils.getField(reader, "scanOptions")); + assertEquals(10, ReflectionTestUtils.getField(reader, "fetchSize")); } }