手敲MyBatis(十三章)-返回Insert操作自增索引值

1.目的

这一章的目的主要是插入语句以后返回插入记录的id，因为插入语句可分为要返回记录id，不要返回记录id的以及不同数据源类型执行的时机也不同（如：oracle不支持主键，需要先插入序列再增加，Mysql支持主键增加一条记录就会有索引）。

如下图，insert里包含selectKey，由selectKey去执行查询此次新增的id记录，我们看到selectKey标签上的属性有keyProperty、order、resultType。

keyProperty这个是把返回索引的索引放入id里
order为after则是在insert后执行selectKey
resultType则是查询返回的类型，示例里为Long

我们最主要的目的就是解析这样的selectKey，然后存入MappedStatement里到时执行时使用，执行Update时则执行selectKey的方法，最后在activity实体id赋值执行完的记录id。

2.设计说明

标黄色的是修改的，其余都是新增的方法。

1.构建时（builder）

在构建（builder）时，我们需要解析下selectKey的标签，而selectKey是在语句里，所以需要更改XMLStatementBuilder类添加解析selectKey的标签。最后把解析好的标签放入到addMappedStatement()，这里的动作是两次addMappedStatement()，为什么是两次呢？ps:可以看下面两个图。

第一次，SqlCommandType是"SELECT"，代表selectKey本身的语句，需要存储一次MappedStatement
第二次，SqlCommandType是"INSERT"，代表是selectKey父级的INSERT标签，它下面要包含这个selectKey的MappedStatement信息，然后如果有selectKey则创建SelectKeyGenerator对象，需要再一次存储MappedStatement。

ps此处看不懂可以调试全局看下，多看几遍就懂了为什么这样设计了。

2.selectKey执行时（executor keygen）

2.1 定义接口：

针对上述情况，需要定义接口KeyGenerator，定义方法为processBefore()（针对Oracle不支持主键的）和processAfter()（支持主键的，如MySql）

2.2 实现接口：

NoKenGenerator和SelectKeyGenerator，不要求返回记录的就执行NoKenGenerator（实现方法什么都不操作）。SelectKeyGenerator的实现则是执行查询索引操作，并将结果添加到insert的操作的实体里，这样就得到了索引。

3.执行时（executor）

我们存储到MappedStatement以后就要流转到执行时，因为SelectKey是查询，所以在Executor时

添加了重载的query()方法。

BaseExecutor实现新的query()方法，拿到了sql语句后，调用原有的doQuery，由SimpleExecutor类实现，此时就会调用BaseStatementHandler构造方法，这里新添加了generateKeys()方法，主要是检查下是否需要在insert前执行selectKey（如：Mysql是insert后，Oracle是insert前）。

执行了新增方法时，则需要添加调用KeyGenerator的processAfter()方法，这个方法依然要检查是否在insert后执行selectKey。

4.结果集的更改(resultset)

由于之前的结果集是只处理bean对象，一般返回selectKey需要的基本类型，如ID是Long，所以这里需要添加createPrimitiveResultObject()方法获取类型处理器。

在createResultObject()此方法里则需要判断是否是基本类型，如果基本类型则调用createPrimitiveResultObject()。其余类型则还按之前代码走。

5.connection，连接的更改

由于要两个语句用一个连接如果两个连接则无法返回id索引，所以需要在获取连接处判断下是否有连接，如果有就不创建，没有就创建。

3.代码

3.1 selectKey执行操作

包：package cn.bugstack.mybatis.executor.keygen;

添加接口KeyGenerator，主要是用来执行SelectKey查询操作的，此接口定义了两个方法，

processBefore：是在不支持主键需要获取序列时调用，执行insert语句前调用

processAfter：是在支持主键在执行insert语句后嗲用

java 复制代码

public interface KeyGenerator {

