001/*
002 * Copyright 2006-2013 the original author or authors.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 *      https://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016
017package org.springframework.batch.core.partition.support;
018
019import java.io.IOException;
020import java.util.HashMap;
021import java.util.Map;
022
023import org.springframework.batch.item.ExecutionContext;
024import org.springframework.core.io.Resource;
025import org.springframework.util.Assert;
026
027/**
028 * Implementation of {@link Partitioner} that locates multiple resources and
029 * associates their file names with execution context keys. Creates an
030 * {@link ExecutionContext} per resource, and labels them as
031 * <code>{partition0, partition1, ..., partitionN}</code>. The grid size is
032 * ignored.
033 *
034 * @author Dave Syer
035 * @since 2.0
036 */
037public class MultiResourcePartitioner implements Partitioner {
038
039        private static final String DEFAULT_KEY_NAME = "fileName";
040
041        private static final String PARTITION_KEY = "partition";
042
043        private Resource[] resources = new Resource[0];
044
045        private String keyName = DEFAULT_KEY_NAME;
046
047        /**
048         * The resources to assign to each partition. In Spring configuration you
049         * can use a pattern to select multiple resources.
050         * @param resources the resources to use
051         */
052        public void setResources(Resource[] resources) {
053                this.resources = resources;
054        }
055
056        /**
057         * The name of the key for the file name in each {@link ExecutionContext}.
058         * Defaults to "fileName".
059         * @param keyName the value of the key
060         */
061        public void setKeyName(String keyName) {
062                this.keyName = keyName;
063        }
064
065        /**
066         * Assign the filename of each of the injected resources to an
067         * {@link ExecutionContext}.
068         *
069         * @see Partitioner#partition(int)
070         */
071        @Override
072        public Map<String, ExecutionContext> partition(int gridSize) {
073                Map<String, ExecutionContext> map = new HashMap<String, ExecutionContext>(gridSize);
074                int i = 0;
075                for (Resource resource : resources) {
076                        ExecutionContext context = new ExecutionContext();
077                        Assert.state(resource.exists(), "Resource does not exist: "+resource);
078                        try {
079                                context.putString(keyName, resource.getURL().toExternalForm());
080                        }
081                        catch (IOException e) {
082                                throw new IllegalArgumentException("File could not be located for: "+resource, e);
083                        }
084                        map.put(PARTITION_KEY + i, context);
085                        i++;
086                }
087                return map;
088        }
089
090}