Mybatis源码阅读（三）：结果集映射3.3 —— 主键生成策略

*************************************优雅的分割线 **********************************

分享一波:程序员赚外快-必看的巅峰干货

如果以上内容对你觉得有用,并想获取更多的赚钱方式和免费的技术教程

请关注微信公众号:HB荷包

一个能让你学习技术和赚钱方法的公众号,持续更新
*************************************优雅的分割线 **********************************
前言

在前面两篇博客中，我们介绍了对于select语句的简单映射和嵌套映射。mybatis中使用ResultHandler等一系列的类，将查询结果封装到实体类中，可以说是mybatis中最复杂的过程，而剩下的insert、update、delete语句的操作则显得较为简单，没有复杂的映射逻辑。这里需要提的是在insert语句中，关于主键自增的问题。
KeyGenerator

在我们实际的开发中，自增主键还是使用比较多的。而mybatis在处理insert语句时，并不会将自增主键给返回，而是返回插入的条数。当我们有业务需求要获取插入时产生的自增主键（或者其他类型不由程序生成的主键）时，则可以使用KeyGenerator。

KeyGenerator是个接口，定义了processBefore和processAfter两个方法，分别在insert之前和之后执行，代码如下。

/**

• insert语句不会反回自动生成的主键

• 该接口用于在插入记录时获取自增主键

• @author Clinton Begin
  */
  public interface KeyGenerator {

  /**

  • 执行insert之前执行，设置属性order=“BEFORE”
  • @param executor
  • @param ms
  • @param stmt
  • @param parameter
    */
    void processBefore(Executor executor, MappedStatement ms, Statement stmt, Object parameter);

  /**

  • 执行insert之后执行，设置属性order=“AFTER”
  • @param executor
  • @param ms
  • @param stmt
  • @param parameter
    */
    void processAfter(Executor executor, MappedStatement ms, Statement stmt, Object parameter);

}

KeyGenerator有三个实现类

其中NoKeyGenerator的方法都是空实现，本文不再讲解。
Jdbc3KeyGenerator

Jdbc3KeyGenerator用于取回数据库生成的自增id，它对应于mybatis-config.xml配置的useGeneratedKeys，以及insert节点中配置的useGeneratedKeys属性。

Jdbc3KeyGenerator的processBefore是空实现，只实现了processAfter方法，该方法会调用processBatch方法将SQL语句执行后生成的主键记录到用户传递的实参中，而不是直接返回。代码如下。

/**
 * 核心方法，将sql语句生成的主键存放到参数中
 *
 * @param ms
 * @param stmt
 * @param parameter
 */
public void processBatch(MappedStatement ms, Statement stmt, Object parameter) {
    // 获取sql节点配置的keyProperties，指定主键的字段
    final String[] keyProperties = ms.getKeyProperties();
    if (keyProperties == null || keyProperties.length == 0) {
        return;
    }
    // 获取数据库生成的主键。如果没有生成主键则结果集为空
    try (ResultSet rs = stmt.getGeneratedKeys()) {
        // 获取resultSet元信息
        final ResultSetMetaData rsmd = rs.getMetaData();
        // 获取Configuration
        final Configuration configuration = ms.getConfiguration();
        // 查到的列数小于主键的列数，说明查询有误。mybatis不进行处理
        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);
    }
}

在assignKeys方法中，处理对主键赋值的逻辑。该方法会判断参数的类型，将Map、集合、对象三种情况分别进行处理。代码如下。

private void assignKeys(Configuration configuration, ResultSet rs, ResultSetMetaData rsmd, String[] keyProperties,
                        Object parameter) throws SQLException {
    if (parameter instanceof ParamMap || parameter instanceof StrictMap) {
        // 参数是map
        assignKeysToParamMap(configuration, rs, rsmd, keyProperties, (Map<String, ?>) parameter);
    } else if (parameter instanceof ArrayList && !((ArrayList<?>) parameter).isEmpty()
            && ((ArrayList<?>) parameter).get(0) instanceof ParamMap) {
        // 参数是集合
        assignKeysToParamMapList(configuration, rs, rsmd, keyProperties, ((ArrayList<ParamMap<?>>) parameter));
    } else {
        // 参数是对象
        assignKeysToParam(configuration, rs, rsmd, keyProperties, parameter);
    }
}