    /**
     * 针对Sequence主键而言，在执行insert sql前必须指定一个主键值给要插入的记录，
     * 如Oracle、DB2，KeyGenerator提供了processBefore()方法。
     */
    void processBefore(Executor executor, MappedStatement ms, Statement stmt, Object parameter);

    /**
     * 针对自增主键的表，在插入时不需要主键，而是在插入过程自动获取一个自增的主键，
     * 比如MySQL、PostgreSQL，KeyGenerator提供了processAfter()方法
     */
    void processAfter(Executor executor, MappedStatement ms, Statement stmt, Object parameter);
}

实现类SelectKeyGenera，有selectKey时实现执行selectKey的逻辑。

java 复制代码

// step13新增
public class SelectKeyGenerator implements KeyGenerator {


    public static final String SELECT_KEY_SUFFIX = "!selectKey";
    private boolean executeBefore;
    private MappedStatement keyStatement;

    public SelectKeyGenerator(MappedStatement keyStatement, boolean executeBefore) {
        this.executeBefore = executeBefore;
        this.keyStatement = keyStatement;
    }

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

    /**
     * 执行selectKey的SQL语句，执行完毕后把返回的id结果放入对应实体中。
     */
    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) {
                    Executor keyExecutor = configuration.newExecutor(executor.getTransaction());
                    List<Object> values = keyExecutor.query(keyStatement, parameter, RowBounds.DEFAULT, Executor.NO_RESULT_HANDLER);
                    if (values.size() == 0) {
                        throw new RuntimeException("SelectKey returned no data.");
                    } else if (values.size() > 1) {
                        throw new RuntimeException("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 {
                                setValue(metaParam, keyProperties[0], values.get(0));
                            }
                        } else {
                            handleMultipleProperties(keyProperties, metaParam, metaResult);
                        }
                    }
                }
            }
        } catch (Exception e) {
            throw new RuntimeException("Error selecting key or setting result to parameter object. Cause: " + e);
        }
    }

    private void handleMultipleProperties(String[] keyProperties,
                                          MetaObject metaParam, MetaObject metaResult) {
        String[] keyColumns = keyStatement.getKeyColumns();

        if (keyColumns == null || keyColumns.length == 0) {
            for (String keyProperty : keyProperties) {
                setValue(metaParam, keyProperty, metaResult.getValue(keyProperty));
            }
        } else {
            if (keyColumns.length != keyProperties.length) {
                throw new RuntimeException("If SelectKey has key columns, the number must match the number of key properties.");
            }
            for (int i = 0; i < keyProperties.length; i++) {
                setValue(metaParam, keyProperties[i], metaResult.getValue(keyColumns[i]));
            }
        }
    }

    private void setValue(MetaObject metaParam, String property, Object value) {
        if (metaParam.hasSetter(property)) {
            metaParam.setValue(property, value);
        } else {
            throw new RuntimeException("No setter found for the keyProperty '" + property + "' in " + metaParam.getOriginalObject().getClass().getName() + ".");
        }
    }
}

实现类NoKeyGenerator，在没有selectKey时的逻辑也就是不写具体的逻辑。

java 复制代码

public class NoKeyGenerator implements KeyGenerator {
    @Override
    public void processBefore(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {
        // Do Nothing
    }

    @Override
    public void processAfter(Executor executor, MappedStatement ms, Statement stmt, Object parameter) {
        // Do Nothing
    }
}

3.2 builder操作部分的修改

XMLStatementBuilder修改parseStatementNode()，添加了解析selectKey的节点方法processSelectKeyNodes()；

如果有多个就遍历每个节点信息存储到MappedStatement下

java 复制代码

public class XMLStatementBuilder extends BaseBuilder {
   // 省略其他
   
   public void parseStatementNode() {
        // 省略其他
        
        //  解析<selectKey> step-13 新增
        processSelectKeyNodes(id, parameterTypeClass, langDriver);
    
   }

