1 | /* |
2 | * Copyright 2006-2007 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 | package org.springframework.batch.item.database.support; |
17 | |
18 | import java.sql.PreparedStatement; |
19 | import java.sql.SQLException; |
20 | import java.util.List; |
21 | |
22 | import org.springframework.batch.item.ExecutionContext; |
23 | import org.springframework.batch.item.ExecutionContextUserSupport; |
24 | import org.springframework.batch.item.database.DrivingQueryItemReader; |
25 | import org.springframework.batch.item.database.ItemPreparedStatementSetter; |
26 | import org.springframework.batch.item.database.KeyCollector; |
27 | import org.springframework.jdbc.core.ColumnMapRowMapper; |
28 | import org.springframework.jdbc.core.JdbcTemplate; |
29 | import org.springframework.jdbc.core.PreparedStatementSetter; |
30 | import org.springframework.jdbc.core.RowMapper; |
31 | import org.springframework.util.Assert; |
32 | import org.springframework.util.ClassUtils; |
33 | import org.springframework.util.StringUtils; |
34 | |
35 | /** |
36 | * <p> |
37 | * JDBC implementation of the {@link KeyCollector} interface that works for |
38 | * composite keys. (i.e. keys represented by multiple columns) A SQL query to be |
39 | * used to return the keys and a {@link ItemPreparedStatementSetter} to map each |
40 | * row in the result set to an Object must be set in order to work correctly. |
41 | * </p> |
42 | * |
43 | * @author Lucas Ward |
44 | * @see DrivingQueryItemReader |
45 | * @see ItemPreparedStatementSetter |
46 | */ |
47 | public class MultipleColumnJdbcKeyCollector extends ExecutionContextUserSupport implements KeyCollector { |
48 | |
49 | private static final String CURRENT_KEY = "current.key"; |
50 | |
51 | private JdbcTemplate jdbcTemplate; |
52 | |
53 | private RowMapper keyMapper = new ColumnMapRowMapper(); |
54 | |
55 | private ItemPreparedStatementSetter preparedStatementSetter = new ColumnMapItemPreparedStatementSetter(); |
56 | |
57 | private String sql; |
58 | |
59 | private String restartSql; |
60 | |
61 | public MultipleColumnJdbcKeyCollector() { |
62 | setName(ClassUtils.getShortName(MultipleColumnJdbcKeyCollector.class)); |
63 | } |
64 | |
65 | /** |
66 | * Construct a new ItemReader. |
67 | * |
68 | * @param jdbcTemplate |
69 | * @param sql - SQL statement that returns all keys to process. |
70 | * object. |
71 | */ |
72 | public MultipleColumnJdbcKeyCollector(JdbcTemplate jdbcTemplate, String sql) { |
73 | this(); |
74 | Assert.notNull(jdbcTemplate, "The JdbcTemplate must not be null."); |
75 | Assert.hasText(sql, "The sql statement must not be null or empty."); |
76 | this.jdbcTemplate = jdbcTemplate; |
77 | this.sql = sql; |
78 | } |
79 | |
80 | /* |
81 | * (non-Javadoc) |
82 | * @see org.springframework.batch.io.sql.scratch.AbstractDrivingQueryItemReader#retrieveKeys() |
83 | */ |
84 | public List retrieveKeys(ExecutionContext executionContext) { |
85 | |
86 | Assert.state(keyMapper != null, "KeyMapper must not be null."); |
87 | Assert.state(StringUtils.hasText(restartSql), "The RestartQuery must not be null or empty" |
88 | + " in order to restart."); |
89 | |
90 | if (executionContext.size() > 0) { |
91 | Object key = executionContext.get(getKey(CURRENT_KEY)); |
92 | return jdbcTemplate.query(restartSql, new PreparedStatementSetterKeyWrapper(key, preparedStatementSetter), keyMapper); |
93 | } |
94 | else { |
95 | return jdbcTemplate.query(sql, keyMapper); |
96 | } |
97 | } |
98 | |
99 | /* |
100 | * (non-Javadoc) |
101 | * @see org.springframework.batch.io.driving.KeyGenerator#getKeyAsExecutionContext(java.lang.Object) |
102 | */ |
103 | public void updateContext(Object key, ExecutionContext executionContext) { |
104 | Assert.notNull(key, "The key must not be null"); |
105 | Assert.notNull(executionContext, "The ExecutionContext must not be null"); |
106 | executionContext.put(getKey(CURRENT_KEY), key); |
107 | } |
108 | |
109 | /** |
110 | * Set the query to use to retrieve keys in order to restore the previous |
111 | * state for restart. |
112 | * |
113 | * @param restartQuery |
114 | */ |
115 | public void setRestartSql(String restartQuery) { |
116 | this.restartSql = restartQuery; |
117 | } |
118 | |
119 | /* |
120 | * (non-Javadoc) |
121 | * @see org.springframework.beans.factory.InitializingBean#afterPropertiesSet() |
122 | */ |
123 | public void afterPropertiesSet() throws Exception { |
124 | Assert.notNull(jdbcTemplate, "The JdbcTemplate must not be null."); |
125 | Assert.hasText(sql, "The DrivingQuery must not be null or empty."); |
126 | Assert.notNull(keyMapper, "The key RowMapper must not be null."); |
127 | } |
128 | |
129 | /** |
130 | * Set the {@link RowMapper} to be used to map a result set |
131 | * to keys. |
132 | * |
133 | * @param keyMapper |
134 | */ |
135 | public void setKeyMapper(RowMapper keyMapper) { |
136 | this.keyMapper = keyMapper; |
137 | } |
138 | |
139 | /** |
140 | * Set the sql statement used to generate the keys list. |
141 | * |
142 | * @param sql |
143 | */ |
144 | public void setSql(String sql) { |
145 | this.sql = sql; |
146 | } |
147 | |
148 | public void setJdbcTemplate(JdbcTemplate jdbcTemplate) { |
149 | this.jdbcTemplate = jdbcTemplate; |
150 | } |
151 | |
152 | public void setPreparedStatementSetter( |
153 | ItemPreparedStatementSetter preparedStatementSetter) { |
154 | this.preparedStatementSetter = preparedStatementSetter; |
155 | } |
156 | |
157 | private static class PreparedStatementSetterKeyWrapper implements PreparedStatementSetter{ |
158 | |
159 | private Object key; |
160 | private ItemPreparedStatementSetter pss; |
161 | |
162 | public PreparedStatementSetterKeyWrapper(Object key, ItemPreparedStatementSetter pss) { |
163 | this.key = key; |
164 | this.pss = pss; |
165 | } |
166 | |
167 | public void setValues(PreparedStatement ps) throws SQLException { |
168 | pss.setValues(key, ps); |
169 | } |
170 | } |
171 | } |