Java tutorial
/** * Copyright 2009-2019 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.apache.ibatis.executor.keygen; import java.sql.ResultSet; import java.sql.ResultSetMetaData; import java.sql.SQLException; import java.sql.Statement; import java.util.AbstractMap; import java.util.ArrayList; import java.util.Arrays; import java.util.Collection; import java.util.HashMap; import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.Map.Entry; import java.util.Set; import org.apache.ibatis.binding.MapperMethod.ParamMap; import org.apache.ibatis.executor.Executor; import org.apache.ibatis.executor.ExecutorException; import org.apache.ibatis.mapping.MappedStatement; import org.apache.ibatis.reflection.ArrayUtil; import org.apache.ibatis.reflection.MetaObject; import org.apache.ibatis.reflection.ParamNameResolver; import org.apache.ibatis.session.Configuration; import org.apache.ibatis.session.defaults.DefaultSqlSession.StrictMap; import org.apache.ibatis.type.JdbcType; import org.apache.ibatis.type.TypeHandler; import org.apache.ibatis.type.TypeHandlerRegistry; /** * @author Clinton Begin * @author Kazuki Shimizu */ public class Jdbc3KeyGenerator implements KeyGenerator { private static final String SECOND_GENERIC_PARAM_NAME = ParamNameResolver.GENERIC_NAME_PREFIX + "2"; /** * A shared instance. * * @since 3.4.3 */ public static final Jdbc3KeyGenerator INSTANCE = new Jdbc3KeyGenerator(); private static final String MSG_TOO_MANY_KEYS = "Too many keys are generated. There are only %d target objects. " + "You either specified a wrong 'keyProperty' or encountered a driver bug like #1523."; @Override public void processBefore(Executor executor, MappedStatement ms, Statement stmt, Object parameter) { // do nothing } @Override public void processAfter(Executor executor, MappedStatement ms, Statement stmt, Object parameter) { processBatch(ms, stmt, parameter); } public void processBatch(MappedStatement ms, Statement stmt, Object parameter) { final String[] keyProperties = ms.getKeyProperties(); if (keyProperties == null || keyProperties.length == 0) { return; } try (ResultSet rs = stmt.getGeneratedKeys()) { final ResultSetMetaData rsmd = rs.getMetaData(); final Configuration configuration = ms.getConfiguration(); if (rsmd.getColumnCount() < keyProperties.length) { // Error? } else { assignKeys(configuration, rs, rsmd, keyProperties, parameter); } } catch (Exception e) { throw new ExecutorException( "Error getting generated key or setting result to parameter object. Cause: " + e, e); } } @SuppressWarnings("unchecked") private void assignKeys(Configuration configuration, ResultSet rs, ResultSetMetaData rsmd, String[] keyProperties, Object parameter) throws SQLException { if (parameter instanceof ParamMap || parameter instanceof StrictMap) { // Multi-param or single param with @Param assignKeysToParamMap(configuration, rs, rsmd, keyProperties, (Map<String, ?>) parameter); } else if (parameter instanceof ArrayList && !((ArrayList<?>) parameter).isEmpty() && ((ArrayList<?>) parameter).get(0) instanceof ParamMap) { // Multi-param or single param with @Param in batch operation assignKeysToParamMapList(configuration, rs, rsmd, keyProperties, ((ArrayList<ParamMap<?>>) parameter)); } else { // Single param without @Param assignKeysToParam(configuration, rs, rsmd, keyProperties, parameter); } } private void assignKeysToParam(Configuration configuration, ResultSet rs, ResultSetMetaData rsmd, String[] keyProperties, Object parameter) throws SQLException { Collection<?> params = collectionize(parameter); if (params.isEmpty()) { return; } List<KeyAssigner> assignerList = new ArrayList<>(); for (int i = 0; i < keyProperties.length; i++) { assignerList.add(new KeyAssigner(configuration, rsmd, i + 1, null, keyProperties[i])); } Iterator<?> iterator = params.iterator(); while (rs.next()) { if (!iterator.hasNext()) { throw new ExecutorException(String.format(MSG_TOO_MANY_KEYS, params.size())); } Object param = iterator.next(); assignerList.forEach(x -> x.assign(rs, param)); } } private void assignKeysToParamMapList(Configuration configuration, ResultSet rs, ResultSetMetaData rsmd, String[] keyProperties, ArrayList<ParamMap<?>> paramMapList) throws SQLException { Iterator<ParamMap<?>> iterator = paramMapList.iterator(); List<KeyAssigner> assignerList = new ArrayList<>(); long counter = 0; while (rs.next()) { if (!iterator.hasNext()) { throw new ExecutorException(String.format(MSG_TOO_MANY_KEYS, counter)); } ParamMap<?> paramMap = iterator.next(); if (assignerList.isEmpty()) { for (int i = 0; i < keyProperties.length; i++) { assignerList.add(getAssignerForParamMap(configuration, rsmd, i + 1, paramMap, keyProperties[i], keyProperties, false).getValue()); } } assignerList.forEach(x -> x.assign(rs, paramMap)); counter++; } } private void assignKeysToParamMap(Configuration configuration, ResultSet rs, ResultSetMetaData rsmd, String[] keyProperties, Map<String, ?> paramMap) throws SQLException { if (paramMap.isEmpty()) { return; } Map<String, Entry<Iterator<?>, List<KeyAssigner>>> assignerMap = new HashMap<>(); for (int i = 0; i < keyProperties.length; i++) { Entry<String, KeyAssigner> entry = getAssignerForParamMap(configuration, rsmd, i + 1, paramMap, keyProperties[i], keyProperties, true); Entry<Iterator<?>, List<KeyAssigner>> iteratorPair = assignerMap.computeIfAbsent(entry.getKey(), k -> entry(collectionize(paramMap.get(k)).iterator(), new ArrayList<>())); iteratorPair.getValue().add(entry.getValue()); } long counter = 0; while (rs.next()) { for (Entry<Iterator<?>, List<KeyAssigner>> pair : assignerMap.values()) { if (!pair.getKey().hasNext()) { throw new ExecutorException(String.format(MSG_TOO_MANY_KEYS, counter)); } Object param = pair.getKey().next(); pair.getValue().forEach(x -> x.assign(rs, param)); } counter++; } } private Entry<String, KeyAssigner> getAssignerForParamMap(Configuration config, ResultSetMetaData rsmd, int columnPosition, Map<String, ?> paramMap, String keyProperty, String[] keyProperties, boolean omitParamName) { Set<String> keySet = paramMap.keySet(); // A caveat : if the only parameter has {@code @Param("param2")} on it, // it must be referenced with param name e.g. 'param2.x'. boolean singleParam = !keySet.contains(SECOND_GENERIC_PARAM_NAME); int firstDot = keyProperty.indexOf('.'); if (firstDot == -1) { if (singleParam) { return getAssignerForSingleParam(config, rsmd, columnPosition, paramMap, keyProperty, omitParamName); } throw new ExecutorException("Could not determine which parameter to assign generated keys to. " + "Note that when there are multiple parameters, 'keyProperty' must include the parameter name (e.g. 'param.id'). " + "Specified key properties are " + ArrayUtil.toString(keyProperties) + " and available parameters are " + keySet); } String paramName = keyProperty.substring(0, firstDot); if (keySet.contains(paramName)) { String argParamName = omitParamName ? null : paramName; String argKeyProperty = keyProperty.substring(firstDot + 1); return entry(paramName, new KeyAssigner(config, rsmd, columnPosition, argParamName, argKeyProperty)); } else if (singleParam) { return getAssignerForSingleParam(config, rsmd, columnPosition, paramMap, keyProperty, omitParamName); } else { throw new ExecutorException("Could not find parameter '" + paramName + "'. " + "Note that when there are multiple parameters, 'keyProperty' must include the parameter name (e.g. 'param.id'). " + "Specified key properties are " + ArrayUtil.toString(keyProperties) + " and available parameters are " + keySet); } } private Entry<String, KeyAssigner> getAssignerForSingleParam(Configuration config, ResultSetMetaData rsmd, int columnPosition, Map<String, ?> paramMap, String keyProperty, boolean omitParamName) { // Assume 'keyProperty' to be a property of the single param. String singleParamName = nameOfSingleParam(paramMap); String argParamName = omitParamName ? null : singleParamName; return entry(singleParamName, new KeyAssigner(config, rsmd, columnPosition, argParamName, keyProperty)); } private static String nameOfSingleParam(Map<String, ?> paramMap) { // There is virtually one parameter, so any key works. return paramMap.keySet().iterator().next(); } private static Collection<?> collectionize(Object param) { if (param instanceof Collection) { return (Collection<?>) param; } else if (param instanceof Object[]) { return Arrays.asList((Object[]) param); } else { return Arrays.asList(param); } } private static <K, V> Entry<K, V> entry(K key, V value) { // Replace this with Map.entry(key, value) in Java 9. return new AbstractMap.SimpleImmutableEntry<>(key, value); } private class KeyAssigner { private final Configuration configuration; private final ResultSetMetaData rsmd; private final TypeHandlerRegistry typeHandlerRegistry; private final int columnPosition; private final String paramName; private final String propertyName; private TypeHandler<?> typeHandler; protected KeyAssigner(Configuration configuration, ResultSetMetaData rsmd, int columnPosition, String paramName, String propertyName) { super(); this.configuration = configuration; this.rsmd = rsmd; this.typeHandlerRegistry = configuration.getTypeHandlerRegistry(); this.columnPosition = columnPosition; this.paramName = paramName; this.propertyName = propertyName; } protected void assign(ResultSet rs, Object param) { if (paramName != null) { // If paramName is set, param is ParamMap param = ((ParamMap<?>) param).get(paramName); } MetaObject metaParam = configuration.newMetaObject(param); try { if (typeHandler == null) { if (metaParam.hasSetter(propertyName)) { Class<?> propertyType = metaParam.getSetterType(propertyName); typeHandler = typeHandlerRegistry.getTypeHandler(propertyType, JdbcType.forCode(rsmd.getColumnType(columnPosition))); } else { throw new ExecutorException("No setter found for the keyProperty '" + propertyName + "' in '" + metaParam.getOriginalObject().getClass().getName() + "'."); } } if (typeHandler == null) { // Error? } else { Object value = typeHandler.getResult(rs, columnPosition); metaParam.setValue(propertyName, value); } } catch (SQLException e) { throw new ExecutorException( "Error getting generated key or setting result to parameter object. Cause: " + e, e); } } } }