     /**
     * 解析selectKey标签
     */
    private void processSelectKeyNodes(String id, Class<?> parameterTypeClass, LanguageDriver langDriver) {
        // 得到selectKey标签
        List<Element> selectKeyNodes = element.elements("selectKey");
        parseSelectKeyNodes(id, selectKeyNodes, parameterTypeClass, langDriver);
    }

    /**
     * for循环遍历selectKey标签
     * 取出selectKey的父id拼接select的Key标识
     */
    private void parseSelectKeyNodes(String parentId, List<Element> list, Class<?> parameterTypeClass, LanguageDriver langDriver) {
        for (Element nodeToHandle : list) {
            String id = parentId + SelectKeyGenerator.SELECT_KEY_SUFFIX;
            parseSelectKeyNode(id, nodeToHandle, parameterTypeClass, langDriver);
        }
    }


    /**
     * 开始解析每一个selectKey标签的内容,并存储MappedStatement里
     * <selectKey keyProperty="id" order="AFTER" resultType="long">
     * SELECT LAST_INSERT_ID()
     * </selectKey>
     */
    private void parseSelectKeyNode(String id, Element nodeToHandle, Class<?> parameterTypeClass, LanguageDriver langDriver) {
        String resultType = nodeToHandle.attributeValue("resultType");
        // 得到resultType类型
        Class<?> resultTypeClass = resolveClass(resultType);
        // 执行前还是执行后
        boolean executeBefore = "BEFORE".equals(nodeToHandle.attributeValue("order", "AFTER"));
        String keyProperty = nodeToHandle.attributeValue("keyProperty");
        // default
        String resultMap = null;
        KeyGenerator keyGenerator = new NoKeyGenerator();

        // 解析成SqlSource，DynamicSqlSource/RawSqlSource
        SqlSource sqlSource = langDriver.createSqlSource(configuration, nodeToHandle, parameterTypeClass);
        // SELECT标签
        SqlCommandType sqlCommandType = SqlCommandType.SELECT;

        // 调用助手类
        builderAssistant.addMappedStatement(id,
                sqlSource,
                sqlCommandType,
                parameterTypeClass,
                resultMap,
                resultTypeClass,
                keyGenerator,
                keyProperty,
                langDriver);

        // 给id加上namespace前缀
        id = builderAssistant.applyCurrentNamespace(id, false);
        // 存放键值生成器配置
        MappedStatement keyStatement = configuration.getMappedStatement(id);

        // 将KeyGenerator放入map中
        configuration.addKeyGenerator(id, new SelectKeyGenerator(keyStatement, executeBefore));
    }
}

MapperAnnotationBuilder类，注解版的构建也是一样的，需要修改解析的语句，方法为parseStatement()

java 复制代码

public class MapperAnnotationBuilder {
  private void parseStatement(Method method) {
           // 省略其他
          
           // step-13 新增-----------------------------------------------------------
            KeyGenerator keyGenerator;
            String keyProperty = "id";
            if (SqlCommandType.INSERT.equals(sqlCommandType) || SqlCommandType.UPDATE.equals(sqlCommandType)) {
                keyGenerator = configuration.isUseGeneratedKeys() ? new Jdbc3KeyGenerator() : new NoKeyGenerator();
            } else {
                keyGenerator = new NoKeyGenerator();
            }
            // step-14 新增-----------------------------------------------------------
            boolean isSelect = sqlCommandType == SqlCommandType.SELECT;

            String resultMapId = null;
            if (isSelect) {
                resultMapId = parseResultMap(method);
            }

            // step-13 部分参数新增( keyGenerator,keyProperty,)-------------------------
            // 调用助手类
            assistant.addMappedStatement(
                    mappedStatementId,
                    sqlSource,
                    sqlCommandType,
                    parameterTypeClass,
                    resultMapId,
                    getReturnType(method),
                    keyGenerator,
                    keyProperty,
                    languageDriver
            );
  }
}

因为MappedStatement要存储数据，所以MappedStatement要加相应的字段，keyGenerator以及keyProperties和keyColumns

java 复制代码

public class MappedStatement {
    // 其余省略
    // step-13 新增
    private String resource;
    private KeyGenerator keyGenerator;
    private String[] keyProperties;
    private String[] keyColumns;

