| 1 | /* |
| 2 | * Copyright 2012 the original author or authors. |
| 3 | * |
| 4 | * Licensed under the Apache License, Version 2.0 (the "License"); |
| 5 | * you may not use this file except in compliance with the License. |
| 6 | * You may obtain a copy of the License at |
| 7 | * |
| 8 | * http://www.apache.org/licenses/LICENSE-2.0 |
| 9 | * |
| 10 | * Unless required by applicable law or agreed to in writing, software |
| 11 | * distributed under the License is distributed on an "AS IS" BASIS, |
| 12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 13 | * See the License for the specific language governing permissions and |
| 14 | * limitations under the License. |
| 15 | */ |
| 16 | |
| 17 | package org.springframework.batch.item.data; |
| 18 | |
| 19 | import com.mongodb.util.JSON; |
| 20 | import org.springframework.batch.item.ExecutionContext; |
| 21 | import org.springframework.batch.item.ItemReader; |
| 22 | import org.springframework.beans.factory.InitializingBean; |
| 23 | import org.springframework.data.domain.PageRequest; |
| 24 | import org.springframework.data.domain.Pageable; |
| 25 | import org.springframework.data.domain.Sort; |
| 26 | import org.springframework.data.mongodb.core.MongoOperations; |
| 27 | import org.springframework.data.mongodb.core.query.BasicQuery; |
| 28 | import org.springframework.data.mongodb.core.query.Query; |
| 29 | import org.springframework.util.Assert; |
| 30 | import org.springframework.util.ClassUtils; |
| 31 | import org.springframework.util.StringUtils; |
| 32 | |
| 33 | import java.util.ArrayList; |
| 34 | import java.util.Iterator; |
| 35 | import java.util.List; |
| 36 | import java.util.Map; |
| 37 | import java.util.regex.Matcher; |
| 38 | import java.util.regex.Pattern; |
| 39 | |
| 40 | /** |
| 41 | * <p> |
| 42 | * Restartable {@link ItemReader} that reads documents from MongoDB |
| 43 | * via a paging technique. |
| 44 | * </p> |
| 45 | * |
| 46 | * <p> |
| 47 | * It executes the JSON {@link #setQuery(String)} to retrieve the requested |
| 48 | * documents. The query is executed using paged requests specified in the |
| 49 | * {@link #setPageSize(int)}. Additional pages are requested as needed to |
| 50 | * provide data when the {@link #read()} method is called. |
| 51 | * </p> |
| 52 | * |
| 53 | * <p> |
| 54 | * The JSON query provided supports parameter substitution via ?<index> |
| 55 | * placeholders where the <index> indicates the index of the |
| 56 | * parameterValue to substitute. |
| 57 | * </p> |
| 58 | * |
| 59 | * <p> |
| 60 | * The implementation is thread-safe between calls to |
| 61 | * {@link #open(ExecutionContext)}, but remember to use <code>saveState=false</code> |
| 62 | * if used in a multi-threaded client (no restart available. |
| 63 | * </p> |
| 64 | * |
| 65 | * |
| 66 | * @author Michael Minella |
| 67 | */ |
| 68 | public class MongoItemReader<T> extends AbstractPaginatedDataItemReader<T> implements InitializingBean { |
| 69 | |
| 70 | private static final Pattern PLACEHOLDER = Pattern.compile("\\?(\\d+)"); |
| 71 | private MongoOperations template; |
| 72 | private String query; |
| 73 | private Class<? extends T> type; |
| 74 | private Sort sort; |
| 75 | private String hint; |
| 76 | private String fields; |
| 77 | private String collection; |
| 78 | private List<Object> parameterValues; |
| 79 | |
| 80 | public MongoItemReader() { |
| 81 | super(); |
| 82 | setName(ClassUtils.getShortName(MongoItemReader.class)); |
| 83 | } |
| 84 | |
| 85 | /** |
| 86 | * Used to perform operations against the MongoDB instance. Also |
| 87 | * handles the mapping of documents to objects. |
| 88 | * |
| 89 | * @param template the MongoOperations instance to use |
| 90 | * @see MongoOperations |
| 91 | */ |
| 92 | public void setTemplate(MongoOperations template) { |
| 93 | this.template = template; |
| 94 | } |
| 95 | |
| 96 | /** |
| 97 | * A JSON formatted MongoDB query. Parameterization of the provided query is allowed |
| 98 | * via ?<index> placeholders where the <index> indicates the index of the |
| 99 | * parameterValue to substitute. |
| 100 | * |
| 101 | * @param query JSON formatted Mongo query |
| 102 | */ |
| 103 | public void setQuery(String query) { |
| 104 | this.query = query; |
| 105 | } |
| 106 | |
| 107 | /** |
| 108 | * The type of object to be returned for each {@link #read()} call. |
| 109 | * |
| 110 | * @param type the type of object to return |
| 111 | */ |
| 112 | public void setTargetType(Class<? extends T> type) { |
| 113 | this.type = type; |
| 114 | } |
| 115 | |
| 116 | /** |
| 117 | * {@link List} of values to be substituted in for each of the |
| 118 | * parameters in the query. |
| 119 | * |
| 120 | * @param parameterValues |
| 121 | */ |
| 122 | public void setParameterValues(List<Object> parameterValues) { |
| 123 | this.parameterValues = parameterValues; |
| 124 | } |
| 125 | |
| 126 | /** |
| 127 | * JSON defining the fields to be returned from the matching documents |
| 128 | * by MongoDB. |
| 129 | * |
| 130 | * @param fields JSON string that identifies the fields to sorty by. |
| 131 | */ |
| 132 | public void setFields(String fields) { |
| 133 | this.fields = fields; |
| 134 | } |
| 135 | |
| 136 | /** |
| 137 | * {@link Map} of property names/{@link org.springframework.data.domain.Sort.Direction} values to |
| 138 | * sort the input by. |
| 139 | * |
| 140 | * @param sorts map of properties and direction to sort each. |
| 141 | */ |
| 142 | public void setSort(Map<String, Sort.Direction> sorts) { |
| 143 | this.sort = convertToSort(sorts); |
| 144 | } |
| 145 | |
| 146 | /** |
| 147 | * @param collection Mongo collection to be queried. |
| 148 | */ |
| 149 | public void setCollection(String collection) { |
| 150 | this.collection = collection; |
| 151 | } |
| 152 | |
| 153 | /** |
| 154 | * JSON String telling MongoDB what index to use. |
| 155 | * |
| 156 | * @param hint string indicating what index to use. |
| 157 | */ |
| 158 | public void setHint(String hint) { |
| 159 | this.hint = hint; |
| 160 | } |
| 161 | |
| 162 | @Override |
| 163 | @SuppressWarnings("unchecked") |
| 164 | protected Iterator<T> doPageRead() { |
| 165 | |
| 166 | Pageable pageRequest = new PageRequest(page, pageSize, sort); |
| 167 | |
| 168 | String populatedQuery = replacePlaceholders(query, parameterValues); |
| 169 | |
| 170 | Query mongoQuery = null; |
| 171 | |
| 172 | if(StringUtils.hasText(fields)) { |
| 173 | mongoQuery = new BasicQuery(populatedQuery, fields); |
| 174 | } |
| 175 | else { |
| 176 | mongoQuery = new BasicQuery(populatedQuery); |
| 177 | } |
| 178 | |
| 179 | mongoQuery.with(pageRequest); |
| 180 | |
| 181 | if(StringUtils.hasText(hint)) { |
| 182 | mongoQuery.withHint(hint); |
| 183 | } |
| 184 | |
| 185 | if(StringUtils.hasText(collection)) { |
| 186 | return (Iterator<T>) template.find(mongoQuery, type, collection).iterator(); |
| 187 | } else { |
| 188 | return (Iterator<T>) template.find(mongoQuery, type).iterator(); |
| 189 | } |
| 190 | } |
| 191 | |
| 192 | /** |
| 193 | * Checks mandatory properties |
| 194 | * |
| 195 | * @see InitializingBean#afterPropertiesSet() |
| 196 | */ |
| 197 | public void afterPropertiesSet() throws Exception { |
| 198 | Assert.state(template != null, "An implementation of MongoOperations is required."); |
| 199 | Assert.state(type != null, "A type to convert the input into is required."); |
| 200 | Assert.state(query != null, "A query is required."); |
| 201 | Assert.state(sort != null, "A sort is required."); |
| 202 | } |
| 203 | |
| 204 | // Copied from StringBasedMongoQuery...is there a place where this type of logic is already exposed? |
| 205 | private String replacePlaceholders(String input, List<Object> values) { |
| 206 | Matcher matcher = PLACEHOLDER.matcher(input); |
| 207 | String result = input; |
| 208 | |
| 209 | while (matcher.find()) { |
| 210 | String group = matcher.group(); |
| 211 | int index = Integer.parseInt(matcher.group(1)); |
| 212 | result = result.replace(group, getParameterWithIndex(values, index)); |
| 213 | } |
| 214 | |
| 215 | return result; |
| 216 | } |
| 217 | |
| 218 | // Copied from StringBasedMongoQuery...is there a place where this type of logic is already exposed? |
| 219 | private String getParameterWithIndex(List<Object> values, int index) { |
| 220 | return JSON.serialize(values.get(index)); |
| 221 | } |
| 222 | |
| 223 | private Sort convertToSort(Map<String, Sort.Direction> sorts) { |
| 224 | List<Sort.Order> sortValues = new ArrayList<Sort.Order>(); |
| 225 | |
| 226 | for (Map.Entry<String, Sort.Direction> curSort : sorts.entrySet()) { |
| 227 | sortValues.add(new Sort.Order(curSort.getValue(), curSort.getKey())); |
| 228 | } |
| 229 | |
| 230 | return new Sort(sortValues); |
| 231 | } |
| 232 | } |