RESOLVED - issue BATCH-620: Flat File Writer that writes a Header

Added headers to flat file and xml item writers

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;
-	}
-
 }

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
 	 * <li>xml declaration - defines encoding and XML version</li>
 	 * <li>opening tag of the root element and its attributes</li>
 	 * </ul>
-	 * 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(("</" + getRootTagName() + ">").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;
-	}
-
 }

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();
 	}

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 = "<!--" + ClassUtils.getShortName(StaxEventItemWriter.class)
+			+ "-testString-->";
 
 	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, "<!--" + header1 + "-->"));
+		assertTrue("Wrong content: "+content, contains(content, "<!--" + header2 + "-->"));
+		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, "<!--" + header + "-->"));
+		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, "<!--" + header + "-->"));
+		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;
+	}
 }