SpringJdbcTemplate 批量新增返回自增编号功能

因为业务上有些功能是希望批量新增的数据能够返回对应的自增主键，单条的JdbcTemplate实现很简单，但是批量的话，可能需要稍微改造一点点。

因为我们的框架内部是在jdbctemplate上封装了一点点，所以就这一部分上加入了通用的批量添加功能，和批量修改隔离开。

/**
* 批量返回新增结果
*
* @param sql             要执行的SQL
* @param parameterSource 参数对象
* @return
* @throws SQLException
*/
private int[] batchInsert(String sql, SqlParameterSource[] parameterSource) throws SQLException {
    ParsedSql parsedSql = NamedParameterUtils.parseSqlStatement(sql);
    PreparedStatementCreatorFactory pscf = getPreparedStatementCreatorFactory(parsedSql, parameterSource[0]);

    /**
     * 参数执行处理
     */
    BatchPreparedStatementSetter pss = new BatchPreparedStatementSetter() {
        @Override
        public void setValues(PreparedStatement ps, int i) throws SQLException {
            // 将SQL的中的:参数转化成数组
            Object[] values = NamedParameterUtils.buildValueArray(parsedSql, parameterSource[i], null);
            // 转化成数组之后一个个和?对应填充
            pscf.newPreparedStatementSetter(values).setValues(ps);
        }

        @Override
        public int getBatchSize() {
            return parameterSource.length;
        }
    };
    /**
     * 创建Statement对象的时候，加上Statement.RETURN_GENERATED_KEYS参数
     */
    BatchPreparedStatementCreator batchPreparedStatementCreator = new BatchPreparedStatementCreator(pscf.getSql());
    return jdbcTemplate.execute(batchPreparedStatementCreator, (PreparedStatementCallback<int[]>) ps -> {
        try {
            int batchSize = pss.getBatchSize();
            int[] keys = new int[batchSize];

            InterruptibleBatchPreparedStatementSetter ipss =
                    (pss instanceof InterruptibleBatchPreparedStatementSetter ?
                            (InterruptibleBatchPreparedStatementSetter) pss : null);
            if (JdbcUtils.supportsBatchUpdates(ps.getConnection())) {
                for (int i = 0; i < batchSize; i++) {
                    pss.setValues(ps, i);
                    if (ipss != null && ipss.isBatchExhausted(i)) {
                        break;
                    }
                    ps.addBatch();
                }
                ps.executeBatch();
                // 将插入的主键值返回出来
                ResultSet rs = ps.getGeneratedKeys();
                int index = 0;
                while (rs.next() && index < batchSize) {
                    keys[index] = rs.getInt(1);
                    index++;
                }
                return keys;
            } else {
                List<Integer> rowsAffected = new ArrayList<>();
                for (int i = 0; i < batchSize; i++) {
                    pss.setValues(ps, i);
                    if (ipss != null && ipss.isBatchExhausted(i)) {
                        break;
                    }
                    rowsAffected.add(ps.executeUpdate());
                }
                int[] rowsAffectedArray = new int[rowsAffected.size()];
                for (int i = 0; i < rowsAffectedArray.length; i++) {
                    rowsAffectedArray[i] = rowsAffected.get(i);
                }
                return rowsAffectedArray;
            }
        } finally {
            if (pss instanceof ParameterDisposer) {
                ((ParameterDisposer) pss).cleanupParameters();
            }
        }
    });
}

protected PreparedStatementCreatorFactory getPreparedStatementCreatorFactory(
            ParsedSql parsedSql, SqlParameterSource paramSource) {
   String sqlToUse = NamedParameterUtils.substituteNamedParameters(parsedSql, paramSource);
    List<SqlParameter> declaredParameters = NamedParameterUtils.buildSqlParameterList(parsedSql, paramSource);
    return new PreparedStatementCreatorFactory(sqlToUse, declaredParameters);
}

BatchPreparedStatementCreator

import org.springframework.jdbc.core.PreparedStatementCreator;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;

/**
 * 为每个结果集加上Statement.RETURN_GENERATED_KEYS参数
 *
 * @author ： liukx
 * @time ： 2020/5/13 - 13:32
 */
public class BatchPreparedStatementCreator implements PreparedStatementCreator {
    private String[] generatedKeysColumnNames;
    private boolean returnGeneratedKeys = true;
    private String actualSql;

    public BatchPreparedStatementCreator(String actualSql) {
        this.actualSql = actualSql;
    }

    public void setGeneratedKeysColumnNames(String[] generatedKeysColumnNames) {
        this.generatedKeysColumnNames = generatedKeysColumnNames;
    }

    public void setReturnGeneratedKeys(boolean returnGeneratedKeys) {
        this.returnGeneratedKeys = returnGeneratedKeys;
    }

    @Override
    public PreparedStatement createPreparedStatement(Connection con) throws SQLException {
        PreparedStatement ps;
        if (generatedKeysColumnNames != null) {
            ps = con.prepareStatement(this.actualSql, generatedKeysColumnNames);
        } else {
            ps = con.prepareStatement(this.actualSql, PreparedStatement.RETURN_GENERATED_KEYS);
        }
        return ps;
    }
}

测试
第一个SQL参数参考 :

insert into t_test (username, name, sex, status, created, time, test_id, love_name) values (:username, :name, :sex, :status, :created, :time, :test_id, :love_name)

第二个参数:
将指定的对象转化成Map或者object 大概意思就是这样。o是代表本次操作的对象

SqlParameterSource[] parameterSource = null;
if (o instanceof Map) {
    parameterSource = new MapSqlParameterSource[list.size()];
 } else {
     parameterSource = new BeanSqlParameterSource[list.size()];
 }

大概讲一下几个关键地方:

1. BatchPreparedStatementCreator : 主要是为创建的结果集加入PreparedStatement.RETURN_GENERATED_KEYS 参数。
2. BatchPreparedStatementSetter : 主要是将:对应的name转化成 ? 最底层的语句然后执行。
3. 将生成的key拿到并返回出来

// 将插入的主键值返回出来
ResultSet rs = ps.getGeneratedKeys();
int index = 0;
while (rs.next() && index < batchSize) {
    keys[index] = rs.getInt(1);
    index++;
}

这个时候查看数组中的数据就已经加入了自增主键的集合。

有问题欢迎留言。