diff --git a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/file/FlatFileItemWriter.java b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/file/FlatFileItemWriter.java index 5b4029ec9..836c03cda 100644 --- a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/file/FlatFileItemWriter.java +++ b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/file/FlatFileItemWriter.java @@ -25,6 +25,7 @@ import java.nio.channels.Channels; import java.nio.channels.FileChannel; import java.nio.charset.UnsupportedCharsetException; import java.util.ArrayList; +import java.util.Arrays; import java.util.Iterator; import java.util.List; @@ -95,6 +96,8 @@ public class FlatFileItemWriter extends ExecutionContextUserSupport implements I private List lineBuffer = new ArrayList(); + private List headerLines = new ArrayList(); + public FlatFileItemWriter() { setName(ClassUtils.getShortName(FlatFileItemWriter.class)); } @@ -139,6 +142,50 @@ public class FlatFileItemWriter extends ExecutionContextUserSupport implements I this.resource = resource; } + /** + * Sets encoding for output template. + */ + public void setEncoding(String newEncoding) { + this.encoding = newEncoding; + } + + /** + * Sets buffer size for output template + */ + public void setBufferSize(int newSize) { + this.bufferSize = newSize; + } + + /** + * @param shouldDeleteIfExists the shouldDeleteIfExists to set + */ + public void setShouldDeleteIfExists(boolean shouldDeleteIfExists) { + this.shouldDeleteIfExists = shouldDeleteIfExists; + } + + /** + * Set the flag indicating whether or not state should be saved in the + * provided {@link ExecutionContext} during the {@link ItemStream} call to + * update. Setting this to false means that it will always start at the + * beginning on a restart. + * + * @param saveState + */ + public void setSaveState(boolean saveState) { + this.saveState = saveState; + } + + /** + * Public setter for the header lines. These will be output at the head of + * the file before any calls to {@link #write(Object)} (and not on restart + * unless the restart is after a failure before the first flush). + * + * @param headerLines the header lines to set + */ + public void setHeaderLines(String[] headerLines) { + this.headerLines = Arrays.asList(headerLines); + } + /** * Writes out a string followed by a "new line", where the format of the new * line separator is determined by the underlying operating system. If the @@ -167,27 +214,6 @@ public class FlatFileItemWriter extends ExecutionContextUserSupport implements I } } - /** - * Sets encoding for output template. - */ - public void setEncoding(String newEncoding) { - this.encoding = newEncoding; - } - - /** - * Sets buffer size for output template - */ - public void setBufferSize(int newSize) { - this.bufferSize = newSize; - } - - /** - * @param shouldDeleteIfExists the shouldDeleteIfExists to set - */ - public void setShouldDeleteIfExists(boolean shouldDeleteIfExists) { - this.shouldDeleteIfExists = shouldDeleteIfExists; - } - /** * Initialize the Output Template. * @@ -199,6 +225,12 @@ public class FlatFileItemWriter extends ExecutionContextUserSupport implements I outputState.restoreFrom(executionContext); } outputState.initializeBufferedWriter(); + if (outputState.lastMarkedByteOffsetPosition == 0) { + for (Iterator iterator = headerLines.iterator(); iterator.hasNext();) { + String line = (String) iterator.next(); + lineBuffer.add(line + LINE_SEPARATOR); + } + } } /** @@ -225,6 +257,21 @@ public class FlatFileItemWriter extends ExecutionContextUserSupport implements I } } + public void flush() throws FlushFailedException { + OutputState state = getOutputState(); + for (Iterator iterator = lineBuffer.listIterator(); iterator.hasNext();) { + String line = (String) iterator.next(); + try { + state.write(line); + } + catch (IOException e) { + throw new FlushFailedException("Failed to write line to output file: " + line, e); + } + } + lineBuffer.clear(); + state.mark(); + } + // Returns object representing state. private OutputState getOutputState() { if (state == null) { @@ -505,31 +552,4 @@ public class FlatFileItemWriter extends ExecutionContextUserSupport implements I lineBuffer.clear(); } - public void flush() throws FlushFailedException { - OutputState state = getOutputState(); - for (Iterator iterator = lineBuffer.listIterator(); iterator.hasNext();) { - String line = (String) iterator.next(); - try { - state.write(line); - } - catch (IOException e) { - throw new FlushFailedException("Failed to write line to output file: " + line, e); - } - } - lineBuffer.clear(); - state.mark(); - } - - /** - * Set the boolean indicating whether or not state should be saved in the - * provided {@link ExecutionContext} during the {@link ItemStream} call to - * update. Setting this to false means that it will always start at the - * beginning. - * - * @param saveState - */ - public void setSaveState(boolean saveState) { - this.saveState = saveState; - } - } 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 f9aab331f..58f0580c5 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 @@ -6,6 +6,7 @@ import java.io.IOException; import java.nio.ByteBuffer; import java.nio.channels.FileChannel; import java.util.ArrayList; +import java.util.Arrays; import java.util.Iterator; import java.util.List; import java.util.Map; @@ -31,19 +32,21 @@ import org.springframework.util.ClassUtils; import org.springframework.util.CollectionUtils; /** - * An implementation of {@link ItemWriter} which uses StAX and {@link EventWriterSerializer} for serializing object to - * XML. + * An implementation of {@link ItemWriter} which uses StAX and + * {@link EventWriterSerializer} for serializing object to XML. * - * This item writer also provides restart, statistics and transaction features by implementing corresponding interfaces. + * This item writer also provides restart, statistics and transaction features + * by implementing corresponding interfaces. * - * Output is buffered until {@link #flush()} is called - only then the actual writing to file takes place. + * Output is buffered until {@link #flush()} is called - only then the actual + * writing to file takes place. * * @author Peter Zozom * @author Robert Kasanicky * */ public class StaxEventItemWriter extends ExecutionContextUserSupport implements ItemWriter, ItemStream, - InitializingBean { + InitializingBean { // default encoding private static final String DEFAULT_ENCODING = "UTF-8"; @@ -110,6 +113,8 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implements // #flush() private List buffer = new ArrayList(); + private List headers = new ArrayList(); + public StaxEventItemWriter() { setName(ClassUtils.getShortName(StaxEventItemWriter.class)); } @@ -178,7 +183,8 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implements } /** - * Set the tag name of the root element. If not set, default name is used ("root"). + * Set the tag name of the root element. If not set, default name is used + * ("root"). * * @param rootTagName the tag name to be used for the root element */ @@ -205,7 +211,8 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implements } /** - * Set "overwrite" flag for the output file. Flag is ignored when output file processing is restarted. + * Set "overwrite" flag for the output file. Flag is ignored when output + * file processing is restarted. * * @param overwriteOutput */ @@ -213,6 +220,19 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implements this.overwriteOutput = overwriteOutput; } + /** + * Setter for the headers. This list will be marshalled and output before + * any calls to {@link #write(Object)}. + * @param headers + */ + public void setHeaderItems(Object[] headers) { + this.headers = Arrays.asList(headers); + } + + public void setSaveState(boolean saveState) { + this.saveState = saveState; + } + /** * @throws Exception * @see org.springframework.beans.factory.InitializingBean#afterPropertiesSet() @@ -238,6 +258,13 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implements } open(startAtPosition); + + if (startAtPosition==0) { + for (Iterator iterator = headers.iterator(); iterator.hasNext();) { + Object header = (Object) iterator.next(); + write(header); + } + } } /** @@ -254,7 +281,8 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implements os = new FileOutputStream(file, true); channel = os.getChannel(); setPosition(position); - } catch (IOException ioe) { + } + catch (IOException ioe) { throw new DataAccessResourceFailureException("Unable to write to file resource: [" + resource + "]", ioe); } @@ -266,7 +294,8 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implements if (!restarted) { startDocument(delegateEventWriter); } - } catch (XMLStreamException xse) { + } + catch (XMLStreamException xse) { throw new DataAccessResourceFailureException("Unable to write to file resource: [" + resource + "]", xse); } @@ -278,8 +307,8 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implements *
  • xml declaration - defines encoding and XML version
    • *
  • opening tag of the root element and its attributes
    • * - * If this is not sufficient for you, simply override this method. Encoding, version and root tag name can be - * retrieved with corresponding getters. + * If this is not sufficient for you, simply override this method. Encoding, + * version and root tag name can be retrieved with corresponding getters. * * @param writer XML event writer * @throws XMLStreamException @@ -307,7 +336,8 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implements } /** - * Finishes the XML document. It closes any start tag and writes corresponding end tags. + * Finishes the XML document. It closes any start tag and writes + * corresponding end tags. * * @param writer XML event writer * @throws XMLStreamException @@ -316,17 +346,18 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implements // writer.writeEndDocument(); <- this doesn't work after restart // we need to write end tag of the root element manually - - //harmless event to close the root tag if there were no items + + // harmless event to close the root tag if there were no items XMLEventFactory factory = XMLEventFactory.newInstance(); writer.add(factory.createCharacters("")); - + writer.flush(); - + ByteBuffer bbuf = ByteBuffer.wrap(("").getBytes()); try { channel.write(bbuf); - } catch (IOException ioe) { + } + catch (IOException ioe) { throw new DataAccessResourceFailureException("Unable to close file resource: [" + resource + "]", ioe); } } @@ -342,9 +373,11 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implements endDocument(delegateEventWriter); eventWriter.close(); channel.close(); - } catch (XMLStreamException xse) { + } + catch (XMLStreamException xse) { throw new DataAccessResourceFailureException("Unable to close file resource: [" + resource + "]", xse); - } catch (IOException ioe) { + } + catch (IOException ioe) { throw new DataAccessResourceFailureException("Unable to close file resource: [" + resource + "]", ioe); } } @@ -376,7 +409,8 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implements } /* - * Get the actual position in file channel. This method flushes any buffered data before position is read. + * Get the actual position in file channel. This method flushes any buffered + * data before position is read. * * @return byte offset in file channel */ @@ -387,7 +421,8 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implements try { eventWriter.flush(); position = channel.position(); - } catch (Exception e) { + } + catch (Exception e) { throw new DataAccessResourceFailureException("Unable to write to file resource: [" + resource + "]", e); } @@ -403,10 +438,11 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implements try { Assert.state(channel.size() >= lastCommitPointPosition, - "Current file size is smaller than size at last commit"); + "Current file size is smaller than size at last commit"); channel.truncate(newPosition); channel.position(newPosition); - } catch (IOException e) { + } + catch (IOException e) { throw new DataAccessResourceFailureException("Unable to write to file resource: [" + resource + "]", e); } @@ -435,8 +471,4 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implements buffer.clear(); } - public void setSaveState(boolean saveState) { - this.saveState = saveState; - } - } diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/file/FlatFileItemWriterTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/file/FlatFileItemWriterTests.java index 478c1279d..695258aa4 100644 --- a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/file/FlatFileItemWriterTests.java +++ b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/file/FlatFileItemWriterTests.java @@ -336,6 +336,70 @@ public class FlatFileItemWriterTests extends TestCase { assertEquals(TEST_STRING, lineFromFile); } + public void testWriteHeader() throws Exception { + writer.setHeaderLines(new String[] {"a", "b"}); + writer.open(executionContext); + writer.write(TEST_STRING); + writer.flush(); + writer.close(null); + String lineFromFile = readLine(); + assertEquals("a", lineFromFile); + lineFromFile = readLine(); + assertEquals("b", lineFromFile); + lineFromFile = readLine(); + assertEquals(TEST_STRING, lineFromFile); + } + + public void testWriteHeaderAfterRestartOnFirstChunk() throws Exception { + writer.setHeaderLines(new String[] {"a", "b"}); + writer.open(executionContext); + writer.write(TEST_STRING); + writer.clear(); + writer.close(executionContext); + writer.open(executionContext); + writer.write(TEST_STRING); + writer.flush(); + writer.close(executionContext); + String lineFromFile = readLine(); + assertEquals("a", lineFromFile); + lineFromFile = readLine(); + assertEquals("b", lineFromFile); + lineFromFile = readLine(); + assertEquals(TEST_STRING, lineFromFile); + lineFromFile = readLine(); + assertEquals(null, lineFromFile); + } + + public void testWriteHeaderAfterRestartOnSecondChunk() throws Exception { + writer.setHeaderLines(new String[] {"a", "b"}); + writer.open(executionContext); + writer.write(TEST_STRING); + writer.flush(); + writer.update(executionContext); + writer.write(TEST_STRING); + writer.clear(); + writer.close(executionContext); + String lineFromFile = readLine(); + assertEquals("a", lineFromFile); + lineFromFile = readLine(); + assertEquals("b", lineFromFile); + lineFromFile = readLine(); + assertEquals(TEST_STRING, lineFromFile); + writer.open(executionContext); + writer.write(TEST_STRING); + writer.flush(); + writer.close(executionContext); + reader = null; + lineFromFile = readLine(); + assertEquals("a", lineFromFile); + lineFromFile = readLine(); + assertEquals("b", lineFromFile); + lineFromFile = readLine(); + assertEquals(TEST_STRING, lineFromFile); + lineFromFile = readLine(); + assertEquals(TEST_STRING, lineFromFile); + } + private void commit() throws Exception { writer.flush(); } diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/xml/StaxEventItemWriterTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/xml/StaxEventItemWriterTests.java index f4d9c104c..843dc8019 100644 --- a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/xml/StaxEventItemWriterTests.java +++ b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/xml/StaxEventItemWriterTests.java @@ -37,11 +37,12 @@ public class StaxEventItemWriterTests extends TestCase { // test item for writing to output private Object item = new Object() { public String toString() { - return TEST_STRING; + return ClassUtils.getShortName(StaxEventItemWriter.class)+"-testString"; } }; - private static final String TEST_STRING = ClassUtils.getShortName(StaxEventItemWriter.class) + "-testString"; + private static final String TEST_STRING = ""; private static final int NOT_FOUND = -1; @@ -96,9 +97,11 @@ public class StaxEventItemWriterTests extends TestCase { public void testWriteAndFlush() throws Exception { writer.open(executionContext); writer.write(item); - assertEquals("", outputFileContent()); + String content = outputFileContent(); + assertEquals("", content); writer.flush(); - assertTrue(contains(outputFileContent(), TEST_STRING)); + content = outputFileContent(); + assertTrue("Wrong content: "+content, contains(content, TEST_STRING)); } /** @@ -131,6 +134,66 @@ public class StaxEventItemWriterTests extends TestCase { assertEquals(NOT_FOUND, thirdRecord); } + /** + * Item is written to the output file only after flush. + */ + public void testWriteWithHeader() throws Exception { + Object header1 = new Object(); + Object header2 = new Object(); + writer.setHeaderItems(new Object[] {header1, header2}); + writer.open(executionContext); + writer.write(item); + String content = outputFileContent(); + assertEquals("", content); + writer.flush(); + content = outputFileContent(); + assertTrue("Wrong content: "+content, contains(content, "")); + assertTrue("Wrong content: "+content, contains(content, "")); + assertTrue("Wrong content: "+content, contains(content, TEST_STRING)); + } + + /** + * Item is written to the output file only after flush. + */ + public void testWriteWithHeaderAfterRollback() throws Exception { + Object header = new Object(); + writer.setHeaderItems(new Object[] {header}); + writer.open(executionContext); + writer.write(item); + String content = outputFileContent(); + assertEquals("", content); + writer.clear(); + writer.flush(); + writer.open(executionContext); + writer.write(item); + writer.flush(); + content = outputFileContent(); + assertEquals("Wrong content: "+content, 1, countContains(content, "")); + assertEquals("Wrong content: "+content, 1, countContains(content, TEST_STRING)); + } + + /** + * Item is written to the output file only after flush. + */ + public void testWriteWithHeaderAfterFlushAndRollback() throws Exception { + Object header = new Object(); + writer.setHeaderItems(new Object[] {header}); + writer.open(executionContext); + writer.write(item); + String content = outputFileContent(); + assertEquals("", content); + writer.flush(); + writer.update(executionContext); + writer.close(executionContext); + writer.open(executionContext); + writer.write(item); + writer.clear(); + writer.flush(); + content = outputFileContent(); + assertEquals("Wrong content: "+content, 1, countContains(content, "")); + assertEquals("Wrong content: "+content, 1, countContains(content, TEST_STRING)); + } + /** * Count of 'records written so far' is returned as statistics. */ @@ -141,7 +204,7 @@ public class StaxEventItemWriterTests extends TestCase { writer.write(item); writer.update(executionContext); long writeStatistics = executionContext.getLong(ClassUtils.getShortName(StaxEventItemWriter.class) - + ".record.count"); + + ".record.count"); assertEquals(i, writeStatistics); } @@ -195,7 +258,8 @@ public class StaxEventItemWriterTests extends TestCase { StaxResult staxResult = (StaxResult) result; try { staxResult.getXMLEventWriter().add(XMLEventFactory.newInstance().createComment(graph.toString())); - } catch (XMLStreamException e) { + } + catch (XMLStreamException e) { throw new RuntimeException("Exception while writing to output file", e); } } @@ -237,4 +301,14 @@ public class StaxEventItemWriterTests extends TestCase { private boolean contains(String str, String searchStr) { return str.indexOf(searchStr) != -1; } + + private int countContains(String str, String searchStr) { + int begin = -1; + int count = 0; + while (str.indexOf(searchStr, begin+1) > begin) { + count++; + begin = str.indexOf(searchStr, begin); + } + return count; + } }