mybatis源码分析：KeyGenerator

概述

默认情况下，insert语句并不会返回自动生成的主键，而是返回插入记录的条数。如果业务逻辑需要获取插入记录时产生的自增主键，则可以使用Mybatis提供的KeyGenerator接口。
不同的数据库产品对应的主键生成策略不一样，例如，Oracle、DB2等数据库铲平是通过sequence实现自增id的，在执行insert语句之前必须明确指定主键的值；而Mysql、Postgresql等数据库在执行insert语句时，可以不指定主键，在插入过程中由数据库自动生成自增主键。KeyGenerator接口针对这些不同的数据库产品提供了对应的处理方法，KeyGenerator接口的定义如下所示：

public interface KeyGenerator {
  // 在执行insert之前执行，设置属性order="BEFORE"
  void processBefore(Executor executor, MappedStatement ms, Statement stmt, Object parameter);
  // 在执行insert之后执行，设置属性order="AFTER"
  void processAfter(Executor executor, MappedStatement ms, Statement stmt, Object parameter);

}

Mybatis提供了三个KeyGenerator接口的实现：

Jdbc3KeyGenerator

Jdbc3KeyGenerator用于取回数据库生成的自增id，它对应于mybatis-config.xml配置文件中的useGeneratedKeys全局配置，以及映射配置文件中SQL节点的useGeneratedKeys属性配合keyProperty属性一起使用。
Jdbc3KeyGenerator.processBefore()方法是空实现，只实现了processAfter()方法，该方法会调用Jdbc3KeyGenerator.processBatch()方法将SQL语句执行后生成的主键记录到用户传递的实参中。一般情况下，对于单行插入操作，传入实参是一个JavaBean对象或者是Map对象，则该对象对应一次插入操作的内容；对于多行插入，传入的实参可以是对象或Map对象的数组或集合，集合中每一个元素都对应一次插入操作。
processAfter()方法如下：

@Override
public void processAfter(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {
    processBatch(ms, stmt, getParameters(parameter));
}

getParameters()方法如下：

// 将用户传入的实参转化成集合
private Collection<Object> getParameters(Object parameter) {
    Collection<Object> parameters = null;
    if (parameter instanceof Collection) {
      parameters = (Collection) parameter;
    } else if (parameter instanceof Map) {
      Map parameterMap = (Map) parameter;
      if (parameterMap.containsKey("collection")) {
        parameters = (Collection) parameterMap.get("collection");
      } else if (parameterMap.containsKey("list")) {
        parameters = (List) parameterMap.get("list");
      } else if (parameterMap.containsKey("array")) {
        parameters = Arrays.asList((Object[]) parameterMap.get("array"));
      }
    }
    if (parameters == null) {
      parameters = new ArrayList<Object>();
      parameters.add(parameter);
    }
    return parameters;
}

processBatch()方法如下：

public void processBatch(MappedStatement ms, Statement stmt, Collection<Object> parameters) {
    ResultSet rs = null;
    try {
      // 从结果集中获取数据库自动生成的主键
      rs = stmt.getGeneratedKeys();
      final Configuration configuration = ms.getConfiguration();
      final TypeHandlerRegistry typeHandlerRegistry = configuration.getTypeHandlerRegistry();
      final String[] keyProperties = ms.getKeyProperties();
      final ResultSetMetaData rsmd = rs.getMetaData();
      TypeHandler<?>[] typeHandlers = null;
      if (keyProperties != null && rsmd.getColumnCount() >= keyProperties.length) {
        for (Object parameter : parameters) {
          // there should be one row for each statement (also one for each parameter)
          if (!rs.next()) {
            break;
          }
          final MetaObject metaParam = configuration.newMetaObject(parameter);
          if (typeHandlers == null) {
            typeHandlers = getTypeHandlers(typeHandlerRegistry, metaParam, keyProperties, rsmd);
          }
          // 将主键值设置到用户提供的参数中返回
          populateKeys(rs, metaParam, keyProperties, typeHandlers);
        }
      }
    } catch (Exception e) {
      throw new ExecutorException("Error getting generated key or setting result to parameter object. Cause: " + e, e);
    } finally {
      if (rs != null) {
        try {
          rs.close();
        } catch (Exception e) {
          // ignore
        }
      }
    }
}

SelectKeyGenerator

对于不支持自动生成自增主键的数据库，用户可以利用Mybatis提供的SelectKeyGenerator来生成主键。它会执行映射配置文件中定义的selectKey节点的SQL语句，该语句会获取insert语句所需要的主键。

private void processGeneratedKeys(Executor executor, MappedStatement ms, Object parameter) {
    try {
      if (parameter != null && keyStatement != null && keyStatement.getKeyProperties() != null) {
        // 获取selectKey的keyProperties的属性
        String[] keyProperties = keyStatement.getKeyProperties();
        final Configuration configuration = ms.getConfiguration();
        final MetaObject metaParam = configuration.newMetaObject(parameter);
        if (keyProperties != null) {
          // Do not close keyExecutor.
          // The transaction will be closed by parent executor.
          Executor keyExecutor = configuration.newExecutor(executor.getTransaction(), ExecutorType.SIMPLE);
          List<Object> values = keyExecutor.query(keyStatement, parameter, RowBounds.DEFAULT, Executor.NO_RESULT_HANDLER);
          if (values.size() == 0) {
            throw new ExecutorException("SelectKey returned no data.");            
          } else if (values.size() > 1) {
            throw new ExecutorException("SelectKey returned more than one value.");
          } else {
            MetaObject metaResult = configuration.newMetaObject(values.get(0));
            if (keyProperties.length == 1) {
              if (metaResult.hasGetter(keyProperties[0])) {
                setValue(metaParam, keyProperties[0], metaResult.getValue(keyProperties[0]));
              } else {
                // no getter for the property - maybe just a single value object
                // so try that
                setValue(metaParam, keyProperties[0], values.get(0));
              }
            } else {
              handleMultipleProperties(keyProperties, metaParam, metaResult);
            }
          }
        }
      }
    } catch (ExecutorException e) {
      throw e;
    } catch (Exception e) {
      throw new ExecutorException("Error selecting key or setting result to parameter object. Cause: " + e, e);
    }
}