diff --git a/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/file/builder/MultiResourceItemWriterBuilder.java b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/file/builder/MultiResourceItemWriterBuilder.java new file mode 100644 index 000000000..bbcda08ac --- /dev/null +++ b/spring-batch-infrastructure/src/main/java/org/springframework/batch/item/file/builder/MultiResourceItemWriterBuilder.java @@ -0,0 +1,157 @@ +/* + * Copyright 2017 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.batch.item.file.builder; + +import org.springframework.batch.item.ExecutionContext; +import org.springframework.batch.item.file.MultiResourceItemWriter; +import org.springframework.batch.item.file.ResourceAwareItemWriterItemStream; +import org.springframework.batch.item.file.ResourceSuffixCreator; +import org.springframework.core.io.Resource; +import org.springframework.util.Assert; + +/** + * A builder implementation for the {@link MultiResourceItemWriter}. + * + * @author Glenn Renfro + * @since 4.0 + * @see MultiResourceItemWriter + */ +public class MultiResourceItemWriterBuilder { + + private Resource resource; + + private ResourceAwareItemWriterItemStream delegate; + + private int itemCountLimitPerResource = Integer.MAX_VALUE; + + private ResourceSuffixCreator suffixCreator; + + private boolean saveState = true; + + private String name; + + /** + * Allows customization of the suffix of the created resources based on the index. + * + * @param suffixCreator the customizable ResourceSuffixCreator to use. + * @return The current instance of the builder. + * @see MultiResourceItemWriter#setResourceSuffixCreator(ResourceSuffixCreator) + */ + public MultiResourceItemWriterBuilder resourceSuffixCreator(ResourceSuffixCreator suffixCreator) { + this.suffixCreator = suffixCreator; + + return this; + } + + /** + * After this limit is exceeded the next chunk will be written into newly created + * resource. + * + * @param itemCountLimitPerResource the max numbers of items to be written per chunk. + * @return The current instance of the builder. + * @see MultiResourceItemWriter#setItemCountLimitPerResource(int) + */ + public MultiResourceItemWriterBuilder itemCountLimitPerResource(int itemCountLimitPerResource) { + this.itemCountLimitPerResource = itemCountLimitPerResource; + + return this; + } + + /** + * Delegate used for actual writing of the output. + * @param delegate The delegate to use for writing. + * @return The current instance of the builder. + * @see MultiResourceItemWriter#setDelegate(ResourceAwareItemWriterItemStream) + */ + public MultiResourceItemWriterBuilder delegate(ResourceAwareItemWriterItemStream delegate) { + this.delegate = delegate; + + return this; + } + + /** + * 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 MultiResourceItemWriter#setResourceSuffixCreator(ResourceSuffixCreator)}. + * + * @param resource the prototype resource to use as the basis for creating resources. + * @return The current instance of the builder. + * @see MultiResourceItemWriter#setResource(Resource) + */ + public MultiResourceItemWriterBuilder resource(Resource resource) { + this.resource = resource; + + return this; + } + + /** + * Set the boolean indicating whether or not state should be saved in the provided + * {@link ExecutionContext} during the delegate call to update. + * + * @param saveState true to update ExecutionContext. False do not update + * ExecutionContext. + * @return The current instance of the builder. + * @see MultiResourceItemWriter#setSaveState(boolean) + */ + public MultiResourceItemWriterBuilder saveState(boolean saveState) { + this.saveState = saveState; + + return this; + } + + /** + * The name of the component which will be used as a stem for keys in the + * {@link ExecutionContext}. Subclasses should provide a default value, e.g. + * the short form of the class name. + * + * @param name the name for the component. + * @return The current instance of the builder. + * @see MultiResourceItemWriter#setName(String) + */ + public MultiResourceItemWriterBuilder name(String name) { + this.name = name; + + return this; + } + + /** + * Builds the {@link MultiResourceItemWriter}. + * + * @return a {@link MultiResourceItemWriter} + */ + public MultiResourceItemWriter build() { + Assert.notNull(this.resource, "resource is required."); + Assert.notNull(this.delegate, "delegate is required."); + + if(this.saveState) { + org.springframework.util.Assert.hasText(this.name, "A name is required when saveState is true."); + } + + MultiResourceItemWriter writer = new MultiResourceItemWriter<>(); + writer.setResource(this.resource); + writer.setDelegate(this.delegate); + writer.setItemCountLimitPerResource(this.itemCountLimitPerResource); + if(this.suffixCreator != null) { + writer.setResourceSuffixCreator(this.suffixCreator); + } + writer.setSaveState(this.saveState); + writer.setName(this.name); + + return writer; + } + +} diff --git a/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/file/builder/MultiResourceItemWriterBuilderTests.java b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/file/builder/MultiResourceItemWriterBuilderTests.java new file mode 100644 index 000000000..99d6a2733 --- /dev/null +++ b/spring-batch-infrastructure/src/test/java/org/springframework/batch/item/file/builder/MultiResourceItemWriterBuilderTests.java @@ -0,0 +1,266 @@ +/* + * Copyright 2017 the original author or authors. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package org.springframework.batch.item.file.builder; + +import java.io.BufferedReader; +import java.io.File; +import java.io.FileReader; +import java.util.Arrays; + +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +import org.springframework.batch.item.ExecutionContext; +import org.springframework.batch.item.file.FlatFileItemWriter; +import org.springframework.batch.item.file.MultiResourceItemWriter; +import org.springframework.batch.item.file.MultiResourceItemWriterFlatFileTests; +import org.springframework.batch.item.file.ResourceSuffixCreator; +import org.springframework.batch.item.file.SimpleResourceSuffixCreator; +import org.springframework.batch.item.file.transform.PassThroughLineAggregator; +import org.springframework.core.io.FileSystemResource; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +/** + * @author Glenn Renfro + */ +public class MultiResourceItemWriterBuilderTests { + + private MultiResourceItemWriter writer; + + private File file; + + private ResourceSuffixCreator suffixCreator = new ResourceSuffixCreator() { + @Override + public String getSuffix(int index) { + return "A" + index; + } + }; + + private ExecutionContext executionContext = new ExecutionContext(); + + private FlatFileItemWriter delegate; + + @Before + public void setUp() throws Exception { + this.delegate = new FlatFileItemWriter<>(); + this.delegate.setLineAggregator(new PassThroughLineAggregator<>()); + this.file = File.createTempFile(MultiResourceItemWriterFlatFileTests.class.getSimpleName(), null); + this.writer = null; + } + + @After + public void tearDown() { + if (this.writer != null) { + this.writer.close(); + } + } + + @Test + public void testBasicMultiResourceWriteScenario() throws Exception { + + this.writer = new MultiResourceItemWriterBuilder().delegate(this.delegate) + .resource(new FileSystemResource(this.file)).resourceSuffixCreator(this.suffixCreator) + .itemCountLimitPerResource(2).saveState(true).name("foo").build(); + + this.writer.open(this.executionContext); + + this.writer.write(Arrays.asList("1", "2", "3")); + + File part1 = new File(this.file.getAbsolutePath() + this.suffixCreator.getSuffix(1)); + assertTrue(part1.exists()); + assertEquals("123", readFile(part1)); + + this.writer.write(Arrays.asList("4")); + File part2 = new File(this.file.getAbsolutePath() + this.suffixCreator.getSuffix(2)); + assertTrue(part2.exists()); + assertEquals("4", readFile(part2)); + + this.writer.write(Arrays.asList("5")); + assertEquals("45", readFile(part2)); + + this.writer.write(Arrays.asList("6", "7", "8", "9")); + File part3 = new File(this.file.getAbsolutePath() + this.suffixCreator.getSuffix(3)); + assertTrue(part3.exists()); + assertEquals("6789", readFile(part3)); + } + + @Test + public void testBasicDefaultSuffixCreator() throws Exception { + + SimpleResourceSuffixCreator simpleResourceSuffixCreator = new SimpleResourceSuffixCreator(); + this.writer = new MultiResourceItemWriterBuilder().delegate(this.delegate) + .resource(new FileSystemResource(this.file)).itemCountLimitPerResource(2).saveState(true).name("foo") + .build(); + + this.writer.open(this.executionContext); + + this.writer.write(Arrays.asList("1", "2", "3")); + + File part1 = new File(this.file.getAbsolutePath() + simpleResourceSuffixCreator.getSuffix(1)); + assertTrue(part1.exists()); + assertEquals("123", readFile(part1)); + + this.writer.write(Arrays.asList("4")); + File part2 = new File(this.file.getAbsolutePath() + simpleResourceSuffixCreator.getSuffix(2)); + assertTrue(part2.exists()); + assertEquals("4", readFile(part2)); + } + + @Test + public void testUpdateAfterDelegateClose() throws Exception { + + this.writer = new MultiResourceItemWriterBuilder().delegate(this.delegate) + .resource(new FileSystemResource(this.file)).resourceSuffixCreator(this.suffixCreator) + .itemCountLimitPerResource(2).saveState(true).name("foo").build(); + + this.writer.update(this.executionContext); + assertEquals(0, this.executionContext.getInt(this.writer.getExecutionContextKey("resource.item.count"))); + assertEquals(1, this.executionContext.getInt(this.writer.getExecutionContextKey("resource.index"))); + this.writer.write(Arrays.asList("1", "2", "3")); + this.writer.update(this.executionContext); + assertEquals(0, this.executionContext.getInt(this.writer.getExecutionContextKey("resource.item.count"))); + assertEquals(2, this.executionContext.getInt(this.writer.getExecutionContextKey("resource.index"))); + } + + @Test + public void testRestart() throws Exception { + + this.writer = new MultiResourceItemWriterBuilder().delegate(this.delegate) + .resource(new FileSystemResource(this.file)).resourceSuffixCreator(this.suffixCreator) + .itemCountLimitPerResource(2).saveState(true).name("foo").build(); + + this.writer.write(Arrays.asList("1", "2", "3")); + + File part1 = new File(this.file.getAbsolutePath() + this.suffixCreator.getSuffix(1)); + assertTrue(part1.exists()); + assertEquals("123", readFile(part1)); + + this.writer.write(Arrays.asList("4")); + File part2 = new File(this.file.getAbsolutePath() + this.suffixCreator.getSuffix(2)); + assertTrue(part2.exists()); + assertEquals("4", readFile(part2)); + + this.writer.update(this.executionContext); + this.writer.close(); + this.writer.open(this.executionContext); + + this.writer.write(Arrays.asList("5")); + assertEquals("45", readFile(part2)); + + this.writer.write(Arrays.asList("6", "7", "8", "9")); + File part3 = new File(this.file.getAbsolutePath() + this.suffixCreator.getSuffix(3)); + assertTrue(part3.exists()); + assertEquals("6789", readFile(part3)); + } + + @Test + public void testRestartNoSaveState() throws Exception { + + this.writer = new MultiResourceItemWriterBuilder().delegate(this.delegate) + .resource(new FileSystemResource(this.file)).resourceSuffixCreator(this.suffixCreator) + .itemCountLimitPerResource(2).saveState(false).name("foo").build(); + + this.writer.write(Arrays.asList("1", "2", "3")); + + File part1 = new File(this.file.getAbsolutePath() + this.suffixCreator.getSuffix(1)); + assertTrue(part1.exists()); + assertEquals("123", readFile(part1)); + + this.writer.write(Arrays.asList("4")); + File part2 = new File(this.file.getAbsolutePath() + this.suffixCreator.getSuffix(2)); + assertTrue(part2.exists()); + assertEquals("4", readFile(part2)); + + this.writer.update(this.executionContext); + this.writer.close(); + this.writer.open(this.executionContext); + + this.writer.write(Arrays.asList("5")); + assertEquals("4", readFile(part2)); + + this.writer.write(Arrays.asList("6", "7", "8", "9")); + File part3 = new File(this.file.getAbsolutePath() + this.suffixCreator.getSuffix(1)); + assertTrue(part3.exists()); + assertEquals("56789", readFile(part3)); + } + + @Test + public void testSaveStateNoName() { + try { + new MultiResourceItemWriterBuilder().delegate(this.delegate) + .resource(new FileSystemResource(this.file)).resourceSuffixCreator(this.suffixCreator) + .itemCountLimitPerResource(2).saveState(true).build(); + + fail("IllegalArgumentException should have been thrown"); + } + catch (IllegalArgumentException ise) { + assertEquals("IllegalArgumentException message did not match the expected result.", + "A name is required when saveState is true.", ise.getMessage()); + } + } + + @Test + public void testNoResource() throws Exception { + try { + new MultiResourceItemWriterBuilder().delegate(this.delegate) + .resourceSuffixCreator(this.suffixCreator).itemCountLimitPerResource(2).build(); + + fail("IllegalArgumentException should have been thrown"); + } + catch (IllegalArgumentException ise) { + assertEquals("IllegalArgumentException message did not match the expected result.", "resource is required.", + ise.getMessage()); + } + } + + @Test + public void testNoDelegateNoName() { + try { + new MultiResourceItemWriterBuilder().resource(new FileSystemResource(this.file)) + .resourceSuffixCreator(this.suffixCreator).itemCountLimitPerResource(2).saveState(false).build(); + + fail("IllegalArgumentException should have been thrown"); + } + catch (IllegalArgumentException ise) { + assertEquals("IllegalArgumentException message did not match the expected result.", "delegate is required.", + ise.getMessage()); + } + } + + private String readFile(File f) throws Exception { + BufferedReader reader = new BufferedReader(new FileReader(f)); + StringBuilder result = new StringBuilder(); + try { + while (true) { + String line = reader.readLine(); + if (line == null) { + break; + } + result.append(line); + } + } + finally { + reader.close(); + } + return result.toString(); + } + +}