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 00b6e2bd7..dbe577e5f 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 @@ -7,7 +7,6 @@ import java.io.OutputStreamWriter; import java.io.UnsupportedEncodingException; import java.nio.ByteBuffer; import java.nio.channels.FileChannel; -import java.util.ArrayList; import java.util.List; import java.util.Map; @@ -110,9 +109,9 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implemen private boolean saveState = true; - private List headerItems = new ArrayList(); + private StaxWriterCallback headerCallback; - private List footerItems = new ArrayList(); + private StaxWriterCallback footerCallback; public StaxEventItemWriter() { setName(ClassUtils.getShortName(StaxEventItemWriter.class)); @@ -136,6 +135,21 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implemen this.marshaller = marshaller; } + /** + * headerCallback is called before writing any items. + */ + public void setHeaderCallback(StaxWriterCallback headerCallback) { + this.headerCallback = headerCallback; + } + + /** + * footerCallback is called after writing all items but before closing the + * file + */ + public void setFooterCallback(StaxWriterCallback footerCallback) { + this.footerCallback = footerCallback; + } + /** * Get used encoding. * @@ -219,28 +233,6 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implemen this.overwriteOutput = overwriteOutput; } - /** - * Setter for the headerItems. This list will be marshalled and output - * before any calls to {@link #write(List)}. Header item type is not - * restricted, but note the {@link #setMarshaller(Marshaller)} needs to - * support the type. - * @param headerItems - */ - public void setHeaderItems(List headerItems) { - this.headerItems = headerItems; - } - - /** - * Setter for the footerItems. This list will be marshalled and output - * immediately before the writer is closed. Footer item type is not - * restricted, but note the {@link #setMarshaller(Marshaller)} needs to - * support the type. - * @param footerItems - */ - public void setFooterItems(List footerItems) { - this.footerItems = footerItems; - } - public void setSaveState(boolean saveState) { this.saveState = saveState; } @@ -275,7 +267,9 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implemen if (startAtPosition == 0) { try { - doWrite(headerItems); + if (headerCallback != null) { + headerCallback.write(delegateEventWriter); + } } catch (IOException e) { throw new ItemStreamException("Failed to write headerItems", e); @@ -396,7 +390,9 @@ public class StaxEventItemWriter extends ExecutionContextUserSupport implemen } try { - doWrite(footerItems); + if (footerCallback != null) { + footerCallback.write(delegateEventWriter); + } delegateEventWriter.flush(); endDocument(delegateEventWriter); } diff --git a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/xml/StaxWriterCallback.java b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/xml/StaxWriterCallback.java new file mode 100644 index 000000000..6b6065b63 --- /dev/null +++ b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/xml/StaxWriterCallback.java @@ -0,0 +1,20 @@ +package org.springframework.batch.item.xml; + +import java.io.IOException; + +import javax.xml.stream.XMLEventWriter; + +/** + * Callback interface for writing to an XML file - useful e.g. for handling headers + * and footers. + * + * @author Robert Kasanicky + */ +public interface StaxWriterCallback { + + /** + * Write contents using the supplied {@link XMLEventWriter}. It is not + * required to flush the writer inside this method. + */ + void write(XMLEventWriter writer) throws IOException; +} 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 5f1f495c4..14ac9107d 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 @@ -10,11 +10,11 @@ import static org.junit.Assert.fail; import java.io.File; import java.io.IOException; -import java.util.ArrayList; import java.util.Collections; import java.util.List; import javax.xml.stream.XMLEventFactory; +import javax.xml.stream.XMLEventWriter; import javax.xml.stream.XMLStreamException; import javax.xml.transform.Result; @@ -95,7 +95,6 @@ public class StaxEventItemWriterTests { // check the output is concatenation of 'before restart' and 'after // restart' writes. String outputFile = outputFileContent(); - System.out.println(outputFile); assertEquals(2, StringUtils.countOccurrencesOf(outputFile, TEST_STRING)); assertTrue(outputFile.contains("" + TEST_STRING + TEST_STRING + "")); } @@ -105,19 +104,26 @@ public class StaxEventItemWriterTests { */ @Test public void testWriteWithHeader() throws Exception { - final Object header1 = new Object(); - final Object header2 = new Object(); - writer.setHeaderItems(new ArrayList() { - { - add(header1); - add(header2); + + writer.setHeaderCallback(new StaxWriterCallback(){ + + public void write(XMLEventWriter writer) throws IOException { + XMLEventFactory factory = XMLEventFactory.newInstance(); + try { + writer.add(factory.createStartElement("", "", "header")); + writer.add(factory.createEndElement("", "", "header")); + } + catch (XMLStreamException e) { + throw new RuntimeException(e); + } + } + }); writer.open(executionContext); writer.write(items); String content = outputFileContent(); - assertTrue("Wrong content: " + content, content.contains((""))); - assertTrue("Wrong content: " + content, content.contains((""))); + assertTrue("Wrong content: " + content, content.contains(("
"))); assertTrue("Wrong content: " + content, content.contains(TEST_STRING)); } @@ -144,12 +150,45 @@ public class StaxEventItemWriterTests { */ @Test public void testOpenAndClose() throws Exception { + writer.setHeaderCallback(new StaxWriterCallback(){ + + public void write(XMLEventWriter writer) throws IOException { + XMLEventFactory factory = XMLEventFactory.newInstance(); + try { + writer.add(factory.createStartElement("", "", "header")); + writer.add(factory.createEndElement("", "", "header")); + } + catch (XMLStreamException e) { + throw new RuntimeException(e); + } + + } + + }); + writer.setFooterCallback(new StaxWriterCallback() { + + public void write(XMLEventWriter writer) throws IOException { + XMLEventFactory factory = XMLEventFactory.newInstance(); + try { + writer.add(factory.createStartElement("", "", "footer")); + writer.add(factory.createEndElement("", "", "footer")); + } + catch (XMLStreamException e) { + throw new RuntimeException(e); + } + + } + + }); writer.setRootTagName("testroot"); writer.setRootElementAttributes(Collections. singletonMap("attribute", "value")); writer.open(executionContext); writer.close(null); String content = outputFileContent(); + assertTrue(content.contains("")); + assertTrue(content.contains("
")); + assertTrue(content.contains("
")); assertTrue(content.endsWith("
")); } diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/xml/TransactionalStaxEventItemWriterTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/xml/TransactionalStaxEventItemWriterTests.java index 272e94ed5..6a8ec11a7 100644 --- a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/xml/TransactionalStaxEventItemWriterTests.java +++ b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/xml/TransactionalStaxEventItemWriterTests.java @@ -10,6 +10,7 @@ import java.util.Collections; import java.util.List; import javax.xml.stream.XMLEventFactory; +import javax.xml.stream.XMLEventWriter; import javax.xml.stream.XMLStreamException; import javax.xml.transform.Result; @@ -92,8 +93,21 @@ public class TransactionalStaxEventItemWriterTests { */ @Test public void testWriteWithHeaderAfterRollback() throws Exception { - Object header = new Object(); - writer.setHeaderItems(Collections.singletonList(header)); + writer.setHeaderCallback(new StaxWriterCallback(){ + + public void write(XMLEventWriter writer) throws IOException { + XMLEventFactory factory = XMLEventFactory.newInstance(); + try { + writer.add(factory.createStartElement("", "", "header")); + writer.add(factory.createEndElement("", "", "header")); + } + catch (XMLStreamException e) { + throw new RuntimeException(e); + } + + } + + }); writer.open(executionContext); try { new TransactionTemplate(transactionManager).execute(new TransactionCallback() { @@ -127,7 +141,7 @@ public class TransactionalStaxEventItemWriterTests { }); writer.close(executionContext); String content = outputFileContent(); - assertEquals("Wrong content: " + content, 1, StringUtils.countOccurrencesOf(content, (""))); + assertEquals("Wrong content: " + content, 1, StringUtils.countOccurrencesOf(content, ("
"))); assertEquals("Wrong content: " + content, 1, StringUtils.countOccurrencesOf(content, TEST_STRING)); } @@ -136,8 +150,21 @@ public class TransactionalStaxEventItemWriterTests { */ @Test public void testWriteWithHeaderAfterFlushAndRollback() throws Exception { - Object header = new Object(); - writer.setHeaderItems(Collections.singletonList(header)); + writer.setHeaderCallback(new StaxWriterCallback(){ + + public void write(XMLEventWriter writer) throws IOException { + XMLEventFactory factory = XMLEventFactory.newInstance(); + try { + writer.add(factory.createStartElement("", "", "header")); + writer.add(factory.createEndElement("", "", "header")); + } + catch (XMLStreamException e) { + throw new RuntimeException(e); + } + + } + + }); writer.open(executionContext); new TransactionTemplate(transactionManager).execute(new TransactionCallback() { public Object doInTransaction(TransactionStatus status) { @@ -172,7 +199,7 @@ public class TransactionalStaxEventItemWriterTests { } writer.close(executionContext); String content = outputFileContent(); - assertEquals("Wrong content: " + content, 1, StringUtils.countOccurrencesOf(content, (""))); + assertEquals("Wrong content: " + content, 1, StringUtils.countOccurrencesOf(content, ("
"))); assertEquals("Wrong content: " + content, 1, StringUtils.countOccurrencesOf(content, TEST_STRING)); }