fixes #1485 Improved consistency in assigning generated keys when using BatchExecutor.

harawata · harawata · commit 36d3a4a7e2e4 · 2019-02-28T01:15:38.000+09:00
Jdbc3KeyGenerator is fully rewritten and it now supports assigning generated keys to different parameters. See the test 'assigningMultipleKeysToDifferentParams'.
diff --git a/src/main/java/org/apache/ibatis/executor/keygen/Jdbc3KeyGenerator.java b/src/main/java/org/apache/ibatis/executor/keygen/Jdbc3KeyGenerator.java
@@ -19,9 +19,15 @@
 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 org.apache.ibatis.binding.MapperMethod.ParamMap;
 import org.apache.ibatis.executor.Executor;
@@ -64,122 +70,190 @@ public void processBatch(MappedStatement ms, Statement stmt, Object parameter) {
       return;
     }
     try (ResultSet rs = stmt.getGeneratedKeys()) {
+      final ResultSetMetaData rsmd = rs.getMetaData();
       final Configuration configuration = ms.getConfiguration();
-      if (rs.getMetaData().getColumnCount() >= keyProperties.length) {
-        Object soleParam = getSoleParameter(parameter);
-        if (soleParam != null) {
-          assignKeysToParam(configuration, rs, keyProperties, soleParam);
-        } else {
-          assignKeysToOneOfParams(configuration, rs, keyProperties, (Map<?, ?>) parameter);
-        }
+      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);
     }
   }
 
-  protected void assignKeysToOneOfParams(final Configuration configuration, ResultSet rs, final String[] keyProperties,
-      Map<?, ?> paramMap) throws SQLException {
-    // Assuming 'keyProperty' includes the parameter name. e.g. 'param.id'.
-    int firstDot = keyProperties[0].indexOf('.');
-    if (firstDot == -1) {
-      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 "
-              + paramMap.keySet());
-    }
-    String paramName = keyProperties[0].substring(0, firstDot);
-    Object param;
-    if (paramMap.containsKey(paramName)) {
-      param = paramMap.get(paramName);
+  @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 {
-      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 "
-          + paramMap.keySet());
+      // 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;
     }
-    // Remove param name from 'keyProperty' string. e.g. 'param.id' -> 'id'
-    String[] modifiedKeyProperties = new String[keyProperties.length];
+    List<KeyAssigner> assignerList = new ArrayList<>();
     for (int i = 0; i < keyProperties.length; i++) {
-      if (keyProperties[i].charAt(firstDot) == '.' && keyProperties[i].startsWith(paramName)) {
-        modifiedKeyProperties[i] = keyProperties[i].substring(firstDot + 1);
-      } else {
-        throw new ExecutorException("Assigning generated keys to multiple parameters is not supported. "
-            + "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 "
-            + paramMap.keySet());
-      }
+      assignerList.add(new KeyAssigner(configuration, rsmd, i + 1, null, keyProperties[i]));
+    }
+    Iterator<?> iterator = params.iterator();
+    while (rs.next()) {
+      Object param = iterator.next();
+      assignerList.forEach(x -> x.assign(rs, param));
     }
-    assignKeysToParam(configuration, rs, modifiedKeyProperties, param);
   }
 
-  private void assignKeysToParam(final Configuration configuration, ResultSet rs, final String[] keyProperties,
-      Object param)
-      throws SQLException {
-    final TypeHandlerRegistry typeHandlerRegistry = configuration.getTypeHandlerRegistry();
-    final ResultSetMetaData rsmd = rs.getMetaData();
-    // Wrap the parameter in Collection to normalize the logic.
-    Collection<?> paramAsCollection;
-    if (param instanceof Object[]) {
-      paramAsCollection = Arrays.asList((Object[]) param);
-    } else if (!(param instanceof Collection)) {
-      paramAsCollection = Arrays.asList(param);
-    } else {
-      paramAsCollection = (Collection<?>) param;
-    }
-    TypeHandler<?>[] typeHandlers = null;
-    for (Object obj : paramAsCollection) {
-      if (!rs.next()) {
-        break;
-      }
-      MetaObject metaParam = configuration.newMetaObject(obj);
-      if (typeHandlers == null) {
-        typeHandlers = getTypeHandlers(typeHandlerRegistry, metaParam, keyProperties, rsmd);
+  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<>();
+    while (rs.next()) {
+      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());
+        }
       }
-      populateKeys(rs, metaParam, keyProperties, typeHandlers);
+      assignerList.forEach(x -> x.assign(rs, paramMap));
     }
   }
 
