mybatis 源码分析（七）KeyGenerator 详解

一、KeyGenerator 概述

在平时开发的时候经常会有这样的需求，插入数据返回主键，或者插入数据之前需要获取主键，这样的需求在 mybatis 中也是支持的，其中主要的逻辑部分就在 KeyGenerator 中，下面是他的类图：

其中：

• NoKeyGenerator：默认空实现，不需要对主键单独处理；
• Jdbc3KeyGenerator：主要用于数据库的自增主键，比如 MySQL、PostgreSQL；
• SelectKeyGenerator：主要用于数据库不支持自增主键的情况，比如 Oracle、DB2；

接口方法如下：

public interface KeyGenerator {
  void processBefore(Executor executor, MappedStatement ms, Statement stmt, Object parameter);
  void processAfter(Executor executor, MappedStatement ms, Statement stmt, Object parameter);
}

如代码所见 KeyGenerator 非常的简单，主要是通过两个拦截方法实现的：

• Jdbc3KeyGenerator：主要基于 java.sql.Statement.getGeneratedKeys 的主键返回接口实现的，所以他不需要 processBefore 方法，只需要在获取到结果后使用 processAfter 拦截，然后用反射将主键设置到参数中即可；
• SelectKeyGenerator：主要是通过 XML 配置或者注解设置 selectKey ，然后单独发出查询语句，在返回拦截方法中使用反射设置主键，其中两个拦截方法只能使用其一，在 selectKey.order 属性中设置 AFTER|BEFORE 来确定；

拦截时机：

processBefore 是在生成 StatementHandler 的时候；

protected BaseStatementHandler(Executor executor, MappedStatement mappedStatement, Object parameterObject, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) {
  ...
  if (boundSql == null) { // issue #435, get the key before calculating the statement
    generateKeys(parameterObject);
    boundSql = mappedStatement.getBoundSql(parameterObject);
  }
  ...
}

protected void generateKeys(Object parameter) {
  KeyGenerator keyGenerator = mappedStatement.getKeyGenerator();
  ErrorContext.instance().store();
  keyGenerator.processBefore(executor, mappedStatement, null, parameter);
  ErrorContext.instance().recall();
}

processAfter 则是在完成插入返回结果之前，但是 PreparedStatementHandler、SimpleStatementHandler、CallableStatementHandler 的代码稍微有一点不同，但是位置是不变的，这里以 PreparedStatementHandler 举例：

@Override
public int update(Statement statement) throws SQLException {
  PreparedStatement ps = (PreparedStatement) statement;
  ps.execute();
  int rows = ps.getUpdateCount();
  Object parameterObject = boundSql.getParameterObject();
  KeyGenerator keyGenerator = mappedStatement.getKeyGenerator();
  keyGenerator.processAfter(executor, mappedStatement, ps, parameterObject);
  return rows;
}

二、Jdbc3KeyGenerator

上面也将了 Jdbc3KeyGenerator 是主要基于 java.sql.Statement.getGeneratedKeys 的主键返回接口实现的，但是 Statement 和 PreparedStatement 稍有不同，所以导致了 PreparedStatementHandler、SimpleStatementHandler 的 update 方法稍有不同：

// java.sql.Connection
PreparedStatement prepareStatement(String sql, int autoGeneratedKeys) throws SQLException;
PreparedStatement prepareStatement(String sql, String columnNames[]) throws SQLException;
PreparedStatement prepareStatement(String sql, int columnIndexes[]) throws SQLException;

// java.sql.Statement
boolean execute(String sql, int autoGeneratedKeys) throws SQLException;
boolean execute(String sql, int columnIndexes[]) throws SQLException;
boolean execute(String sql, String columnNames[]) throws SQLException;
// 其中 autoGenerateKeys - Statement.RETURN_GENERATED_KEYS、Statement.NO_GENERATED_KEYS

可以看到 PreparedStatement 是在实例化的时候就指定了，而 Statement 是在执行 sql 的时候才指定但实质是一样的，这里就以 PreparedStatement 举例：

public void testJDBC3() {
  try {
    String url = "jdbc:mysql://localhost:3306/mybatis?serverTimezone=GMT";
    String sql = "INSERT INTO user(username,password,address) VALUES (?,?,?)";
    Class.forName("com.mysql.jdbc.Driver");
    Connection conn = DriverManager.getConnection(url, "root", "root");
    String[] columnNames = {"ids", "name"};
    PreparedStatement stmt = conn.prepareStatement(sql, columnNames);
    stmt.setString(1, "test");
    stmt.setString(2, "123456");
    stmt.setString(3, "test");
    stmt.executeUpdate();
    ResultSet rs = stmt.getGeneratedKeys();
    int id = 0;
    if (rs.next()) {
      id = rs.getInt(1);
      System.out.println("----------" + id);
    }
  } catch (Exception e) {
    e.printStackTrace();
  }
}

