SpringDataJpa IDENTITY 主键生成方式批量插入

JPA IDENTITY 主键批量插入

问题描述

当JPA实体主键生成方式为IDENTITY时，hibernate不支持批量插入，即使开启批量插入，hibernate 也会进行单条插入

解决方法

• 通过JdbcTemplate 进行插入
• 通过原生Jdbc原始SQL
• Mybatis实现
• 其他框架
  这里主要使用JdbcTemplate实现

JdbcTemplate 批量插入

适用范围

• JPA 2.X
• Java 8
• hibernate 5.X
• 主键生成方式为IDENTITY

实现方式

利用hibernate实现的JPA功能组件，生成需要的SQL。这种方式不需要手动写SQL且不需要考虑数据库差异。

• AbstractEntityPersister ： 存储实体类信息
  • 获取插入语句：getSQLIdentityInsertString
  • 获取实体插入参数列表：getPropertyValuesToInsert(Object entity, Map mergeMap, SessionImplementor session)
• CallbackRegistry -> entity增强
  • 实体创建前处理 ：preCreate(Object entity)
  • 实体创建后处理 ：postCreate(Object entity)

示例代码

package net.reduck.jpa;

import org.hibernate.internal.SessionFactoryImpl;
import org.hibernate.jpa.event.internal.CallbacksFactory;
import org.hibernate.jpa.event.spi.CallbackRegistry;
import org.hibernate.metamodel.internal.MetamodelImpl;
import org.hibernate.persister.entity.AbstractEntityPersister;
import org.springframework.jdbc.core.JdbcTemplate;
import org.springframework.orm.jpa.LocalContainerEntityManagerFactoryBean;
import org.springframework.stereotype.Repository;
import org.springframework.transaction.annotation.Transactional;

import javax.persistence.EntityManagerFactory;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;

/**
 * @author Reduck
 * @since 2023/5/22 1:26
 */
@Repository
@Transactional
public class BatchRepository {

    private final JdbcTemplate jdbcTemplate;
    private final MetamodelImpl metamodel;
    private final CallbackRegistry callbackRegistry;

    public BatchRepository(LocalContainerEntityManagerFactoryBean localContainerEntityManagerFactoryBean, EntityManagerFactory factory, JdbcTemplate jdbcTemplate) {
        this.jdbcTemplate = jdbcTemplate;
        this.callbackRegistry = CallbacksFactory.buildCallbackRegistry((SessionFactoryImpl) localContainerEntityManagerFactoryBean.getNativeEntityManagerFactory());
        this.metamodel = ((MetamodelImpl) factory.getMetamodel());
    }


    public <T> void insert(List<T> entities, Class<T> entityType) {
        int batchSize = 3000;

        if (entities == null || entities.size() == 0) {
            return;
        }

        // 获取持久化类
        AbstractEntityPersister persister = ((AbstractEntityPersister) metamodel.locateEntityPersister(entityType.getName()));
        // 获取IDENTITY类型insertSQL
        String insertSQL = persister.getSQLIdentityInsertString();

        if (insertSQL == null) {
            throw new RuntimeException("Just support id generate by IDENTITY");
        }

        String valuesTemplate = insertSQL.substring(insertSQL.lastIndexOf("("));
        StringBuilder stringBuilder = new StringBuilder(insertSQL);

        System.out.println("Total insert " + entities.size());
        for (int i = 0; i < entities.size(); i += batchSize) {
            List<T> batch = entities.subList(i, Math.min(i + batchSize, entities.size()));
            if (batch.size() == 0) {
                continue;
            }

            List<Object> params = new ArrayList<>();
            boolean first = true;
            for (T entity : batch) {
                // 参数预处理
                callbackRegistry.preCreate(entity);
                if (!first) {
                    stringBuilder.append(",").append(valuesTemplate);
                }
                first = false;
                // 参数填充
                params.addAll(Arrays.asList(persister.getPropertyValuesToInsert(entity, null, null)));
            }
            jdbcTemplate.update(stringBuilder.toString(), params.toArray());
            System.out.println("Finished insert " + i);
        }
    }
}