三个方法逻辑类似，这里只对assignKeysToParam方法进行分析，。

该方法会在进入方法时，将参数转为集合，并遍历集合去给主键进行赋值，代码如下。

/**
 * 当参数是对象时的处理逻辑
 *
 * @param configuration
 * @param rs
 * @param rsmd
 * @param keyProperties
 * @param parameter
 * @throws SQLException
 */
private void assignKeysToParam(Configuration configuration, ResultSet rs, ResultSetMetaData rsmd,
                               String[] keyProperties, Object parameter) throws SQLException {
    // 将参数转为集合，此时这个集合最多只有一个对象
    Collection<?> params = collectionize(parameter);
    if (params.isEmpty()) {
        return;
    }
    // 存放主键属性的信息
    List<KeyAssigner> assignerList = new ArrayList<>();
    for (int i = 0; i < keyProperties.length; i++) {
        assignerList.add(new KeyAssigner(configuration, rsmd, i + 1, null, keyProperties[i]));
    }
    Iterator<?> iterator = params.iterator();
    while (rs.next()) {
        if (!iterator.hasNext()) {
            throw new ExecutorException(String.format(MSG_TOO_MANY_KEYS, params.size()));
        }
        Object param = iterator.next();
        assignerList.forEach(x -> x.assign(rs, param));
    }
}

SelectKeyGenerator

MySql、PostgreSql等数据库支持自增主键，在执行insert时可以不指定主键，插入的过程中由数据库去生成自增主键。而像Oracle、DB2等数据库产品则需要使用sequence实现自增id，因此在insert之前必须明确指定主键的值。SelectKeyGenerator就用来处理不支持自增主键的数据库。

SelectKeyGenerator主要用于生成主键，它会执行映射配置文件中定义的selectKey节点的sql，该语句会获取insert语句所需要的主键。

SelectKeyGenerator中有两个核心字段，keyStatement用于存放解析后的selectKey节点，executeBefore用于标识该节点是在insert之前还是之后执行。

/**
 * 标识selectKey是在insert之前还是之后执行
 * true为之前，false为之后
 */
private final boolean executeBefore;
/**
 * 存放selectKey节点，用于获取insert语句使用的主键
 */
private final MappedStatement keyStatement;

SelectKeyGenerator的processBefore和processAfter都是执行processGeneratorKeys方法，根据executeBefore字段的值决定是在insert之前还是之后执行。

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

processGeneratorKeys方法是处理主键的核心代码。该方法会先获取selectKey节点的keyProperties配置的属性名称，该属性对应着主键。接着，创建Executor执行器，执行selectKey节点中的sql，查询主键。Executor是mybatis的核心执行器，后面的博客会对其进行介绍。

selectKey执行完毕后，会判断查询结果是否合法。主键的查询结果肯定只会返回一条，因此查询结果条数不为1的全部视为非法查询。该方法代码如下。

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();
            // 创建用户传入的实参对应的MetaObject对象
            final MetaObject metaParam = configuration.newMetaObject(parameter);
            // 创建Executor对象，执行keyStatement记录的sql。Executor是mybatis中的执行器，后面会讲
            Executor keyExecutor = configuration.newExecutor(executor.getTransaction(), ExecutorType.SIMPLE);
            // 执行selectKey节点的sql，查询主键
            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 {
                // 主键查询出来肯定只有一条，因此values的size不等于1的情况都是错误的
                // 创建主键对象对应的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 {
                        // 如果主键对象不包含get方法，就可能是基本类型或者String，直接复制
                        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);
    }
}

该方法中有一个handleMultipleProperties方法，用于处理一张表多个主键的情况。由于复合主键耦合性高、影响性能，并且操作起来较为繁琐，因此不推荐数据库中给一张表设置多个主键，这里也就不对该方法进行介绍。
结语

KeyGenerator的代码较为简单，阅读起来也不会吃力。本周我公司已经开始了上班，因此博客也需要开始写了，mybatis的源码已经阅读了一大半，争取在7月份之前把剩下代码的博客全部编写完毕
*************************************优雅的分割线 **********************************

分享一波:程序员赚外快-必看的巅峰干货

如果以上内容对你觉得有用,并想获取更多的赚钱方式和免费的技术教程

请关注微信公众号:HB荷包

一个能让你学习技术和赚钱方法的公众号,持续更新
*************************************优雅的分割线 **********************************