Spring Data Jpa 实现批量插入或更新

1. BatchConsumer

package com.demo.common.hibernate.batch;
 
import com.demo.common.hibernate.querydsl.QueryParameterBuilder;
 
/**
 * 批量数据消费者接口，用于设置 SQL 参数并执行操作。
 *
 * @param  记录类型的泛型
 * @author xm.z
 */
@FunctionalInterface
public interface BatchConsumer {
 
    /**
     * 设置 SQL 参数并执行操作。
     *
     * @param builder     参数构建对象
     * @param record      要处理的记录
     */
    void accept(QueryParameterBuilder builder, T record);
 
}

2. QueryParameterBuilder

package com.demo.common.hibernate.querydsl;
 
import lombok.AccessLevel;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;
import org.hibernate.jpa.TypedParameterValue;
import org.hibernate.type.*;
import org.springframework.util.Assert;
 
import javax.persistence.Query;
import java.math.BigDecimal;
import java.time.LocalDate;
import java.time.LocalDateTime;
import java.time.LocalTime;
import java.util.Date;
import java.util.concurrent.atomic.AtomicInteger;
 
/**
 * QueryParameterBuilder
 * 
 * A utility class for building parameters for query.
 *
 * @author xm.z
 */
@Slf4j
@Getter
public class QueryParameterBuilder {
 
    /**
     * The native query object to be used for parameter setting
     */
    private final Query nativeQuery;
 
    /**
     * The counter for parameter position
     */
    @Getter(value = AccessLevel.NONE)
    private final AtomicInteger position;
 
    /**
     * The current date and time when the QueryParameterBuilder instance is created
     */
    private final LocalDateTime now;
 
    /**
     * Private constructor to initialize QueryParameterBuilder
     */
    private QueryParameterBuilder(Query nativeQuery, AtomicInteger position) {
        this.nativeQuery = nativeQuery;
        this.position = position;
        this.now = LocalDateTime.now();
    }
 
    /**
     * Retrieves the current position of the parameter.
     *
     * @return The current position of the parameter.
     */
    public Integer obtainCurrentPosition() {
        return position.get();
    }
 
    /**
     * Create an instance of QueryParameterBuilder.
     *
     * @param nativeQuery The native query object
     * @param position    The parameter position counter
     * @return QueryParameterBuilder instance
     */
    public static QueryParameterBuilder create(Query nativeQuery, AtomicInteger position) {
        Assert.notNull(nativeQuery, "Native query must not be null");
        Assert.notNull(position, "Position must not be null");
        return new QueryParameterBuilder(nativeQuery, position);
    }
 
    /**
     * Set a parameter of type Long.
     *
     * @param value The Long value for the parameter
     * @return The current QueryParameterBuilder instance
     */
    public QueryParameterBuilder setParameter(Long value) {
        return this.setParameter(StandardBasicTypes.LONG, value);
    }
 
    /**
     * Set a parameter of type Integer.
     *
     * @param value The Integer value for the parameter
     * @return The current QueryParameterBuilder instance
     */
    public QueryParameterBuilder setParameter(Integer value) {
        return this.setParameter(StandardBasicTypes.INTEGER, value);
    }
 
    /**
     * Set a parameter of type BigDecimal.
     *
     * @param value The BigDecimal value for the parameter
     * @return The current QueryParameterBuilder instance
     */
    public QueryParameterBuilder setParameter(BigDecimal value) {
        return this.setParameter(StandardBasicTypes.BIG_DECIMAL, value);
    }
 
    /**
     * Set a parameter of type String.
     *
     * @param value The String value for the parameter
     * @return The current QueryParameterBuilder instance
     */
    public QueryParameterBuilder setParameter(String value) {
        return this.setParameter(StandardBasicTypes.STRING, value);
    }
 
    /**
     * Set a parameter of type Boolean.
     *
     * @param value The Boolean value for the parameter
     * @return The current QueryParameterBuilder instance
     */
    public QueryParameterBuilder setParameter(Boolean value) {
        return this.setParameter(StandardBasicTypes.BOOLEAN, value);
    }
 
    /**
     * Set a parameter of type Date.
     *
     * @param value The Date value for the parameter
     * @return The current QueryParameterBuilder instance
     */
    public QueryParameterBuilder setParameter(Date value) {
        return this.setParameter(StandardBasicTypes.DATE, value);
    }
 
    /**
     * Set a parameter of type LocalDate.
     *
     * @param value The LocalDate value for the parameter
     * @return The current QueryParameterBuilder instance
     */
    public QueryParameterBuilder setParameter(LocalDate value) {
        return this.setParameter(LocalDateType.INSTANCE, value);
    }
 
    /**
     * Set a parameter of type LocalTime.
     *
     * @param value The LocalTime value for the parameter
     * @return The current QueryParameterBuilder instance
     */
    public QueryParameterBuilder setParameter(LocalTime value) {
        return this.setParameter(LocalTimeType.INSTANCE, value);
    }
 
