当前位置：首页 > news >正文

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

news 来源：原创 2024/7/1 17:14:55

1. BatchConsumer

package com.demo.common.hibernate.batch;import com.demo.common.hibernate.querydsl.QueryParameterBuilder;/*** 批量数据消费者接口，用于设置 SQL 参数并执行操作。** @param <T> 记录类型的泛型* @author xm.z*/
@FunctionalInterface
public interface BatchConsumer<T> {/*** 设置 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* <p>* 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 <K> 键的类型* @param <V> 值的类型* @author xm.z*/
@Data
@NoArgsConstructor
@AllArgsConstructor
public class KeyValue<K, V> 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 <T>         The type of records* @return The number of records successfully inserted*/public static <T> int batchInsert(@NonNull KeyValue<String, LinkedHashSet<String>> tableFields,@NonNull List<T> records, @NonNull BatchConsumer<? super 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 <T>         the type of records* @return the total number of records successfully inserted*/public static <T> int batchInsert(int batchSize, @NonNull KeyValue<String, LinkedHashSet<String>> tableFields,@NonNull List<T> records, @NonNull BatchConsumer<? super 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 <T>           The type of records* @return The number of records successfully inserted*/public static <T> int batchInsert(EntityManager entityManager,@NonNull KeyValue<String, LinkedHashSet<String>> tableFields,@NonNull List<T> records, @NonNull BatchConsumer<? super 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 <T>          the type of the records to be inserted or updated* @return the total number of rows affected by the batch operation*/public static <T> int batchInsertOrUpdate(@NonNull KeyValue<String, LinkedHashSet<String>> tableFields,@NonNull LinkedHashSet<String> updateFields,@NonNull List<T> records, @NonNull BatchConsumer<? super 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 <T>          the type of the records to be inserted or updated* @return the total number of rows affected by the batch operation*/public static <T> int batchInsertOrUpdate(int batchSize, @NonNull KeyValue<String, LinkedHashSet<String>> tableFields,@NonNull LinkedHashSet<String> updateFields,@NonNull List<T> records, @NonNull BatchConsumer<? super 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 <T>           the type of the records to be inserted or updated* @return the total number of rows affected by the batch operation*/public static <T> int batchInsertOrUpdate(EntityManager entityManager,@NonNull KeyValue<String, LinkedHashSet<String>> tableFields,@NonNull LinkedHashSet<String> updateFields,@NonNull List<T> records, @NonNull BatchConsumer<? super 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 <T>           the type of the records to be inserted* @return the total number of rows affected by the batch operation*/private static <T> int batchExecuteUpdate(int batchSize, EntityManager entityManager,@NonNull KeyValue<String, LinkedHashSet<String>> tableFields,@Nullable LinkedHashSet<String> updateFields,@NonNull List<T> records, @NonNull BatchConsumer<? super 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 sizeList<List<T>> recordBatches = CollUtil.split(records, batchSize);for (List<T> batchRecords : recordBatches) {AtomicInteger position = new AtomicInteger(1);// Generate the appropriate SQL statement for the batchString preparedStatementSql = CollUtil.isEmpty(updateFields) ?generateBatchInsertSql(tableFields, batchRecords.size()) :generateBatchInsertOrUpdateSql(tableFields, updateFields, batchRecords.size());// Create a Query instance for executing native SQL statementsQuery nativeQuery = entityManager.createNativeQuery(preparedStatementSql);// Create a parameter builder instance using QueryParameterBuilderQueryParameterBuilder parameterBuilder = QueryParameterBuilder.create(nativeQuery, position);for (T record : batchRecords) {// Set parameters for the prepared statementconsumer.accept(parameterBuilder, record);}// Execute the SQL statement and accumulate the affected rowstotalRows.addAndGet(nativeQuery.executeUpdate());}// Return the total number of affected rowsreturn totalRows.get();}/*** Generate batch insert SQL statement.** <p>* This method generates an SQL statement for batch insertion into a specified table with the provided fields.* Example SQL statement:* <pre>* {@code INSERT INTO TABLE_NAME ( field_1, field_2 ) VALUES ( value_1, value_2 ), (value_3, value_4); }* </pre>* </p>** @param tableFields The key-value pair representing the table name and its associated field set* @param batchSize   The batch size for insertion* @return The batch insert SQL statement*/private static String generateBatchInsertSql(@NonNull KeyValue<String, LinkedHashSet<String>> tableFields, int batchSize) {String preparedStatementSql = generateInsertStatement(tableFields.getKey(), tableFields.getValue(), batchSize);if (log.isDebugEnabled()) {log.debug("[Batch Insert] Prepared {} records SQL: {}", batchSize, preparedStatementSql);}return preparedStatementSql;}/*** Generates SQL statement for batch insert with on duplicate key update.** @param tableFields  Key-value pair representing table name and its corresponding fields.* @param updateFields Fields to be updated in case of duplicate key.* @param batchSize    Number of records to be inserted in a single batch.* @return SQL statement for batch insert with on duplicate key update.* @throws IllegalArgumentException if updateFields collection is empty.*/private static String generateBatchInsertOrUpdateSql(@NonNull KeyValue<String, LinkedHashSet<String>> tableFields,LinkedHashSet<String> updateFields, int batchSize) {Assert.notEmpty(updateFields, "Update field collection cannot be empty.");// Generate the insert statementString insertStatement = generateInsertStatement(tableFields.getKey(), tableFields.getValue(), batchSize);// Initialize StringBuilder with initial capacityStringBuilder builder = new StringBuilder(insertStatement.length() + 100);// Append insert statementbuilder.append(insertStatement).append(" ON DUPLICATE KEY UPDATE ");// Append update clauseString updateClause = updateFields.stream().map(updateField -> updateField + " = VALUES(" + updateField + ")").collect(Collectors.joining(", "));builder.append(updateClause);String preparedStatementSql = builder.toString();if (log.isDebugEnabled()) {log.debug("[Batch Insert On Duplicate Key Update] Prepared {} records SQL: {}", batchSize, preparedStatementSql);}return preparedStatementSql;}@NotNullprivate static String generateInsertStatement(@NonNull String tableName, @NonNull LinkedHashSet<String> fields, int batchSize) {Assert.hasText(tableName, "Table name cannot be empty.");Assert.notNull(fields, "Field collection cannot be empty.");// Set a reasonable initial capacityStringBuilder builder = new StringBuilder(fields.size() * 100);// Concatenate field namesString fieldNames = String.join(", ", fields);String intoTemplate = String.format("INSERT INTO %s (%s) VALUES ", tableName, fieldNames);// Generate placeholdersString placeholders = "(" + String.join(", ", Collections.nCopies(fields.size(), "?")) + ")";// Construct the insert statementbuilder.append(intoTemplate);for (int i = 0; i < batchSize; i++) {if (i > 0) {builder.append(", ");}builder.append(placeholders);}return builder.toString();}}