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