CODEWITHSUNDEEP
javamysqlsqlspring
hese
hese

Reputation: 3417

BatchSqlUpdate - how to get auto generated keys

I am using spring BatchSqlUpdate to insert a set of rows. How do I get the auto generated keys for all of the rows inserted?

When doing a single insert I get the keys like this - 

SqlUpdate sqlUpdate = new SqlUpdate(dataSource, sqlTemplate.toString());
sqlUpdate.setReturnGeneratedKeys(true);
KeyHolder keyHolder = new GeneratedKeyHolder();
sqlUpdate.update(new Object[] {}, keyHolder);
return keyHolder.getKey().longValue();

Thanks!

Upvotes: 17

Views: 16555

Answers (4)

takeseem
takeseem

Reputation: 130

This solution is merged from the implementations of JdbcTemplate.update(PreparedStatementCreator, KeyHolder) and JdbcTemplate.batchUpdate(String, BatchPreparedStatementSetter) to allow having both batching and the generated keys.

In addition, we directly store the generated keys back to the beans where the inserted rows came from.

a template proc class:

public abstract class BatchPreparedStatementSetterWithKeyHolder<T> implements BatchPreparedStatementSetter {
    private final List<T> beans;
    
    /**
     * @param datas
     * @param returnGeneratedKeys true设置{@linkplain Statement#RETURN_GENERATED_KEYS}
     */
    public BatchPreparedStatementSetterWithKeyHolder(List<T> beans) {
        this.beans = beans;
    }
    
    
    @Override
    public void setValues(PreparedStatement ps, int i) throws SQLException {
        setValues(ps, beans.get(i));
    }

    @Override
    public final int getBatchSize() {
        return beans.size();
    }
    
    public void setPrimaryKey(KeyHolder keyHolder) {
        List<Map<String, Object>> keys = keyHolder.getKeyList();
        for (int i = 0, len = keys.size(); i < len; i++) {
            setPrimaryKey(keys.get(i), beans.get(i));
        }
    }

    protected abstract void setValues(PreparedStatement ps, T bean) throws SQLException;
    protected abstract void setPrimaryKey(Map<String, Object> primaryKey, T bean);
}

a util method batchupdate:

private static void generatedKeys(PreparedStatement ps, KeyHolder keyHolder) throws SQLException {
        List<Map<String, Object>> keys = keyHolder.getKeyList();
        ResultSet rs = ps.getGeneratedKeys();
        if (rs == null) return;
        
        try {
            keys.addAll(new RowMapperResultSetExtractor<Map<String, Object>>(new ColumnMapRowMapper(), 1).extractData(rs));
        } finally {
            rs.close();
        }
    }
    
    /**
     * 批量更新
     * @param jdbcTemplate
     * @param sql
     * @param pss
     * @param keyHolder 存储主键，如果要存储主键，就必须传入此对象
     * @return sql执行结果
     * @see JdbcTemplate#batchUpdate(String, org.springframework.jdbc.core.BatchPreparedStatementSetter)
     */
    public static <T> int[] batchUpdateWithKeyHolder(JdbcTemplate jdbcTemplate, final String sql, final BatchPreparedStatementSetterWithKeyHolder<T> pss) {
        return jdbcTemplate.execute(new PreparedStatementCreator() {
            @Override
            public PreparedStatement createPreparedStatement(Connection con) throws SQLException {
                return con.prepareStatement(sql, Statement.RETURN_GENERATED_KEYS);
            }
        }, new PreparedStatementCallback<int[]>() {
            @Override
            public int[] doInPreparedStatement(PreparedStatement ps) throws SQLException, DataAccessException {
                try {
                    int batchSize = pss.getBatchSize();
                    InterruptibleBatchPreparedStatementSetter ipss =
                            (pss instanceof InterruptibleBatchPreparedStatementSetter ?
                            (InterruptibleBatchPreparedStatementSetter) pss : null);
                    int[] result;
                    KeyHolder keyHolder = new GeneratedKeyHolder();
                    
                    try {
                        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();
                            }
                            result = ps.executeBatch();
                            
                            generatedKeys(ps, keyHolder);
                        } else {
                            List<Integer> rowsAffected = new ArrayList<Integer>();
                            for (int i = 0; i < batchSize; i++) {
                                pss.setValues(ps, i);
                                if (ipss != null && ipss.isBatchExhausted(i)) break;
                                
                                rowsAffected.add(ps.executeUpdate());
                                generatedKeys(ps, keyHolder);
                            }
                            
                            result =  rowsAffected.stream().mapToInt(Integer::intValue).toArray();
                        }
                    } finally {
                        pss.setPrimaryKey(keyHolder);
                    }
                    
                    return result;
                } finally {
                    if (pss instanceof ParameterDisposer) ((ParameterDisposer) pss).cleanupParameters();
                }
            }
        });
    }

e.x:

UtilJdbc.batchUpdateWithKeyHolder(jdbcTemplate,
                "insert into tbe_vm_node (creator_id, host, ssh_ip, ssh_user, ssh_passwd)" +
                " values (?, ?, ?, ?, ?)",
                new BatchPreparedStatementSetterWithKeyHolder<VmNode>(vmNodes) {
                    @Override
                    protected void setValues(PreparedStatement ps, VmNode vmNode) throws SQLException {
                        UtilJdbc.setValues(ps, vmNode.getCreatorId(), vmNode.getHost(),
                                vmNode.getSshIp(), vmNode.getSshUser(), vmNode.getSshPasswd());
                    }
                    @Override
                    protected void setPrimaryKey(Map<String, Object> primaryKey, VmNode vmNode) {
                        vmNode.setId((Long) primaryKey.get("abc"));
                    }
                });

Upvotes: 13

Dean Thomsen
Dean Thomsen

Reputation: 149

Thanks for the provided solution, I implemented that and it worked great. However, you have references to private classes (UtilIo, UtilObj) where I had to interpret what you were doing. I replaced "UtilIO.close(rs);" with:

if (rs != null) {
    rs.close();
}

And I replaced "result = UtilObj.tointArray(rowsAffected);" with this Java 8 code:

result = rowsAffected.stream().mapToInt(Integer::intValue).toArray();

Upvotes: 1

BrianC
BrianC

Reputation: 1822

    KeyHolder keyHolder = new GeneratedKeyHolder();
    BatchSqlUpdate updateBatch = new BatchSqlUpdate();
    updateBatch.setSql("Some INSERT SQL");
    updateBatch.setJdbcTemplate(jdbcTemp);
    updateBatch.setReturnGeneratedKeys(true);
    updateBatch.updateByNamedParam( paramsMap ,keyHolder);
    updateBatch.flush();

Upvotes: 1

abalogh
abalogh

Reputation: 8281

There is no provided solution for this using BatchSqlUpdate as far as I know, but you can always

  • query the last key before the insert
  • using this information, query all new keys after the insert

Upvotes: -6

Related Questions