这里的 User 表以 id 为主键，但是代码中我传的 columnNames 都不符合，而结果仍然可以正确的返回主键，主要是因为在 mybatis 的驱动中只要 columnNames.length > 1就可以了，所以在具体使用的时候还要注意不同数据库驱动实现不同所带来的影响；

上面将了 Statement 和 PreparedStatement 指定返回主键的位置不同，在下面就能很清楚的看到：

// org.apache.ibatis.executor.statement.SimpleStatementHandler
public int update(Statement statement) throws SQLException {
  String sql = boundSql.getSql();
  Object parameterObject = boundSql.getParameterObject();
  KeyGenerator keyGenerator = mappedStatement.getKeyGenerator();
  int rows;
  if (keyGenerator instanceof Jdbc3KeyGenerator) {
    statement.execute(sql, Statement.RETURN_GENERATED_KEYS);
    rows = statement.getUpdateCount();
    keyGenerator.processAfter(executor, mappedStatement, statement, parameterObject);
  } else if (keyGenerator instanceof SelectKeyGenerator) {
    statement.execute(sql);
    rows = statement.getUpdateCount();
    keyGenerator.processAfter(executor, mappedStatement, statement, parameterObject);
  } else {
    //如果没有keyGenerator,直接调用Statement.execute和Statement.getUpdateCount
    statement.execute(sql);
    rows = statement.getUpdateCount();
  }
  return rows;
}

// org.apache.ibatis.executor.statement.PreparedStatementHandler
protected Statement instantiateStatement(Connection connection) throws SQLException {
  String sql = boundSql.getSql();
  if (mappedStatement.getKeyGenerator() instanceof Jdbc3KeyGenerator) {
    String[] keyColumnNames = mappedStatement.getKeyColumns();
    if (keyColumnNames == null) {
      return connection.prepareStatement(sql, PreparedStatement.RETURN_GENERATED_KEYS);
    } else {
      return connection.prepareStatement(sql, keyColumnNames);
    }
  } else if (mappedStatement.getResultSetType() != null) {
    return connection.prepareStatement(sql, mappedStatement.getResultSetType().getValue(), ResultSet.CONCUR_READ_ONLY);
  } else {
    return connection.prepareStatement(sql);
  }
}

在完成初始化后，下面来看 Jdbc3KeyGenerator 中最主要的拦截方法：

public void processAfter(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {
  List<Object> parameters = new ArrayList<Object>();
  parameters.add(parameter);
  processBatch(ms, stmt, parameters);
}

public void processBatch(MappedStatement ms, Statement stmt, List<Object> parameters) {
  ResultSet rs = null;
  try {
    //核心是使用JDBC3的Statement.getGeneratedKeys
    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);
        }
        //填充键值
        populateKeys(rs, metaParam, keyProperties, typeHandlers);
      }
    }
  } catch (Exception e) {
    ...
  }
}

这里就很清楚了，直接获取返回的主键，然后一次使用反射设置到参数中；

三、SelectKeyGenerator

上面也讲了 SelectKeyGenerator 主要是配置 selectKey 使用的，默认 使用 processBefore，但是可以配置 order 属性（AFTER|BEFORE）；

<insert id="insertUser2" parameterType="u" useGeneratedKeys="true" keyProperty="id">
  <selectKey keyProperty="id" resultType="long" order="BEFORE">
    SELECT if(max(id) is null,1,max(id)+2) as newId FROM user2
  </selectKey>
  INSERT INTO user2(id,username,password,address) VALUES (#{id},#{userName},#{password},#{address})
</insert>

这里直接看源码：

public void processBefore(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {
  if (executeBefore) processGeneratedKeys(executor, ms, parameter);
}

public void processAfter(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {
  if (!executeBefore) processGeneratedKeys(executor, ms, parameter);
}

private void processGeneratedKeys(Executor executor, MappedStatement ms, Object parameter) {
  try {
    if (parameter != null && keyStatement != null && keyStatement.getKeyProperties() != null) {
      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) {
    ...
  }
}

这里代码也很简单，就是用一个新的 Executor 再发一条 SQL，然后反射设置参数即可；