diff --git a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/ItemCountAware.java b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/ItemCountAware.java new file mode 100644 index 000000000..5c6a5e2bb --- /dev/null +++ b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/ItemCountAware.java @@ -0,0 +1,15 @@ +package org.springframework.batch.item; + +import org.springframework.batch.item.support.AbstractItemCountingItemStreamItemReader; + +/** + * Marker interface indicating that an item should have the item count set on it. Typically used within + * an {@link AbstractItemCountingItemStreamItemReader}. + * + * @author Jimmy Praet + */ +public interface ItemCountAware { + + void setItemCount(int count); + +} 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 6d8b28d17..b8443cf50 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 @@ -17,6 +17,7 @@ package org.springframework.batch.item.support; import org.springframework.batch.item.ExecutionContext; +import org.springframework.batch.item.ItemCountAware; import org.springframework.batch.item.ItemReader; import org.springframework.batch.item.ItemStreamException; import org.springframework.batch.item.ParseException; @@ -79,7 +80,11 @@ public abstract class AbstractItemCountingItemStreamItemReader extends Abstra return null; } currentItemCount++; - return doRead(); + T item = doRead(); + if(item instanceof ItemCountAware) { + ((ItemCountAware) item).setItemCount(currentItemCount); + } + return item; } protected int getCurrentItemCount() { 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 4a78d67f1..117c97b65 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 @@ -12,6 +12,7 @@ import java.io.InputStream; import org.junit.Before; import org.junit.Test; import org.springframework.batch.item.ExecutionContext; +import org.springframework.batch.item.ItemCountAware; import org.springframework.batch.item.ItemStreamException; import org.springframework.batch.item.file.mapping.PassThroughLineMapper; import org.springframework.batch.item.file.separator.RecordSeparatorPolicy; @@ -31,6 +32,8 @@ public class FlatFileItemReaderTests { private String TEST_STRING = "FlatFileInputTemplate-TestData"; private FlatFileItemReader reader = new FlatFileItemReader(); + + private FlatFileItemReader itemReader = new FlatFileItemReader(); private ExecutionContext executionContext = new ExecutionContext(); @@ -39,6 +42,9 @@ public class FlatFileItemReaderTests { reader.setResource(getInputResource("testLine1\ntestLine2\ntestLine3\ntestLine4\ntestLine5\ntestLine6")); reader.setLineMapper(new PassThroughLineMapper()); + + itemReader.setResource(getInputResource("testLine1\ntestLine2\ntestLine3\ntestLine4\ntestLine5\ntestLine6")); + itemReader.setLineMapper(new ItemLineMapper()); } @Test @@ -427,6 +433,66 @@ public class FlatFileItemReaderTests { assertEquals("Parsing error at line: 2 in resource=[resource loaded from byte array], input=[testLine2]", expected.getMessage()); } } + + @Test + public void testItemCountAware() throws Exception { + itemReader.open(executionContext); + Item item1 = itemReader.read(); + assertEquals("testLine1", item1.getValue()); + assertEquals(1, item1.getItemCount()); + Item item2 = itemReader.read(); + assertEquals("testLine2", item2.getValue()); + assertEquals(2, item2.getItemCount()); + itemReader.update(executionContext); + itemReader.close(); + + itemReader.open(executionContext); + Item item3 = itemReader.read(); + assertEquals("testLine3", item3.getValue()); + assertEquals(3, item3.getItemCount()); + } + + @Test + public void testItemCountAwareMultiLine() throws Exception { + itemReader.setRecordSeparatorPolicy(new RecordSeparatorPolicy() { + + // 1 record = 2 lines + boolean pair = true; + + @Override + public boolean isEndOfRecord(String line) { + if (StringUtils.hasText(line)) { + pair = !pair; + } + return pair; + } + + @Override + public String postProcess(String record) { + return StringUtils.hasText(record) ? record : null; + } + + @Override + public String preProcess(String record) { + return record; + } + }); + + itemReader.open(executionContext); + Item item1 = itemReader.read(); + assertEquals("testLine1testLine2", item1.getValue()); + assertEquals(1, item1.getItemCount()); + Item item2 = itemReader.read(); + assertEquals("testLine3testLine4", item2.getValue()); + assertEquals(2, item2.getItemCount()); + itemReader.update(executionContext); + itemReader.close(); + + itemReader.open(executionContext); + Item item3 = itemReader.read(); + assertEquals("testLine5testLine6", item3.getValue()); + assertEquals(3, item3.getItemCount()); + } private Resource getInputResource(String input) { return new ByteArrayResource(input.getBytes()); @@ -452,4 +518,42 @@ public class FlatFileItemReaderTests { return null; } } + + private static class Item implements ItemCountAware { + + private String value; + + private int itemCount; + + public Item(String value) { + this.value = value; + } + + public void setValue(String value) { + this.value = value; + } + + public String getValue() { + return value; + } + + @Override + public void setItemCount(int count) { + this.itemCount = count; + } + + public int getItemCount() { + return itemCount; + } + + } + + private static final class ItemLineMapper implements LineMapper { + + @Override + public Item mapLine(String line, int lineNumber) throws Exception { + return new Item(line); + } + + } } diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/xml/StaxEventItemReaderTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/xml/StaxEventItemReaderTests.java index 165353b96..f911762c0 100644 --- a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/xml/StaxEventItemReaderTests.java +++ b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/xml/StaxEventItemReaderTests.java @@ -23,6 +23,7 @@ import javax.xml.transform.Source; import org.junit.Before; import org.junit.Test; import org.springframework.batch.item.ExecutionContext; +import org.springframework.batch.item.ItemCountAware; import org.springframework.batch.item.ItemStreamException; import org.springframework.batch.item.NonTransientResourceException; import org.springframework.core.io.AbstractResource; @@ -115,6 +116,35 @@ public class StaxEventItemReaderTests { source.close(); } + + @Test + public void testItemCountAwareFragment() throws Exception { + StaxEventItemReader source = createNewItemCountAwareInputSouce(); + source.afterPropertiesSet(); + source.open(executionContext); + assertEquals(1, source.read().getItemCount()); + assertEquals(2, source.read().getItemCount()); + assertNull(source.read()); // there are only two fragments + + source.close(); + } + + @Test + public void testItemCountAwareFragmentRestart() throws Exception { + StaxEventItemReader source = createNewItemCountAwareInputSouce(); + source.afterPropertiesSet(); + source.open(executionContext); + assertEquals(1, source.read().getItemCount()); + source.update(executionContext); + source.close(); + source = createNewItemCountAwareInputSouce(); + source.afterPropertiesSet(); + source.open(executionContext); + assertEquals(2, source.read().getItemCount()); + assertNull(source.read()); // there are only two fragments + + source.close(); + } @Test public void testFragmentNamespace() throws Exception { @@ -451,6 +481,19 @@ public class StaxEventItemReaderTests { return newSource; } + + private StaxEventItemReader createNewItemCountAwareInputSouce() { + Resource resource = new ByteArrayResource(xml.getBytes()); + + StaxEventItemReader newSource = new StaxEventItemReader(); + newSource.setResource(resource); + + newSource.setFragmentRootElementName(FRAGMENT_ROOT_ELEMENT); + newSource.setUnmarshaller(new ItemCountAwareMockFragmentUnmarshaller()); + newSource.setSaveState(true); + + return newSource; + } /** * A simple XMLEvent unmarshaller mock - check for the start and end document events for the fragment root & end @@ -524,6 +567,40 @@ public class StaxEventItemReaderTests { } } + + private static class ItemCountAwareMockFragmentUnmarshaller extends MockFragmentUnmarshaller { + @Override + public Object unmarshal(Source source) throws XmlMappingException, + IOException { + List fragment = (List) super.unmarshal(source); + if(fragment != null) { + return new ItemCountAwareFragment(fragment); + } else { + return null; + } + } + } + + private static class ItemCountAwareFragment implements ItemCountAware { + + private List fragment; + + private int itemCount; + + public ItemCountAwareFragment(List fragment) { + this.fragment = fragment; + } + + @Override + public void setItemCount(int count) { + this.itemCount = count; + } + + public int getItemCount() { + return itemCount; + } + + } private static class MockStaxEventItemReader extends StaxEventItemReader> {