OPEN - BATCH-34: Support for multiple I/O files in a single jobRun for a particular scheduleDate.

tested multi-resource XML output

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/file/MultiResourceItemWriter.java

@@ -46,6 +46,8 @@ public class MultiResourceItemWriter<T> 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<T> 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<T> 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<T> 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<T> 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();

spring-batch-infrastructure/src/main/java/org/springframework/batch/item/xml/StaxEventItemWriter.java

@@ -87,9 +87,6 @@ public class StaxEventItemWriter<T> extends ExecutionContextUserSupport implemen
 	// root element attributes
 	private Map<String, String> 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<T> 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<T> extends ExecutionContextUserSupport implemen
 			restarted = true;
 		}
 
-		open(startAtPosition);
+		open(startAtPosition, restarted);
 
 		if (startAtPosition == 0) {
 			try {
@@ -281,7 +279,7 @@ public class StaxEventItemWriter<T> 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<T> 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<T> extends ExecutionContextUserSupport implemen
 		catch (XMLStreamException e) {
 			throw new FlushFailedException("Failed to flush the events", e);
 		}
+
 	}
 
 	/**

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 = "<?xml version=\"1.0\" encoding=\"UTF-8\"?><root>";
+
+	final private String xmlDocEnd = "</root>";
+
+	@Override
+	@Before
+	public void setUp() throws Exception {
+		delegate = new StaxEventItemWriter<String>() {
+			{
+				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 + "<prefix:4></prefix:4>" + xmlDocEnd, readFile(part2));
+		assertEquals(xmlDocStart + "<prefix:1></prefix:1><prefix:2></prefix:2><prefix:3></prefix:3>" + 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 + "<prefix:4></prefix:4><prefix:5></prefix:5>" + xmlDocEnd, readFile(part2));
+		assertEquals(xmlDocStart
+				+ "<prefix:6></prefix:6><prefix:7></prefix:7><prefix:8></prefix:8><prefix:9></prefix:9>" + xmlDocEnd,
+				readFile(part3));
+	}
+}

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<String> tested = new MultiResourceItemWriter<String>();
+
+	protected File file;
+
+	protected ResourceSuffixCreator suffixCreator = new SimpleResourceSuffixCreator();
+
+	protected ResourceAwareItemWriterItemStream<String> 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();
+	}
+}

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<String> tested = new MultiResourceItemWriter<String>();
-
-	private File file;
-	
-	private ResourceSuffixCreator suffixCreator = new SimpleResourceSuffixCreator();
-
-	private ResourceAwareItemWriterItemStream<String> delegate = new FlatFileItemWriter<String>() {
-		{
-			setLineAggregator(new PassThroughLineAggregator<String>());
-		}
-	};
-
-	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<String>() {
+			{
+				setLineAggregator(new PassThroughLineAggregator<String>());
+			}
+		};
+		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();
-	}
 }