    /**
     * Set a parameter of type LocalDateTime.
     *
     * @param value The LocalDateTime value for the parameter
     * @return The current QueryParameterBuilder instance
     */
    public QueryParameterBuilder setParameter(LocalDateTime value) {
        return this.setParameter(LocalDateTimeType.INSTANCE, value);
    }
 
    /**
     * Add or include a query condition to the native query object and set the parameter value.
     *
     * @param type  The parameter type
     * @param value The parameter value
     * @return The current QueryParameterBuilder instance
     */
    public QueryParameterBuilder setParameter(Type type, Object value) {
        return this.setParameter(position.getAndIncrement(), type, value);
    }
 
    /**
     * Add or include a query condition to the native query object and set the parameter value at the specified position.
     *
     * @param position The position of the parameter in the query
     * @param type     The parameter type
     * @param value    The parameter value
     * @return The current QueryParameterBuilder instance
     */
    public QueryParameterBuilder setParameter(int position, Type type, Object value) {
        TypedParameterValue typedParameterValue = new TypedParameterValue(type, value);
        if (log.isDebugEnabled()) {
            log.debug("Setting parameter at position {}: {}", position, typedParameterValue);
        }
        nativeQuery.setParameter(position, typedParameterValue);
        return this;
    }
 
}

3. KeyValue

package com.demo.common.model;
 
import io.swagger.v3.oas.annotations.media.Schema;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.NoArgsConstructor;
 
import java.io.Serializable;
 
/**
 * 用于表示键值对的通用类
 *
 * @param  键的类型
 * @param  值的类型
 * @author xm.z
 */
@Data
@NoArgsConstructor
@AllArgsConstructor
public class KeyValue implements Serializable {
 
    private static final long serialVersionUID = 1L;
 
    /**
     * 键
     */
    @Schema(title = "键")
    private K key;
 
    /**
     * 值
     */
    @Schema(title = "值")
    private V value;
 
}

3. SqlUtil

package com.demo.common.hibernate.util;
 
import com.demo.common.hibernate.batch.BatchConsumer;
import com.demo.common.hibernate.querydsl.QueryParameterBuilder;
import com.demo.common.model.KeyValue;
import cn.hutool.extra.spring.SpringUtil;
import cn.hutool.core.collection.CollUtil;
import lombok.extern.slf4j.Slf4j;
import org.jetbrains.annotations.NotNull;
import org.springframework.lang.NonNull;
import org.springframework.lang.Nullable;
import org.springframework.util.Assert;
 
import javax.persistence.EntityManager;
import javax.persistence.Query;
import java.util.Collections;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.stream.Collectors;
 
/**
 * SqlUtil
 *
 * @author xm.z
 */
@Slf4j
@SuppressWarnings("all")
public class SqlUtil {
 
    /**
     * Default batch insert size.
     */
    public static final int DEFAULT_BATCH_SIZE = 100;
 
    /**
     * Private constructor.
     */
    private SqlUtil() {
    }
 
    /**
     * Batch insert records into the database.
     *
     * @param tableFields The table fields information
     * @param records     The list of records to be inserted
     * @param consumer    The consumer function interface for customizing the insert behavior
     * @param          The type of records
     * @return The number of records successfully inserted
     */
    public static  int batchInsert(@NonNull KeyValue> tableFields,
                                      @NonNull List records, @NonNull BatchConsumersuper T> consumer) {
        return batchInsert(DEFAULT_BATCH_SIZE, tableFields, records, consumer);
    }
 
    /**
     * Perform batch insert operation with the specified batch size.
     *
     * @param batchSize   the size of each batch for insertion
     * @param tableFields the key-value pair representing the table fields
     * @param records     the list of records to be inserted
     * @param consumer    the batch consumer for processing each batch of records
     * @param          the type of records
     * @return the total number of records successfully inserted
     */
    public static  int batchInsert(int batchSize, @NonNull KeyValue> tableFields,
                                      @NonNull List records, @NonNull BatchConsumersuper T> consumer) {
        EntityManager entityManager = SpringUtil.getBean(EntityManager.class);
        return batchExecuteUpdate(batchSize, entityManager, tableFields, null, records, consumer);
    }
 
    /**
     * Batch insert records into the database.
     *
     * @param entityManager The entity manager
     * @param tableFields   The table fields information
     * @param records       The list of records to be inserted
     * @param consumer      The consumer function interface for customizing the insert behavior
     * @param            The type of records
     * @return The number of records successfully inserted
     */
    public static  int batchInsert(EntityManager entityManager,
                                      @NonNull KeyValue> tableFields,
                                      @NonNull List records, @NonNull BatchConsumersuper T> consumer) {
        return batchExecuteUpdate(DEFAULT_BATCH_SIZE, entityManager, tableFields, null, records, consumer);
    }
 