    public static class Builder {
          public Builder(Configuration configuration, String id, SqlCommandType 
          sqlCommandType, SqlSource sqlSource, Class<?> resultType) {
               mappedStatement.keyGenerator = configuration.isUseGeneratedKeys() && SqlCommandType.INSERT.equals(sqlCommandType) ? new Jdbc3KeyGenerator() : new NoKeyGenerator();
          }
 
        public Builder resource(String resource) {
            mappedStatement.resource = resource;
            return this;
        }

        public Builder keyGenerator(KeyGenerator keyGenerator) {
            mappedStatement.keyGenerator = keyGenerator;
            return this;
        }

        public Builder keyProperty(String keyProperty) {
            mappedStatement.keyProperties = delimitedStringToArray(keyProperty);
            return this;
        }
    }
    
    public String[] getKeyColumns() {
        return keyColumns;
    }

    public String[] getKeyProperties() {
        return keyProperties;
    }

    public KeyGenerator getKeyGenerator() {
        return keyGenerator;
    }
}

MapperBuilderAssistant类：此类修改了addMappedStatement()方法需要增加几个参数，并把参数放入MappedStatement中。

java 复制代码

  public MappedStatement addMappedStatement(
            String id,
            SqlSource sqlSource,
            SqlCommandType sqlCommandType,
            Class<?> parameterType,
            String resultMap,
            Class<?> resultType,
            KeyGenerator keyGenerator,
            String keyProperty,
            LanguageDriver lang
    ) {
        // 给id加上namespace前缀：cn.bugstack.mybatis.test.dao.IUserDao.queryUserInfoById
        id = applyCurrentNamespace(id, false);
        MappedStatement.Builder statementBuilder = new MappedStatement.Builder(configuration, id, sqlCommandType, sqlSource, resultType);
        // step-13新增/添加三个属性
        statementBuilder.resource(resource);
        statementBuilder.keyGenerator(keyGenerator);
        statementBuilder.keyProperty(keyProperty);

        // 结果映射，给 MappedStatement#resultMaps
        setStatementResultMap(resultMap, resultType, statementBuilder);

        MappedStatement statement = statementBuilder.build();
        // 映射语句信息，建造完存放到配置项中
        configuration.addMappedStatement(statement);

        return statement;
    }

3,3 insert执行SQL部分修改

包：package cn.bugstack.mybatis.executor

Executor接口需要新增个query方法，SelectKeyGenerator执行查询时使用。

java 复制代码

  // step-13新增-----------------------------------
  <E> List<E> query(MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler) throws SQLException;

BaseExecutor：基础实现类去实现这个query方法，然后调用原有的其他的query方法，然后顺着会调用SimpleExecutor的doQuery()这里的处理没有改变，所以就不提供代码拉。这块就和查询逻辑一样。

java 复制代码

    // step-14新增-----------------------------------------
    @Override
    public <E> List<E> query(MappedStatement ms, Object parameter, RowBounds rowBounds, ResultHandler resultHandler) throws SQLException {
        BoundSql boundSql = ms.getBoundSql(parameter);
        return query(ms, parameter, rowBounds, resultHandler, boundSql);
    }

在上面SimpleExecutor执行doQuery()时，需要调用到BaseStatementHandler的构造方法，此时就需要检查一下是否需要在insert前要执行，如果需要就执行SelectKey，selectKey又会调用上边的query进行查询操作，如果否则不会执行任何操作，判断则在keyGenerator.processBefore()处理。

java 复制代码

public abstract class BaseStatementHandler implements StatementHandler {
      public BaseStatementHandler(Executor executor, MappedStatement mappedStatement, Object parameterObject, RowBounds rowBounds, ResultHandler resultHandler, BoundSql boundSql) {
            if (boundSql == null) {
            // 针对Sequence主键而言，在执行insert sql前必须指定一个主键值给要插入的记录。
            generateKeys(parameterObject);
        }
    }