-  private Object getSoleParameter(Object parameter) {
-    if (!(parameter instanceof ParamMap || parameter instanceof StrictMap)) {
-      return parameter;
+  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());
     }
-    Object soleParam = null;
-    for (Object paramValue : ((Map<?, ?>) parameter).values()) {
-      if (soleParam == null) {
-        soleParam = paramValue;
-      } else if (soleParam != paramValue) {
-        soleParam = null;
-        break;
+    while (rs.next()) {
+      for (Entry<Iterator<?>, List<KeyAssigner>> pair : assignerMap.values()) {
+        Object param = pair.getKey().next();
+        pair.getValue().forEach(x -> x.assign(rs, param));
       }
     }
-    return soleParam;
   }
 
-  private TypeHandler<?>[] getTypeHandlers(TypeHandlerRegistry typeHandlerRegistry, MetaObject metaParam, String[] keyProperties, ResultSetMetaData rsmd) throws SQLException {
-    TypeHandler<?>[] typeHandlers = new TypeHandler<?>[keyProperties.length];
-    for (int i = 0; i < keyProperties.length; i++) {
-      if (metaParam.hasSetter(keyProperties[i])) {
-        Class<?> keyPropertyType = metaParam.getSetterType(keyProperties[i]);
-        typeHandlers[i] = typeHandlerRegistry.getTypeHandler(keyPropertyType, JdbcType.forCode(rsmd.getColumnType(i + 1)));
-      } else {
-        throw new ExecutorException("No setter found for the keyProperty '" + keyProperties[i] + "' in '"
-            + metaParam.getOriginalObject().getClass().getName() + "'.");
+  private Entry<String, KeyAssigner> getAssignerForParamMap(Configuration config, ResultSetMetaData rsmd,
+      int columnPosition, Map<String, ?> paramMap, String keyProperty, String[] keyProperties, boolean omitParamName) {
+    boolean singleParam = paramMap.values().stream().distinct().count() == 1;
+    int firstDot = keyProperty.indexOf('.');
+    if (firstDot == -1) {
+      if (singleParam) {
+        // 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));
       }
+      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 "
+          + paramMap.keySet());
+    }
+    String paramName = keyProperty.substring(0, firstDot);
+    if (paramMap.containsKey(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) {
+      // 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));
+    } 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 "
+          + paramMap.keySet());
     }
-    return typeHandlers;
   }
 
-  private void populateKeys(ResultSet rs, MetaObject metaParam, String[] keyProperties, TypeHandler<?>[] typeHandlers) throws SQLException {
-    for (int i = 0; i < keyProperties.length; i++) {
-      String property = keyProperties[i];
-      TypeHandler<?> th = typeHandlers[i];
-      if (th != null) {
-        Object value = th.getResult(rs, i + 1);
-        metaParam.setValue(property, value);
-      }
+  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 {
+    protected final Configuration configuration;
+    protected final ResultSetMetaData rsmd;
+    protected final TypeHandlerRegistry typeHandlerRegistry;
+    protected final int columnPosition;
+    protected final String paramName;
+    protected final String propertyName;
+    protected 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);
+      }
+    }
+  }
 }
diff --git a/src/test/java/org/apache/ibatis/submitted/keygen/CountryMapper.java b/src/test/java/org/apache/ibatis/submitted/keygen/CountryMapper.java
@@ -1,5 +1,5 @@
 /**
- *    Copyright 2009-2018 the original author or authors.
+ *    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.
@@ -33,6 +33,10 @@ public interface CountryMapper {
   @Insert({ "insert into country (countryname,countrycode) values (#{country.countryname},#{country.countrycode})" })
   int insertNamedBean(@Param("country") Country country);
 
+  @Options(useGeneratedKeys = true, keyProperty = "country.id")
+  @Insert({ "insert into country (countryname,countrycode) values (#{country.countryname},#{country.countrycode})" })
+  int insertNamedBean_keyPropertyWithParamName(@Param("country") Country country);
+
   int insertList(List<Country> countries);
 
   int insertNamedList(@Param("countries") List<Country> countries);
diff --git a/src/test/java/org/apache/ibatis/submitted/keygen/Jdbc3KeyGeneratorTest.java b/src/test/java/org/apache/ibatis/submitted/keygen/Jdbc3KeyGeneratorTest.java
