mybatis源码阅读(八) ---KeyGenerator了解一下

不同的数据库产品对应的主键生成策略不一样，比如Oracle，DB2等数据库产品是通过sequence实现主键id自增的，在执行insert语句之前必须有明确的指定主键值，而mysql等数据困在执行insert语句时，可以不指定主键，由数据库自动生成自增主键。

1.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：用于处理数据库支持自增主键的情况，如MySQL的auto_increment。

NoKeyGenerator：空实现，不需要处理主键。

SelectKeyGenerator：用于处理数据库不支持自增主键的情况，比如Oracle，postgres的sequence序列。

先看看jdbc是怎么返回主键的

        Class.forName("com.mysql.jdbc.Driver");

        Connection conn = DriverManager.getConnection("jdbc:mysql://localhost:3306/mydb", "root", "123");
        conn.setAutoCommit(false);
        PreparedStatement pstm = conn.prepareStatement("insert into students(name, email) values(?, ?)",
        Statement.RETURN_GENERATED_KEYS);

        pstm.setString(1, "name1");
        pstm.setString(2, "email1");
        pstm.addBatch();
        pstm.setString(1, "name2");
        pstm.setString(2, "email2");
        pstm.addBatch();
        pstm.executeBatch();
        // 返回自增主键值
        ResultSet rs = pstm.getGeneratedKeys();
        while (rs.next()) {
        Object value = rs.getObject(1);
        System.out.println(value);
        }
        conn.commit();
        rs.close();
        pstm.close();
        conn.close();

以上代码，仅作为演示使用。Mybatis是对JDBC的封装，其Jdbc3KeyGenerator类，就是使用上面的原理，来返回数据库生成的主键值的。

2.KeyGenerator 初始化

在之前介绍的mapper文件解析中XMLStatementBuilder.parseStatementNode()方法中有如下代码片段来初始化KeyGenerator

KeyGenerator keyGenerator;
// 获取selectKey节点对应的selectKeyGenerator的id
String keyStatementId = id + SelectKeyGenerator.SELECT_KEY_SUFFIX;
keyStatementId = builderAssistant.applyCurrentNamespace(keyStatementId, true);
if (configuration.hasKeyGenerator(keyStatementId)) {//SQL节点下是否存在<selectKey>节点
  keyGenerator = configuration.getKeyGenerator(keyStatementId);
} else {
  // 根据SQL节点的useGeneratedKeys属性值、mybatis-config.xml中的全局useGeneratedKeys配置，
  // 以及是否是insert语句 觉得使用哪个KeyGenerator接口
  keyGenerator = context.getBooleanAttribute("useGeneratedKeys",
      configuration.isUseGeneratedKeys() && SqlCommandType.INSERT.equals(sqlCommandType))
      ? Jdbc3KeyGenerator.INSTANCE : NoKeyGenerator.INSTANCE;
}

3. Jdbc3KeyGenerator源码解读

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

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) {
        // 给参数object对象的属性赋主键值（批量插入，可能是多个）
        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
      }
    }
  }
}

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);
    }
  }
}

mapper配置方式。

<insert id="insert"  useGeneratedKeys="true" keyProperty="type_id">

4.SelectKeyGenerator

对于不支持自动生成的自增主键的数据库，例如Oracle，用户可以利用SelectKeyGenerator来获取生成的主键的功能，SelectKeyGenerator中的processBefore()和processAfter()方法的实现都是调用的processGeneratedKeys方法。源码如下

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

@Override
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) {
        // 获取<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);
        // 执行 SQL语句 得到主键值
        List<Object> values = keyExecutor.query(keyStatement, parameter, RowBounds.DEFAULT, Executor.NO_RESULT_HANDLER);
        // 检测values集合的长度，改集合只能维1
        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对象
          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);
  }
}

配置如下

<selectKey order="AFTER" resultType="java.lang.Integer" keyProperty="po.id">
    <choose>
        <when  test="po.DBTYPE =='ORACLE'">
            SELECT mon_deviceindex_id_seq.NEXTVAL-1 from DUAL
        </when>
        <otherwise>
            select last_value from mon_deviceindex_id_seq
       </otherwise>
    </choose>
</selectKey>

（1）selectKey 初始化过程

org.apache.ibatis.builder.xml.XMLStatementBuilder.parseSelectKeyNode()解析<selectKey>元素，构建SelectKeyGenerator的源码。

