From e516cbd4a0352fc3c6a05fcfa04d4f9f80f7f11e Mon Sep 17 00:00:00 2001 From: dsyer Date: Tue, 7 Jul 2009 08:45:10 +0000 Subject: [PATCH] RESOLVED - issue BATCH-1253: Ability to throttle the number of rows returned by a HibernateCursorItemReader --- ...tractItemCountingItemStreamItemReader.java | 81 ++++++++++++++----- .../item/file/FlatFileItemReaderTests.java | 78 ++++++++++++++++++ 2 files changed, 140 insertions(+), 19 deletions(-) diff --git a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/support/AbstractItemCountingItemStreamItemReader.java b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/support/AbstractItemCountingItemStreamItemReader.java index b145539c6..26621a966 100644 --- a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/support/AbstractItemCountingItemStreamItemReader.java +++ b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/support/AbstractItemCountingItemStreamItemReader.java @@ -34,7 +34,8 @@ import org.springframework.util.Assert; * * @author Robert Kasanicky */ -public abstract class AbstractItemCountingItemStreamItemReader implements ItemReader, ItemStream { +public abstract class AbstractItemCountingItemStreamItemReader implements + ItemReader, ItemStream { private static final String READ_COUNT = "read.count"; @@ -50,6 +51,7 @@ public abstract class AbstractItemCountingItemStreamItemReader implements Ite /** * Read next item from input. + * * @return item * @throws Exception */ @@ -76,8 +78,9 @@ public abstract class AbstractItemCountingItemStreamItemReader implements Ite } } - public final T read() throws Exception, UnexpectedInputException, ParseException { - if (currentItemCount >= maxItemCount-1) { + public final T read() throws Exception, UnexpectedInputException, + ParseException { + if (currentItemCount >= maxItemCount) { return null; } currentItemCount++; @@ -88,42 +91,69 @@ public abstract class AbstractItemCountingItemStreamItemReader implements Ite return currentItemCount; } - protected void setCurrentItemCount(int count) { + /** + * The index of the item to start reading from. If the + * {@link ExecutionContext} contains a key [name].read.count + * (where [name] is the name of this component) the value from + * the {@link ExecutionContext} will be used in preference. + * + * @see #setName(String) + * + * @param count + * the value of the current item count + */ + public void setCurrentItemCount(int count) { this.currentItemCount = count; } + /** + * The maximum index of the items to be read. If the + * {@link ExecutionContext} contains a key [name].read.count.max + * (where [name] is the name of this component) the value from + * the {@link ExecutionContext} will be used in preference. + * + * @see #setName(String) + * + * @param count + * the value of the maximum item count + */ + public void setMaxItemCount(int count) { + this.maxItemCount = count; + } + public void close() throws ItemStreamException { currentItemCount = 0; try { doClose(); - } - catch (Exception e) { + } catch (Exception e) { throw new ItemStreamException("Error while closing item reader", e); } } - public void open(ExecutionContext executionContext) throws ItemStreamException { + public void open(ExecutionContext executionContext) + throws ItemStreamException { try { doOpen(); - } - catch (Exception e) { + } catch (Exception e) { throw new ItemStreamException("Failed to initialize the reader", e); } if (executionContext.containsKey(ecSupport.getKey(READ_COUNT_MAX))) { - maxItemCount = executionContext.getInt(ecSupport.getKey(READ_COUNT_MAX)); + maxItemCount = executionContext.getInt(ecSupport + .getKey(READ_COUNT_MAX)); } if (executionContext.containsKey(ecSupport.getKey(READ_COUNT))) { - int itemCount = executionContext.getInt(ecSupport.getKey(READ_COUNT)); + int itemCount = executionContext.getInt(ecSupport + .getKey(READ_COUNT)); if (itemCount < maxItemCount) { try { jumpToItem(itemCount); - } - catch (Exception e) { - throw new ItemStreamException("Could not move to stored position on restart", e); + } catch (Exception e) { + throw new ItemStreamException( + "Could not move to stored position on restart", e); } } currentItemCount = itemCount; @@ -132,17 +162,29 @@ public abstract class AbstractItemCountingItemStreamItemReader implements Ite } - public void update(ExecutionContext executionContext) throws ItemStreamException { + public void update(ExecutionContext executionContext) + throws ItemStreamException { if (saveState) { - Assert.notNull(executionContext, "ExecutionContext must not be null"); - executionContext.putInt(ecSupport.getKey(READ_COUNT), currentItemCount); + Assert.notNull(executionContext, + "ExecutionContext must not be null"); + executionContext.putInt(ecSupport.getKey(READ_COUNT), + currentItemCount); if (maxItemCount < Integer.MAX_VALUE) { - executionContext.putInt(ecSupport.getKey(READ_COUNT_MAX), maxItemCount); + executionContext.putInt(ecSupport.getKey(READ_COUNT_MAX), + maxItemCount); } } } + /** + * The name of the component which will be used as a stem for keys in the + * {@link ExecutionContext}. Subclasses should provide a default value, e.g. + * the short form of the class name. + * + * @param name + * the name for the component + */ public void setName(String name) { ecSupport.setName(name); } @@ -152,7 +194,8 @@ public abstract class AbstractItemCountingItemStreamItemReader implements Ite * {@link ExecutionContext}. Only switch this to false if you don't want to * save any state from this stream, and you don't need it to be restartable. * - * @param saveState flag value (default true). + * @param saveState + * flag value (default true). */ public void setSaveState(boolean saveState) { this.saveState = saveState; diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/file/FlatFileItemReaderTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/file/FlatFileItemReaderTests.java index 2011f1bae..b411ba397 100644 --- a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/file/FlatFileItemReaderTests.java +++ b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/file/FlatFileItemReaderTests.java @@ -103,6 +103,84 @@ public class FlatFileItemReaderTests { assertEquals(4, executionContext.getInt(ClassUtils.getShortName(FlatFileItemReader.class) + ".read.count")); } + @Test + public void testCurrentItemCount() throws Exception { + + reader.setCurrentItemCount(2); + reader.open(executionContext); + + // read some records + reader.read(); + reader.read(); + // get restart data + reader.update(executionContext); + + assertEquals(4, executionContext.getInt(ClassUtils.getShortName(FlatFileItemReader.class) + ".read.count")); + // close input + reader.close(); + + } + + @Test + public void testMaxItemCount() throws Exception { + + reader.setMaxItemCount(2); + reader.open(executionContext); + + // read some records + reader.read(); + reader.read(); + // get restart data + reader.update(executionContext); + assertNull(reader.read()); + + assertEquals(2, executionContext.getInt(ClassUtils.getShortName(FlatFileItemReader.class) + ".read.count")); + // close input + reader.close(); + + } + + @Test + public void testMaxItemCountFromContext() throws Exception { + + reader.setMaxItemCount(2); + executionContext.putInt(reader.getClass().getSimpleName()+".read.count.max", Integer.MAX_VALUE); + reader.open(executionContext); + // read some records + reader.read(); + reader.read(); + assertNotNull(reader.read()); + // close input + reader.close(); + + } + + @Test + public void testCurrentItemCountFromContext() throws Exception { + + reader.setCurrentItemCount(2); + executionContext.putInt(reader.getClass().getSimpleName()+".read.count", 3); + reader.open(executionContext); + // read some records + assertEquals("testLine4", reader.read()); + // close input + reader.close(); + + } + + @Test + public void testMaxAndCurrentItemCount() throws Exception { + + reader.setMaxItemCount(2); + reader.setCurrentItemCount(2); + reader.open(executionContext); + // read some records + assertNull(reader.read()); + // close input + reader.close(); + + } + @Test public void testNonExistentResource() throws Exception {