View Javadoc
1   /**
2    *    Copyright 2009-2020 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.apache.ibatis.executor.keygen;
17  
18  import java.sql.ResultSet;
19  import java.sql.ResultSetMetaData;
20  import java.sql.SQLException;
21  import java.sql.Statement;
22  import java.util.AbstractMap;
23  import java.util.ArrayList;
24  import java.util.Arrays;
25  import java.util.Collection;
26  import java.util.HashMap;
27  import java.util.Iterator;
28  import java.util.List;
29  import java.util.Map;
30  import java.util.Map.Entry;
31  import java.util.Set;
32  
33  import org.apache.ibatis.binding.MapperMethod.ParamMap;
34  import org.apache.ibatis.executor.Executor;
35  import org.apache.ibatis.executor.ExecutorException;
36  import org.apache.ibatis.mapping.MappedStatement;
37  import org.apache.ibatis.reflection.ArrayUtil;
38  import org.apache.ibatis.reflection.MetaObject;
39  import org.apache.ibatis.reflection.ParamNameResolver;
40  import org.apache.ibatis.session.Configuration;
41  import org.apache.ibatis.session.defaults.DefaultSqlSession.StrictMap;
42  import org.apache.ibatis.type.JdbcType;
43  import org.apache.ibatis.type.TypeHandler;
44  import org.apache.ibatis.type.TypeHandlerRegistry;
45  
46  /**
47   * @author Clinton Begin
48   * @author Kazuki Shimizu
49   */
50  public class Jdbc3KeyGenerator implements KeyGenerator {
51  
52    private static final String SECOND_GENERIC_PARAM_NAME = ParamNameResolver.GENERIC_NAME_PREFIX + "2";
53  
54    /**
55     * A shared instance.
56     *
57     * @since 3.4.3
58     */
59    public static final Jdbc3KeyGenerator3KeyGenerator.html#Jdbc3KeyGenerator">Jdbc3KeyGenerator INSTANCE = new Jdbc3KeyGenerator();
60  
61    private static final String MSG_TOO_MANY_KEYS = "Too many keys are generated. There are only %d target objects. "
62        + "You either specified a wrong 'keyProperty' or encountered a driver bug like #1523.";
63  
64    @Override
65    public void processBefore(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {
66      // do nothing
67    }
68  
69    @Override
70    public void processAfter(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {
71      processBatch(ms, stmt, parameter);
72    }
73  
74    public void processBatch(MappedStatement ms, Statement stmt, Object parameter) {
75      final String[] keyProperties = ms.getKeyProperties();
76      if (keyProperties == null || keyProperties.length == 0) {
77        return;
78      }
79      try (ResultSet rs = stmt.getGeneratedKeys()) {
80        final ResultSetMetaData rsmd = rs.getMetaData();
81        final Configuration configuration = ms.getConfiguration();
82        if (rsmd.getColumnCount() < keyProperties.length) {
83          // Error?
84        } else {
85          assignKeys(configuration, rs, rsmd, keyProperties, parameter);
86        }
87      } catch (Exception e) {
88        throw new ExecutorException("Error getting generated key or setting result to parameter object. Cause: " + e, e);
89      }
90    }
91  
92    @SuppressWarnings("unchecked")
93    private void assignKeys(Configuration configuration, ResultSet rs, ResultSetMetaData rsmd, String[] keyProperties,
94        Object parameter) throws SQLException {
95      if (parameter instanceof ParamMap || parameter instanceof StrictMap) {
96        // Multi-param or single param with @Param
97        assignKeysToParamMap(configuration, rs, rsmd, keyProperties, (Map<String, ?>) parameter);
98      } else if (parameter instanceof ArrayList && !((ArrayList<?>) parameter).isEmpty()
99          && ((ArrayList<?>) parameter).get(0) instanceof ParamMap) {
100       // Multi-param or single param with @Param in batch operation
101       assignKeysToParamMapList(configuration, rs, rsmd, keyProperties, (ArrayList<ParamMap<?>>) parameter);
102     } else {
103       // Single param without @Param
104       assignKeysToParam(configuration, rs, rsmd, keyProperties, parameter);
105     }
106   }
107 
108   private void assignKeysToParam(Configuration configuration, ResultSet rs, ResultSetMetaData rsmd,
109       String[] keyProperties, Object parameter) throws SQLException {
110     Collection<?> params = collectionize(parameter);
111     if (params.isEmpty()) {
112       return;
113     }
114     List<KeyAssigner> assignerList = new ArrayList<>();
115     for (int i = 0; i < keyProperties.length; i++) {
116       assignerList.add(new KeyAssigner(configuration, rsmd, i + 1, null, keyProperties[i]));
117     }
118     Iterator<?> iterator = params.iterator();
119     while (rs.next()) {
120       if (!iterator.hasNext()) {
121         throw new ExecutorException(String.format(MSG_TOO_MANY_KEYS, params.size()));
122       }
123       Object param = iterator.next();
124       assignerList.forEach(x -> x.assign(rs, param));
125     }
126   }
127 
128   private void assignKeysToParamMapList(Configuration configuration, ResultSet rs, ResultSetMetaData rsmd,
129       String[] keyProperties, ArrayList<ParamMap<?>> paramMapList) throws SQLException {
130     Iterator<ParamMap<?>> iterator = paramMapList.iterator();
131     List<KeyAssigner> assignerList = new ArrayList<>();
132     long counter = 0;
133     while (rs.next()) {
134       if (!iterator.hasNext()) {
135         throw new ExecutorException(String.format(MSG_TOO_MANY_KEYS, counter));
136       }
137       ParamMap<?> paramMap = iterator.next();
138       if (assignerList.isEmpty()) {
139         for (int i = 0; i < keyProperties.length; i++) {
140           assignerList
141               .add(getAssignerForParamMap(configuration, rsmd, i + 1, paramMap, keyProperties[i], keyProperties, false)
142                   .getValue());
143         }
144       }
145       assignerList.forEach(x -> x.assign(rs, paramMap));
146       counter++;
147     }
148   }
149 
150   private void assignKeysToParamMap(Configuration configuration, ResultSet rs, ResultSetMetaData rsmd,
151       String[] keyProperties, Map<String, ?> paramMap) throws SQLException {
152     if (paramMap.isEmpty()) {
153       return;
154     }
155     Map<String, Entry<Iterator<?>, List<KeyAssigner>>> assignerMap = new HashMap<>();
156     for (int i = 0; i < keyProperties.length; i++) {
157       Entry<String, KeyAssigner> entry = getAssignerForParamMap(configuration, rsmd, i + 1, paramMap, keyProperties[i],
158           keyProperties, true);
159       Entry<Iterator<?>, List<KeyAssigner>> iteratorPair = assignerMap.computeIfAbsent(entry.getKey(),
160           k -> entry(collectionize(paramMap.get(k)).iterator(), new ArrayList<>()));
161       iteratorPair.getValue().add(entry.getValue());
162     }
163     long counter = 0;
164     while (rs.next()) {
165       for (Entry<Iterator<?>, List<KeyAssigner>> pair : assignerMap.values()) {
166         if (!pair.getKey().hasNext()) {
167           throw new ExecutorException(String.format(MSG_TOO_MANY_KEYS, counter));
168         }
169         Object param = pair.getKey().next();
170         pair.getValue().forEach(x -> x.assign(rs, param));
171       }
172       counter++;
173     }
174   }
175 
176   private Entry<String, KeyAssigner> getAssignerForParamMap(Configuration config, ResultSetMetaData rsmd,
177       int columnPosition, Map<String, ?> paramMap, String keyProperty, String[] keyProperties, boolean omitParamName) {
178     Set<String> keySet = paramMap.keySet();
179     // A caveat : if the only parameter has {@code @Param("param2")} on it,
180     // it must be referenced with param name e.g. 'param2.x'.
181     boolean singleParam = !keySet.contains(SECOND_GENERIC_PARAM_NAME);
182     int firstDot = keyProperty.indexOf('.');
183     if (firstDot == -1) {
184       if (singleParam) {
185         return getAssignerForSingleParam(config, rsmd, columnPosition, paramMap, keyProperty, omitParamName);
186       }
187       throw new ExecutorException("Could not determine which parameter to assign generated keys to. "
188           + "Note that when there are multiple parameters, 'keyProperty' must include the parameter name (e.g. 'param.id'). "
189           + "Specified key properties are " + ArrayUtil.toString(keyProperties) + " and available parameters are "
190           + keySet);
191     }
192     String paramName = keyProperty.substring(0, firstDot);
193     if (keySet.contains(paramName)) {
194       String argParamName = omitParamName ? null : paramName;
195       String argKeyProperty = keyProperty.substring(firstDot + 1);
196       return entry(paramName, new KeyAssigner(config, rsmd, columnPosition, argParamName, argKeyProperty));
197     } else if (singleParam) {
198       return getAssignerForSingleParam(config, rsmd, columnPosition, paramMap, keyProperty, omitParamName);
199     } else {
200       throw new ExecutorException("Could not find parameter '" + paramName + "'. "
201           + "Note that when there are multiple parameters, 'keyProperty' must include the parameter name (e.g. 'param.id'). "
202           + "Specified key properties are " + ArrayUtil.toString(keyProperties) + " and available parameters are "
203           + keySet);
204     }
205   }
206 
207   private Entry<String, KeyAssigner> getAssignerForSingleParam(Configuration config, ResultSetMetaData rsmd,
208       int columnPosition, Map<String, ?> paramMap, String keyProperty, boolean omitParamName) {
209     // Assume 'keyProperty' to be a property of the single param.
210     String singleParamName = nameOfSingleParam(paramMap);
211     String argParamName = omitParamName ? null : singleParamName;
212     return entry(singleParamName, new KeyAssigner(config, rsmd, columnPosition, argParamName, keyProperty));
213   }
214 
215   private static String nameOfSingleParam(Map<String, ?> paramMap) {
216     // There is virtually one parameter, so any key works.
217     return paramMap.keySet().iterator().next();
218   }
219 
220   private static Collection<?> collectionize(Object param) {
221     if (param instanceof Collection) {
222       return (Collection<?>) param;
223     } else if (param instanceof Object[]) {
224       return Arrays.asList((Object[]) param);
225     } else {
226       return Arrays.asList(param);
227     }
228   }
229 
230   private static <K, V> Entry<K, V> entry(K key, V value) {
231     // Replace this with Map.entry(key, value) in Java 9.
232     return new AbstractMap.SimpleImmutableEntry<>(key, value);
233   }
234 
235   private class KeyAssigner {
236     private final Configuration configuration;
237     private final ResultSetMetaData rsmd;
238     private final TypeHandlerRegistry typeHandlerRegistry;
239     private final int columnPosition;
240     private final String paramName;
241     private final String propertyName;
242     private TypeHandler<?> typeHandler;
243 
244     protected KeyAssigner(Configuration configuration, ResultSetMetaData rsmd, int columnPosition, String paramName,
245         String propertyName) {
246       super();
247       this.configuration = configuration;
248       this.rsmd = rsmd;
249       this.typeHandlerRegistry = configuration.getTypeHandlerRegistry();
250       this.columnPosition = columnPosition;
251       this.paramName = paramName;
252       this.propertyName = propertyName;
253     }
254 
255     protected void assign(ResultSet rs, Object param) {
256       if (paramName != null) {
257         // If paramName is set, param is ParamMap
258         param = ((ParamMap<?>) param).get(paramName);
259       }
260       MetaObject metaParam = configuration.newMetaObject(param);
261       try {
262         if (typeHandler == null) {
263           if (metaParam.hasSetter(propertyName)) {
264             Class<?> propertyType = metaParam.getSetterType(propertyName);
265             typeHandler = typeHandlerRegistry.getTypeHandler(propertyType,
266                 JdbcType.forCode(rsmd.getColumnType(columnPosition)));
267           } else {
268             throw new ExecutorException("No setter found for the keyProperty '" + propertyName + "' in '"
269                 + metaParam.getOriginalObject().getClass().getName() + "'.");
270           }
271         }
272         if (typeHandler == null) {
273           // Error?
274         } else {
275           Object value = typeHandler.getResult(rs, columnPosition);
276           metaParam.setValue(propertyName, value);
277         }
278       } catch (SQLException e) {
279         throw new ExecutorException("Error getting generated key or setting result to parameter object. Cause: " + e,
280             e);
281       }
282     }
283   }
284 }