     // 针对Sequence主键而言，在执行insert sql前必须指定一个主键值给要插入的记录。KeyGenerator#processBefore
     protected void generateKeys(Object parameter) {
        KeyGenerator keyGenerator = mappedStatement.getKeyGenerator();
        keyGenerator.processBefore(executor, mappedStatement, null, parameter);
    }
}

然后就要修改下执行insert语句了，在Mybatis中所有的insert、del、update，都统一为update方法，所以我们需要修改下update()，修改为执行完insert操作后调用keyGenerator.processAfter()，后置处理（有主键情况下使用），修改如下：

PreparedStatementHandler类

包：package cn.bugstack.mybatis.executor.statement;

java 复制代码

    @Override
    public int update(Statement statement) throws SQLException {
        PreparedStatement ps = (PreparedStatement) statement;
        ps.execute();
        // step-13新增----------------------------------------
        // 执行 selectKey 语句
        int rows = ps.getUpdateCount();
        Object parameterObject = boundSql.getParameterObject();
        KeyGenerator keyGenerator = mappedStatement.getKeyGenerator();
        keyGenerator.processAfter(executor, mappedStatement, ps, parameterObject);
        // step-13新增----------------------------------------
        return rows;
    }

最后最重要的是，执行insert再执行查询当前插入记录的主键，必须是一个连接，如果是两个连接结果就是空，所以需要更改获取连接的操作，这里很简单，判断下就行了。

修改的包：cn.bugstack.mybatis.transaction.jdbc

修改的类：JdbcTransaction

java 复制代码

    @Override
    public Connection getConnection() throws SQLException {
        // step-13新增------------------------------------
        // 本章节新增；多个SQL在同一个JDBC连接下，才能完成事务特性
        if (connection != null) {
            return connection;
        }
        // step-13新增------------------------------------
        connection = dataSource.getConnection();
        connection.setTransactionIsolation(level.getLevel());
        connection.setAutoCommit(autoCommit);
        return connection;
    }

4.准备测试

xml如下：

XML 复制代码

     <insert id="insert" parameterType="cn.bugstack.mybatis.test.po.Activity">
        INSERT INTO activity
        (activity_id, activity_name, activity_desc, create_time, update_time)
        VALUES (#{activityId}, #{activityName}, #{activityDesc}, now(), now())

        <selectKey keyProperty="id" order="AFTER" resultType="long">
            SELECT LAST_INSERT_ID()
        </selectKey>

    </insert>

单元测试如下：

java 复制代码

    private Logger logger = LoggerFactory.getLogger(ApiTest.class);

    private SqlSession sqlSession;

    @Before
    public void init() throws IOException {
        // 1. 从SqlSessionFactory中获取SqlSession
        SqlSessionFactory sqlSessionFactory = new SqlSessionFactoryBuilder().build(Resources.getResourceAsReader("mybatis-config-datasource.xml"));
        sqlSession = sqlSessionFactory.openSession();
    }


    @Test
    public void test_insert() {
        // 1. 获取映射器对象
        IActivityDao dao = sqlSession.getMapper(IActivityDao.class);

        Activity activity = new Activity();
        activity.setActivityId(10004L);
        activity.setActivityName("测试活动");
        activity.setActivityDesc("测试数据插入");
        activity.setCreator("xdf");

        // 2. 测试验证
        Integer res = dao.insert(activity);
        sqlSession.commit();

        logger.info("测试结果：count：{} idx：{}", res, JSON.toJSONString(activity.getId()));
    }

执行结果：

执行的两条都展示了正确的索引。