diff --git a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/file/MultiResourceItemWriter.java b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/file/MultiResourceItemWriter.java index 8c093996b..3fd351a79 100644 --- a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/file/MultiResourceItemWriter.java +++ b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/file/MultiResourceItemWriter.java @@ -46,6 +46,8 @@ public class MultiResourceItemWriter extends ExecutionContextUserSupport impl private ResourceSuffixCreator suffixCreator = new SimpleResourceSuffixCreator(); + private boolean saveState = true; + public MultiResourceItemWriter() { setName(ClassUtils.getShortName(MultiResourceItemWriter.class)); } @@ -55,7 +57,7 @@ public class MultiResourceItemWriter extends ExecutionContextUserSupport impl delegate.close(new ExecutionContext()); resourceIndex++; currentResourceItemCount = 0; - pointDelegateToNextResource(); + setResourceToDelegate(); delegate.open(new ExecutionContext()); } delegate.write(items); @@ -88,12 +90,17 @@ public class MultiResourceItemWriter extends ExecutionContextUserSupport impl /** * Prototype for output resources. Actual output files will be created in * the same directory and use the same name as this prototype with appended - * suffix (according to {@link #setResourceSuffixCreator(ResourceSuffixCreator)}. + * suffix (according to + * {@link #setResourceSuffixCreator(ResourceSuffixCreator)}. */ public void setResource(Resource resource) { this.resource = resource; } + public void setSaveState(boolean saveState) { + this.saveState = saveState; + } + public void close(ExecutionContext executionContext) throws ItemStreamException { resourceIndex = 1; currentResourceItemCount = 0; @@ -103,9 +110,9 @@ public class MultiResourceItemWriter extends ExecutionContextUserSupport impl public void open(ExecutionContext executionContext) throws ItemStreamException { resourceIndex = executionContext.getInt(getKey(RESOURCE_INDEX_KEY), 1); currentResourceItemCount = executionContext.getInt(getKey(CURRENT_RESOURCE_ITEM_COUNT), 0); - + try { - pointDelegateToNextResource(); + setResourceToDelegate(); } catch (IOException e) { throw new ItemStreamException("Couldn't open resource", e); @@ -114,15 +121,17 @@ public class MultiResourceItemWriter extends ExecutionContextUserSupport impl } public void update(ExecutionContext executionContext) throws ItemStreamException { - delegate.update(executionContext); - executionContext.putInt(getKey(CURRENT_RESOURCE_ITEM_COUNT), currentResourceItemCount); - executionContext.putInt(getKey(RESOURCE_INDEX_KEY), resourceIndex); + if (saveState) { + delegate.update(executionContext); + executionContext.putInt(getKey(CURRENT_RESOURCE_ITEM_COUNT), currentResourceItemCount); + executionContext.putInt(getKey(RESOURCE_INDEX_KEY), resourceIndex); + } } /** - * Create next output resource and point the delegate to it. + * Create output resource (if necessary) and point the delegate to it. */ - private void pointDelegateToNextResource() throws IOException { + private void setResourceToDelegate() throws IOException { String path = resource.getFile().getAbsolutePath() + suffixCreator.getSuffix(resourceIndex); File file = new File(path); file.createNewFile(); diff --git a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/xml/StaxEventItemWriter.java b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/xml/StaxEventItemWriter.java index 5b426bfea..916e1ef9f 100644 --- a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/xml/StaxEventItemWriter.java +++ b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/xml/StaxEventItemWriter.java @@ -87,9 +87,6 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implemen // root element attributes private Map rootElementAttributes = null; - // signalizes that marshalling was restarted - private boolean restarted = false; - // TRUE means, that output file will be overwritten if exists - default is // TRUE private boolean overwriteOutput = true; @@ -255,6 +252,7 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implemen Assert.notNull(resource, "The resource must be set"); long startAtPosition = 0; + boolean restarted = false; // if restart data is provided, restart from provided offset // otherwise start from beginning @@ -263,7 +261,7 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implemen restarted = true; } - open(startAtPosition); + open(startAtPosition, restarted); if (startAtPosition == 0) { try { @@ -281,7 +279,7 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implemen /** * Helper method for opening output source at given file position */ - private void open(long position) { + private void open(long position, boolean restarted) { File file; FileOutputStream os = null; @@ -433,12 +431,7 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implemen currentRecordCount += items.size(); - doWrite(items); - - } - - private void doWrite(List objects) throws XmlMappingException, IOException { - for (Object object : objects) { + for (Object object : items) { Assert.state(marshaller.supports(object.getClass()), "Marshaller must support the class of the marshalled object"); marshaller.marshal(object, new StaxResult(eventWriter)); @@ -449,6 +442,7 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implemen catch (XMLStreamException e) { throw new FlushFailedException("Failed to flush the events", e); } + } /** diff --git a/spring-batch-infrastructure/src/test/java/MultiResourceItemWriterXmlTests.java b/spring-batch-infrastructure/src/test/java/MultiResourceItemWriterXmlTests.java new file mode 100644 index 000000000..206a893f7 --- /dev/null +++ b/spring-batch-infrastructure/src/test/java/MultiResourceItemWriterXmlTests.java @@ -0,0 +1,103 @@ +import java.io.File; +import java.io.IOException; +import java.util.Arrays; + +import javax.xml.stream.XMLEventFactory; +import javax.xml.stream.XMLEventWriter; +import javax.xml.stream.XMLStreamException; +import javax.xml.transform.Result; + +import static org.junit.Assert.*; +import org.junit.Before; +import org.junit.Test; +import org.springframework.batch.item.file.AbstractMultiResourceItemWriterTests; +import org.springframework.batch.item.file.MultiResourceItemWriter; +import org.springframework.batch.item.xml.StaxEventItemWriter; +import org.springframework.oxm.Marshaller; +import org.springframework.oxm.XmlMappingException; +import org.springframework.util.Assert; +import org.springframework.xml.transform.StaxResult; + +/** + * Tests for {@link MultiResourceItemWriter} delegating to + * {@link StaxEventItemWriter}. + */ +public class MultiResourceItemWriterXmlTests extends AbstractMultiResourceItemWriterTests { + + final private String xmlDocStart = ""; + + final private String xmlDocEnd = ""; + + @Override + @Before + public void setUp() throws Exception { + delegate = new StaxEventItemWriter() { + { + setMarshaller(new SimpleMarshaller()); + } + }; + super.setUp(); + } + + /** + * Writes object's toString representation as tag. + */ + private static class SimpleMarshaller implements Marshaller { + public void marshal(Object graph, Result result) throws XmlMappingException, IOException { + Assert.isInstanceOf(StaxResult.class, result); + + StaxResult staxResult = (StaxResult) result; + try { + XMLEventFactory factory = XMLEventFactory.newInstance(); + XMLEventWriter writer = staxResult.getXMLEventWriter(); + writer.add(factory.createStartDocument("UTF-8")); + writer.add(factory.createStartElement("prefix", "namespace", graph.toString())); + writer.add(factory.createEndElement("prefix", "namespace", graph.toString())); + writer.add(factory.createEndDocument()); + } + catch (XMLStreamException e) { + throw new RuntimeException("Exception while writing to output file", e); + } + } + + @SuppressWarnings("unchecked") + public boolean supports(Class clazz) { + return true; + } + } + + @Test + public void multiResourceWritingWithRestart() throws Exception { + + tested.write(Arrays.asList("1", "2", "3")); + + File part1 = new File(file.getAbsolutePath() + suffixCreator.getSuffix(1)); + assertTrue(part1.exists()); + + tested.write(Arrays.asList("4")); + File part2 = new File(file.getAbsolutePath() + suffixCreator.getSuffix(2)); + assertTrue(part2.exists()); + + tested.update(executionContext); + tested.close(executionContext); + + assertEquals(xmlDocStart + "" + xmlDocEnd, readFile(part2)); + assertEquals(xmlDocStart + "" + xmlDocEnd, + readFile(part1)); + + tested.open(executionContext); + + tested.write(Arrays.asList("5")); + + tested.write(Arrays.asList("6", "7", "8", "9")); + File part3 = new File(file.getAbsolutePath() + suffixCreator.getSuffix(3)); + assertTrue(part3.exists()); + + tested.close(executionContext); + + assertEquals(xmlDocStart + "" + xmlDocEnd, readFile(part2)); + assertEquals(xmlDocStart + + "" + xmlDocEnd, + readFile(part3)); + } +} diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/file/AbstractMultiResourceItemWriterTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/file/AbstractMultiResourceItemWriterTests.java new file mode 100644 index 000000000..e7c5fb7b1 --- /dev/null +++ b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/file/AbstractMultiResourceItemWriterTests.java @@ -0,0 +1,52 @@ +package org.springframework.batch.item.file; + +import java.io.BufferedReader; +import java.io.File; +import java.io.FileReader; + +import org.junit.Before; +import org.springframework.batch.item.ExecutionContext; +import org.springframework.core.io.FileSystemResource; + +/** + * Tests for {@link MultiResourceItemWriter}. + * + * @see MultiResourceItemWriterTests + * @see MultiResourceItemReaderXmlTests + */ +public class AbstractMultiResourceItemWriterTests { + + protected MultiResourceItemWriter tested = new MultiResourceItemWriter(); + + protected File file; + + protected ResourceSuffixCreator suffixCreator = new SimpleResourceSuffixCreator(); + + protected ResourceAwareItemWriterItemStream delegate; + + protected ExecutionContext executionContext = new ExecutionContext(); + + @Before + public void setUp() throws Exception { + file = File.createTempFile(MultiResourceItemWriterTests.class.getSimpleName(), null); + tested.setResource(new FileSystemResource(file)); + tested.setDelegate(delegate); + tested.setResourceSuffixCreator(suffixCreator); + tested.setItemCountLimitPerResource(2); + tested.setSaveState(true); + tested.open(executionContext); + } + + protected String readFile(File f) throws Exception { + BufferedReader reader = new BufferedReader(new FileReader(f)); + StringBuilder result = new StringBuilder(); + while (true) { + String line = reader.readLine(); + if (line == null) { + break; + } + result.append(line); + } + return result.toString(); + } +} diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/file/MultiResourceItemWriterTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/file/MultiResourceItemWriterTests.java index 5ad5f89b2..3ac147ad5 100644 --- a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/file/MultiResourceItemWriterTests.java +++ b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/file/MultiResourceItemWriterTests.java @@ -1,46 +1,30 @@ package org.springframework.batch.item.file; -import static org.junit.Assert.*; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; -import java.io.BufferedReader; import java.io.File; -import java.io.FileReader; import java.util.Arrays; import org.junit.Before; import org.junit.Test; -import org.springframework.batch.item.ExecutionContext; import org.springframework.batch.item.file.transform.PassThroughLineAggregator; -import org.springframework.core.io.FileSystemResource; /** - * Tests for {@link MultiResourceItemWriter}. + * Tests for {@link MultiResourceItemWriter} delegating to + * {@link FlatFileItemWriter}. */ -public class MultiResourceItemWriterTests { - - private MultiResourceItemWriter tested = new MultiResourceItemWriter(); - - private File file; - - private ResourceSuffixCreator suffixCreator = new SimpleResourceSuffixCreator(); - - private ResourceAwareItemWriterItemStream delegate = new FlatFileItemWriter() { - { - setLineAggregator(new PassThroughLineAggregator()); - } - }; - - private ExecutionContext executionContext = new ExecutionContext(); +public class MultiResourceItemWriterTests extends AbstractMultiResourceItemWriterTests { + @Override @Before public void setUp() throws Exception { - file = File.createTempFile(MultiResourceItemWriterTests.class.getSimpleName(), null); - tested.setResource(new FileSystemResource(file)); - tested.setDelegate(delegate); - tested.setResourceSuffixCreator(suffixCreator); - tested.setItemCountLimitPerResource(2); - - tested.open(executionContext); + delegate = new FlatFileItemWriter() { + { + setLineAggregator(new PassThroughLineAggregator()); + } + }; + super.setUp(); } @Test @@ -82,7 +66,7 @@ public class MultiResourceItemWriterTests { tested.update(executionContext); tested.close(executionContext); tested.open(executionContext); - + tested.write(Arrays.asList("5")); assertEquals("45", readFile(part2)); @@ -91,17 +75,4 @@ public class MultiResourceItemWriterTests { assertTrue(part3.exists()); assertEquals("6789", readFile(part3)); } - - private String readFile(File f) throws Exception { - BufferedReader reader = new BufferedReader(new FileReader(f)); - StringBuilder result = new StringBuilder(); - while (true) { - String line = reader.readLine(); - if (line == null) { - break; - } - result.append(line); - } - return result.toString(); - } }