    /**
     * Executes batch insert or update operations on the database using native SQL with a default batch size.
     *
     * @param tableFields  key-value pair representing the table name and its fields
     * @param updateFields set of fields to be updated if a record with matching primary key exists
     * @param records      the list of records to be inserted or updated
     * @param consumer     functional interface for accepting batch consumer operations
     * @param           the type of the records to be inserted or updated
     * @return the total number of rows affected by the batch operation
     */
    public static  int batchInsertOrUpdate(@NonNull KeyValue> tableFields,
                                              @NonNull LinkedHashSet updateFields,
                                              @NonNull List records, @NonNull BatchConsumersuper T> consumer) {
        return batchInsertOrUpdate(DEFAULT_BATCH_SIZE, tableFields, updateFields, records, consumer);
    }
 
    /**
     * Executes batch insert or update operations on the database using native SQL with a parameterized batch size.
     *
     * @param batchSize    the size of each batch for insertion
     * @param tableFields  key-value pair representing the table name and its fields
     * @param updateFields set of fields to be updated if a record with matching primary key exists
     * @param records      the list of records to be inserted or updated
     * @param consumer     functional interface for accepting batch consumer operations
     * @param           the type of the records to be inserted or updated
     * @return the total number of rows affected by the batch operation
     */
    public static  int batchInsertOrUpdate(int batchSize, @NonNull KeyValue> tableFields,
                                              @NonNull LinkedHashSet updateFields,
                                              @NonNull List records, @NonNull BatchConsumersuper T> consumer) {
        EntityManager entityManager = SpringUtil.getBean(EntityManager.class);
        return batchExecuteUpdate(batchSize, entityManager, tableFields, updateFields, records, consumer);
    }
 
    /**
     * Executes batch insert or update operations on the database using native SQL with a default batch size.
     *
     * @param entityManager The entity manager
     * @param tableFields   key-value pair representing the table name and its fields
     * @param updateFields  set of fields to be updated if a record with matching primary key exists
     * @param records       the list of records to be inserted or updated
     * @param consumer      functional interface for accepting batch consumer operations
     * @param            the type of the records to be inserted or updated
     * @return the total number of rows affected by the batch operation
     */
    public static  int batchInsertOrUpdate(EntityManager entityManager,
                                              @NonNull KeyValue> tableFields,
                                              @NonNull LinkedHashSet updateFields,
                                              @NonNull List records, @NonNull BatchConsumersuper T> consumer) {
        return batchExecuteUpdate(DEFAULT_BATCH_SIZE, entityManager, tableFields, updateFields, records, consumer);
    }
 
    /**
     * Executes batch updates on the database using native SQL with a parameterized batch size.
     *
     * @param batchSize     the size of each batch for inserting records
     * @param entityManager the entity manager for creating and executing queries
     * @param tableFields   key-value pair representing the table name and its fields
     * @param updateFields  set of fields to be updated if a record with matching primary key exists (optional)
     * @param records       the list of records to be inserted
     * @param consumer      functional interface for accepting batch consumer operations
     * @param            the type of the records to be inserted
     * @return the total number of rows affected by the batch operation
     */
    private static  int batchExecuteUpdate(int batchSize, EntityManager entityManager,
                                              @NonNull KeyValue> tableFields,
                                              @Nullable LinkedHashSet updateFields,
                                              @NonNull List records, @NonNull BatchConsumersuper T> consumer) {
        if (records.isEmpty()) {
            log.debug("No records to process. The records list is empty.");
            return 0;
        }
 
        Assert.notNull(entityManager, "The entity manager must not be null.");
        Assert.isTrue(batchSize > 0 && batchSize < 500, "The batch size must be between 1 and 500.");
 
        AtomicInteger totalRows = new AtomicInteger(0);
 
        // Split the records into batches based on the specified batch size
        List> recordBatches = CollUtil.split(records, batchSize);
 
        for (List batchRecords : recordBatches) {
            AtomicInteger position = new AtomicInteger(1);
 
            // Generate the appropriate SQL statement for the batch
            String preparedStatementSql = CollUtil.isEmpty(updateFields) ?
                    generateBatchInsertSql(tableFields, batchRecords.size()) :
                    generateBatchInsertOrUpdateSql(tableFields, updateFields, batchRecords.size());
 
            // Create a Query instance for executing native SQL statements
            Query nativeQuery = entityManager.createNativeQuery(preparedStatementSql);
 
            // Create a parameter builder instance using QueryParameterBuilder
            QueryParameterBuilder parameterBuilder = QueryParameterBuilder.create(nativeQuery, position);
 
            for (T record : batchRecords) {
                // Set parameters for the prepared statement
                consumer.accept(parameterBuilder, record);
            }
 
            // Execute the SQL statement and accumulate the affected rows
            totalRows.addAndGet(nativeQuery.executeUpdate());
        }
 
        // Return the total number of affected rows
        return totalRows.get();
    }
 
    /**
     * Generate batch insert SQL statement.
     *
     * 
     * This method generates an SQL statement for batch insertion into a specified table with the provided fields.
     * Example SQL statement:
     * 
     * {@code INSERT INTO TABLE_NAME ( field_1, field_2 ) VALUES ( value_1, value_2 ), (value_3, value_4); }
     *