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..0abfec542 --- /dev/null +++ b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/ItemCountAware.java @@ -0,0 +1,34 @@ +/* + * Copyright 20013 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. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +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 { + + /** + * Setter for the injection of the current item count. + * + * @param count the number of items that have been processed in this execution. + */ + 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..6e9e20950 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 @@ -1,5 +1,5 @@ /* - * Copyright 2006-2007 the original author or authors. + * Copyright 2006-2013 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. @@ -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; @@ -73,13 +74,17 @@ public abstract class AbstractItemCountingItemStreamItemReader extends Abstra } } - @Override + @Override public final T read() throws Exception, UnexpectedInputException, ParseException { if (currentItemCount >= maxItemCount) { return null; } currentItemCount++; - return doRead(); + T item = doRead(); + if(item instanceof ItemCountAware) { + ((ItemCountAware) item).setItemCount(currentItemCount); + } + return item; } protected int getCurrentItemCount() { @@ -115,9 +120,9 @@ public abstract class AbstractItemCountingItemStreamItemReader extends Abstra this.maxItemCount = count; } - @Override + @Override public void close() throws ItemStreamException { - super.close(); + super.close(); currentItemCount = 0; try { doClose(); @@ -127,9 +132,9 @@ public abstract class AbstractItemCountingItemStreamItemReader extends Abstra } } - @Override + @Override public void open(ExecutionContext executionContext) throws ItemStreamException { - super.open(executionContext); + super.open(executionContext); try { doOpen(); } @@ -161,9 +166,9 @@ public abstract class AbstractItemCountingItemStreamItemReader extends Abstra } - @Override + @Override public void update(ExecutionContext executionContext) throws ItemStreamException { - super.update(executionContext); + super.update(executionContext); if (saveState) { Assert.notNull(executionContext, "ExecutionContext must not be null"); executionContext.putInt(getExecutionContextKey(READ_COUNT), currentItemCount); @@ -182,7 +187,7 @@ public abstract class AbstractItemCountingItemStreamItemReader extends Abstra * @param name the name for the component */ public void setName(String name) { - this.setExecutionContextName(name); + this.setExecutionContextName(name); } /** 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..7c8bd5b94 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; @@ -32,6 +33,8 @@ public class FlatFileItemReaderTests { private FlatFileItemReader reader = new FlatFileItemReader(); + private FlatFileItemReader itemReader = new FlatFileItemReader(); + private ExecutionContext executionContext = new ExecutionContext(); @Before @@ -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 @@ -48,18 +54,18 @@ public class FlatFileItemReaderTests { // 1 record = 2 lines boolean pair = true; - @Override + @Override public boolean isEndOfRecord(String line) { pair = !pair; return pair; } - @Override + @Override public String postProcess(String record) { return record; } - @Override + @Override public String preProcess(String record) { return record; } @@ -86,18 +92,18 @@ public class FlatFileItemReaderTests { // 1 record = 2 lines boolean pair = true; - @Override + @Override public boolean isEndOfRecord(String line) { pair = !pair; return pair; } - @Override + @Override public String postProcess(String record) { return record; } - @Override + @Override public String preProcess(String record) { return record; } @@ -125,17 +131,17 @@ public class FlatFileItemReaderTests { reader.setRecordSeparatorPolicy(new RecordSeparatorPolicy() { - @Override + @Override public boolean isEndOfRecord(String line) { return StringUtils.hasText(line); } - @Override + @Override public String postProcess(String record) { return StringUtils.hasText(record) ? record : null; } - @Override + @Override public String preProcess(String record) { return record; } @@ -159,7 +165,7 @@ public class FlatFileItemReaderTests { // 1 record = 2 lines boolean pair = true; - @Override + @Override public boolean isEndOfRecord(String line) { if (StringUtils.hasText(line)) { pair = !pair; @@ -167,12 +173,12 @@ public class FlatFileItemReaderTests { return pair; } - @Override + @Override public String postProcess(String record) { return StringUtils.hasText(record) ? record : null; } - @Override + @Override public String preProcess(String record) { return record; } @@ -317,17 +323,17 @@ public class FlatFileItemReaderTests { public void testOpenBadIOInput() throws Exception { reader.setResource(new AbstractResource() { - @Override + @Override public String getDescription() { return null; } - @Override + @Override public InputStream getInputStream() throws IOException { throw new IOException(); } - @Override + @Override public boolean exists() { return true; } @@ -402,7 +408,7 @@ public class FlatFileItemReaderTests { @Test public void testMappingExceptionWrapping() throws Exception { LineMapper exceptionLineMapper = new LineMapper() { - @Override + @Override public String mapLine(String line, int lineNumber) throws Exception { if (lineNumber == 2) { throw new Exception("Couldn't map line 2"); @@ -428,6 +434,66 @@ public class FlatFileItemReaderTests { } } + @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()); } @@ -437,19 +503,58 @@ public class FlatFileItemReaderTests { public NonExistentResource() { } - @Override + @Override public boolean exists() { return false; } - @Override + @Override public String getDescription() { return "NonExistentResource"; } - @Override + @Override public InputStream getInputStream() throws IOException { return null; } } + + private static class Item implements ItemCountAware { + + private String value; + + private int itemCount; + + public Item(String value) { + this.value = value; + } + + @SuppressWarnings("unused") + 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..0f2dc0ea0 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; @@ -116,6 +117,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 { @@ -235,14 +265,14 @@ public class StaxEventItemReaderTests { assertNull(source.read()); source.update(executionContext); source.close(); - + assertEquals(3, executionContext.getInt(ClassUtils.getShortName(StaxEventItemReader.class) + ".read.count")); - + source = createNewInputSouce(); source.open(executionContext); assertNull(source.read()); } - + @Test public void testRestoreWorksFromClosedStream() throws Exception { source.close(); @@ -311,17 +341,17 @@ public class StaxEventItemReaderTests { public void testOpenBadIOInput() throws Exception { source.setResource(new AbstractResource() { - @Override + @Override public String getDescription() { return null; } - @Override + @Override public InputStream getInputStream() throws IOException { throw new IOException(); } - @Override + @Override public boolean exists() { return true; } @@ -426,12 +456,12 @@ public class StaxEventItemReaderTests { public static final String MESSAGE = "Unmarshallers on strike."; - @Override + @Override public Object unmarshal(Source source) throws XmlMappingException, IOException { throw new UnmarshallingFailureException(MESSAGE); } - @Override + @Override @SuppressWarnings("rawtypes") public boolean supports(Class clazz) { return true; @@ -452,6 +482,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 * tags + skips the fragment contents. @@ -476,7 +519,7 @@ public class StaxEventItemReaderTests { return events; } - @Override + @Override @SuppressWarnings("rawtypes") public boolean supports(Class clazz) { return true; @@ -488,7 +531,7 @@ public class StaxEventItemReaderTests { * @param source * @return list of the events from fragment body */ - @Override + @Override public Object unmarshal(Source source) throws XmlMappingException, IOException { List fragmentContent; @@ -525,11 +568,43 @@ public class StaxEventItemReaderTests { } + @SuppressWarnings("unchecked") + 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 int itemCount; + + public ItemCountAwareFragment(List fragment) { + } + + @Override + public void setItemCount(int count) { + this.itemCount = count; + } + + public int getItemCount() { + return itemCount; + } + + } + private static class MockStaxEventItemReader extends StaxEventItemReader> { private boolean openCalled = false; - @Override + @Override public void open(ExecutionContext executionContext) { super.open(executionContext); openCalled = true; @@ -549,17 +624,17 @@ public class StaxEventItemReaderTests { public NonExistentResource() { } - @Override + @Override public boolean exists() { return false; } - @Override + @Override public String getDescription() { return "NonExistantResource"; } - @Override + @Override public InputStream getInputStream() throws IOException { return null; }