private void processSelectKeyNodes(String id, Class<?> parameterTypeClass, LanguageDriver langDriver) {
  // 获取全部的selectKey节点
  List<XNode> selectKeyNodes = context.evalNodes("selectKey");
  // 解析selectKey节点
  if (configuration.getDatabaseId() != null) {
    parseSelectKeyNodes(id, selectKeyNodes, parameterTypeClass, langDriver, configuration.getDatabaseId());
  }
  parseSelectKeyNodes(id, selectKeyNodes, parameterTypeClass, langDriver, null);
  // 删除selectKey节点
  removeSelectKeyNodes(selectKeyNodes);
}

private void parseSelectKeyNodes(String parentId, List<XNode> list, Class<?> parameterTypeClass, LanguageDriver langDriver, String skRequiredDatabaseId) {
  for (XNode nodeToHandle : list) {
    String id = parentId + SelectKeyGenerator.SELECT_KEY_SUFFIX;
    String databaseId = nodeToHandle.getStringAttribute("databaseId");
    if (databaseIdMatchesCurrent(id, databaseId, skRequiredDatabaseId)) {
      parseSelectKeyNode(id, nodeToHandle, parameterTypeClass, langDriver, databaseId);
    }
  }
}

private void parseSelectKeyNode(String id, XNode nodeToHandle, Class<?> parameterTypeClass, LanguageDriver langDriver, String databaseId) {
  String resultType = nodeToHandle.getStringAttribute("resultType");
  Class<?> resultTypeClass = resolveClass(resultType);
  StatementType statementType = StatementType.valueOf(nodeToHandle.getStringAttribute("statementType", StatementType.PREPARED.toString()));
  String keyProperty = nodeToHandle.getStringAttribute("keyProperty");
  String keyColumn = nodeToHandle.getStringAttribute("keyColumn");
  boolean executeBefore = "BEFORE".equals(nodeToHandle.getStringAttribute("order", "AFTER"));

  //defaults
  boolean useCache = false;
  boolean resultOrdered = false;
  KeyGenerator keyGenerator = NoKeyGenerator.INSTANCE;
  Integer fetchSize = null;
  Integer timeout = null;
  boolean flushCache = false;
  String parameterMap = null;
  String resultMap = null;
  ResultSetType resultSetTypeEnum = null;

  // 生成SqlSource
  SqlSource sqlSource = langDriver.createSqlSource(configuration, nodeToHandle, parameterTypeClass);
  // selectKey节点中只能配置select语句
  SqlCommandType sqlCommandType = SqlCommandType.SELECT;

  // 创建MappedStatement对象，并添加到configuration的mappedStatements集合中保存
  builderAssistant.addMappedStatement(id, sqlSource, statementType, sqlCommandType,
      fetchSize, timeout, parameterMap, parameterTypeClass, resultMap, resultTypeClass,
      resultSetTypeEnum, flushCache, useCache, resultOrdered,
      keyGenerator, keyProperty, keyColumn, databaseId, langDriver, null);

  id = builderAssistant.applyCurrentNamespace(id, false);

  MappedStatement keyStatement = configuration.getMappedStatement(id, false);
  // 创建对应的KeyGenerator(主键自增策略)，添加到configuration中
  configuration.addKeyGenerator(id, new SelectKeyGenerator(keyStatement, executeBefore));
}

因此，只有SelectKeyGenerator会保存到Configuration对象的Map<String, KeyGenerator> keyGenerators属性当中。<selectKey>元素，会被Mybatis解析为一个MappedStatement对象，并作为构造参数传递至SelectKeyGenerator内保存起来。

至此，每一个MappedStatement对象，都恰当的绑定了一个KeyGenerator对象，就可以开始工作了。

5. KeyGenerator的使用过程

keyGenerator.processBefore()方法调用时机：

org.apache.ibatis.executor.statement.BaseStatementHandler.BaseStatementHandler()构造方法源码。

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();
}

即，创建StatementHandler对象时，就会执行keyGenerator.processBefore()方法。keyGenerator.processAfter()方法，自然就是Statement执行后执行了。

org.apache.ibatis.executor.statement.SimpleStatementHandler.update(Statement)方法源码。其他的StatementHandler都是类似的。

@Override
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 {
    statement.execute(sql);
    rows = statement.getUpdateCount();
  }
  return rows;
}

注意：由于selectKey本身返回单个序列主键值，也就无法支持批量insert操作并返回主键id列表了。如果要执行批量insert，请选择使用for循环执行多次插入操作。