GH-1159 Batch Insert Operations Within a Single Aggregate

pom.xml

@@ -5,7 +5,7 @@
 
 	<groupId>org.springframework.data</groupId>
 	<artifactId>spring-data-relational-parent</artifactId>
-	<version>2.4.0-SNAPSHOT</version>
+	<version>2.4.0-1159-batch-insert-within-aggregate-SNAPSHOT</version>
 	<packaging>pom</packaging>
 
 	<name>Spring Data Relational Parent</name>

spring-data-jdbc-distribution/pom.xml

@@ -14,7 +14,7 @@
 	<parent>
 		<groupId>org.springframework.data</groupId>
 		<artifactId>spring-data-relational-parent</artifactId>
-		<version>2.4.0-SNAPSHOT</version>
+		<version>2.4.0-1159-batch-insert-within-aggregate-SNAPSHOT</version>
 		<relativePath>../pom.xml</relativePath>
 	</parent>
 

spring-data-jdbc/pom.xml

@@ -6,7 +6,7 @@
 	<modelVersion>4.0.0</modelVersion>
 
 	<artifactId>spring-data-jdbc</artifactId>
-	<version>2.4.0-SNAPSHOT</version>
+	<version>2.4.0-1159-batch-insert-within-aggregate-SNAPSHOT</version>
 
 	<name>Spring Data JDBC</name>
 	<description>Spring Data module for JDBC repositories.</description>
@@ -15,7 +15,7 @@
 	<parent>
 		<groupId>org.springframework.data</groupId>
 		<artifactId>spring-data-relational-parent</artifactId>
-		<version>2.4.0-SNAPSHOT</version>
+		<version>2.4.0-1159-batch-insert-within-aggregate-SNAPSHOT</version>
 	</parent>
 
 	<properties>

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/AggregateChangeExecutor.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2020-2021 the original author or authors.
+ * Copyright 2020-2022 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -28,6 +28,7 @@
  *
  * @author Jens Schauder
  * @author Myeonghyeon Lee
+ * @author Chirag Tailor
  * @since 2.0
  */
 class AggregateChangeExecutor {
@@ -66,6 +67,8 @@ private void execute(DbAction<?> action, JdbcAggregateChangeExecutionContext exe
 				executionContext.executeInsertRoot((DbAction.InsertRoot<?>) action);
 			} else if (action instanceof DbAction.Insert) {
 				executionContext.executeInsert((DbAction.Insert<?>) action);
+			} else if (action instanceof DbAction.InsertBatch) {
+				executionContext.executeInsertBatch((DbAction.InsertBatch<?>) action);
 			} else if (action instanceof DbAction.UpdateRoot) {
 				executionContext.executeUpdateRoot((DbAction.UpdateRoot<?>) action);
 			} else if (action instanceof DbAction.Update) {

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/JdbcAggregateChangeExecutionContext.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2019-2021 the original author or authors.
+ * Copyright 2019-2022 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -25,13 +25,15 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.function.BiConsumer;
+import java.util.stream.Collectors;
 
 import org.springframework.dao.IncorrectUpdateSemanticsDataAccessException;
 import org.springframework.dao.OptimisticLockingFailureException;
 import org.springframework.data.jdbc.core.convert.DataAccessStrategy;
 import org.springframework.data.jdbc.core.convert.Identifier;
 import org.springframework.data.jdbc.core.convert.JdbcConverter;
 import org.springframework.data.jdbc.core.convert.JdbcIdentifierBuilder;
+import org.springframework.data.jdbc.core.convert.InsertSubject;
 import org.springframework.data.mapping.PersistentProperty;
 import org.springframework.data.mapping.PersistentPropertyAccessor;
 import org.springframework.data.mapping.PersistentPropertyPath;
@@ -51,6 +53,7 @@
  * @author Jens Schauder
  * @author Umut Erturk
  * @author Myeonghyeon Lee
+ * @author Chirag Tailor
  */
 class JdbcAggregateChangeExecutionContext {
 
@@ -87,11 +90,11 @@ <T> void executeInsertRoot(DbAction.InsertRoot<T> insert) {
 			T rootEntity = RelationalEntityVersionUtils.setVersionNumberOnEntity( //
 					insert.getEntity(), initialVersion, persistentEntity, converter);
 
-			id = accessStrategy.insert(rootEntity, insert.getEntityType(), Identifier.empty());
+			id = accessStrategy.insert(rootEntity, insert.getEntityType(), Identifier.empty(), insert.getIdValueSource());
 
 			setNewVersion(initialVersion);
 		} else {
-			id = accessStrategy.insert(insert.getEntity(), insert.getEntityType(), Identifier.empty());
+			id = accessStrategy.insert(insert.getEntity(), insert.getEntityType(), Identifier.empty(), insert.getIdValueSource());
 		}
 
 		add(new DbActionExecutionResult(insert, id));
@@ -100,10 +103,24 @@ <T> void executeInsertRoot(DbAction.InsertRoot<T> insert) {
 	<T> void executeInsert(DbAction.Insert<T> insert) {
 
 		Identifier parentKeys = getParentKeys(insert, converter);
-		Object id = accessStrategy.insert(insert.getEntity(), insert.getEntityType(), parentKeys);
+		Object id = accessStrategy.insert(insert.getEntity(), insert.getEntityType(), parentKeys, insert.getIdValueSource());
 		add(new DbActionExecutionResult(insert, id));
 	}
 
+	<T> void executeInsertBatch(DbAction.InsertBatch<T> insertBatch) {
+
+		List<DbAction.Insert<T>> inserts = insertBatch.getInserts();
+		List<InsertSubject<T>> insertSubjects = inserts.stream()
+				.map(insert -> InsertSubject.describedBy(insert.getEntity(), getParentKeys(insert, converter)))
+				.collect(Collectors.toList());
+
+		Object[] ids = accessStrategy.insert(insertSubjects, insertBatch.getEntityType(), insertBatch.getIdValueSource());
+
+		for (int i = 0; i < inserts.size(); i++) {
+			add(new DbActionExecutionResult(inserts.get(i), ids.length > 0 ? ids[i] : null));
+		}
+	}
+
 	<T> void executeUpdateRoot(DbAction.UpdateRoot<T> update) {
 
 		RelationalPersistentEntity<T> persistentEntity = getRequiredPersistentEntity(update.getEntityType());
@@ -155,18 +172,6 @@ <T> void executeDeleteAll(DbAction.DeleteAll<T> delete) {
 		accessStrategy.deleteAll(delete.getPropertyPath());
 	}
 
-	<T> void executeMerge(DbAction.Merge<T> merge) {
-
-		// temporary implementation
-		if (!accessStrategy.update(merge.getEntity(), merge.getEntityType())) {
-
-			Object id = accessStrategy.insert(merge.getEntity(), merge.getEntityType(), getParentKeys(merge, converter));
-			add(new DbActionExecutionResult(merge, id));
-		} else {
-			add(new DbActionExecutionResult());
-		}
-	}
-
 	<T> void executeAcquireLock(DbAction.AcquireLockRoot<T> acquireLock) {
 		accessStrategy.acquireLockById(acquireLock.getId(), LockMode.PESSIMISTIC_WRITE, acquireLock.getEntityType());
 	}

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/convert/BatchInsertStrategy.java

@@ -0,0 +1,36 @@
+/*
+ * Copyright 2022 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.jdbc.core.convert;
+
+import org.springframework.jdbc.core.namedparam.SqlParameterSource;
+
+/**
+ * Strategy for executing a batch insert.
+ *
+ * @author Chirag Tailor
+ * @since 2.4
+ */
+interface BatchInsertStrategy {
+
+	/**
+	 * @param sql the insert sql. Must not be {@code null}.
+	 * @param sqlParameterSources the sql parameters for each record to be inserted. Must not be {@code null}.
+	 * @return the ids corresponding to each record that was inserted, if ids were generated. If ids were not generated,
+	 *         elements will be {@code null}.
+	 * @since 2.4
+	 */
+	Object[] execute(String sql, SqlParameterSource[] sqlParameterSources);
+}

spring-data-jdbc/src/main/java/org/springframework/data/jdbc/core/convert/BatchJdbcOperations.java

@@ -0,0 +1,165 @@
+/*
+ * Copyright 2022 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.data.jdbc.core.convert;
+
+import org.springframework.jdbc.core.BatchPreparedStatementSetter;
+import org.springframework.jdbc.core.ColumnMapRowMapper;
+import org.springframework.jdbc.core.JdbcOperations;
+import org.springframework.jdbc.core.PreparedStatementCallback;
+import org.springframework.jdbc.core.PreparedStatementCreator;
+import org.springframework.jdbc.core.PreparedStatementCreatorFactory;
+import org.springframework.jdbc.core.RowMapperResultSetExtractor;
+import org.springframework.jdbc.core.SqlParameter;
+import org.springframework.jdbc.core.namedparam.NamedParameterUtils;
+import org.springframework.jdbc.core.namedparam.ParsedSql;
+import org.springframework.jdbc.core.namedparam.SqlParameterSource;
+import org.springframework.jdbc.support.JdbcUtils;
+import org.springframework.jdbc.support.KeyHolder;
+import org.springframework.lang.Nullable;
+import org.springframework.util.Assert;
+
+import java.sql.PreparedStatement;
+import java.sql.ResultSet;
+import java.sql.SQLException;
+import java.util.ArrayList;
+import java.util.List;
+import java.util.Map;
+
+/**
+ * Counterpart to {@link org.springframework.jdbc.core.namedparam.NamedParameterJdbcOperations} containing
+ * methods for performing batch updates with generated keys.
+ *
+ * @author Chirag Tailor
+ * @since 2.4
+ */
+public class BatchJdbcOperations {
+	private final JdbcOperations jdbcOperations;
+
+	public BatchJdbcOperations(JdbcOperations jdbcOperations) {
+		this.jdbcOperations = jdbcOperations;
+	}
+
+	/**
+	 * Execute a batch using the supplied SQL statement with the batch of supplied arguments,
+	 * returning generated keys.
+	 * @param sql the SQL statement to execute
+	 * @param batchArgs the array of {@link SqlParameterSource} containing the batch of
+	 * arguments for the query
+	 * @param generatedKeyHolder a {@link KeyHolder} that will hold the generated keys
+	 * @return an array containing the numbers of rows affected by each update in the batch
+	 * (may also contain special JDBC-defined negative values for affected rows such as
+	 * {@link java.sql.Statement#SUCCESS_NO_INFO}/{@link java.sql.Statement#EXECUTE_FAILED})
+	 * @throws org.springframework.dao.DataAccessException if there is any problem issuing the update
+	 * @see org.springframework.jdbc.support.GeneratedKeyHolder
+	 * @since 2.4
+	 */
+	int[] batchUpdate(String sql, SqlParameterSource[] batchArgs, KeyHolder generatedKeyHolder) {
+		return batchUpdate(sql, batchArgs, generatedKeyHolder, null);
+	}
+
+	/**
+	 * Execute a batch using the supplied SQL statement with the batch of supplied arguments,
+	 * returning generated keys.
+	 * @param sql the SQL statement to execute
+	 * @param batchArgs the array of {@link SqlParameterSource} containing the batch of
+	 * arguments for the query
+	 * @param generatedKeyHolder a {@link KeyHolder} that will hold the generated keys
+	 * @param keyColumnNames names of the columns that will have keys generated for them
+	 * @return an array containing the numbers of rows affected by each update in the batch
+	 * (may also contain special JDBC-defined negative values for affected rows such as
+	 * {@link java.sql.Statement#SUCCESS_NO_INFO}/{@link java.sql.Statement#EXECUTE_FAILED})
+	 * @throws org.springframework.dao.DataAccessException if there is any problem issuing the update
+	 * @see org.springframework.jdbc.support.GeneratedKeyHolder
+	 * @since 2.4
+	 */
+	int[] batchUpdate(String sql, SqlParameterSource[] batchArgs, KeyHolder generatedKeyHolder,
+			@Nullable String[] keyColumnNames) {
+
+		if (batchArgs.length == 0) {
+			return new int[0];
+		}
+
+		ParsedSql parsedSql = NamedParameterUtils.parseSqlStatement(sql);
+		SqlParameterSource paramSource = batchArgs[0];
+		String sqlToUse = NamedParameterUtils.substituteNamedParameters(parsedSql, paramSource);
+		List<SqlParameter> declaredParameters = NamedParameterUtils.buildSqlParameterList(parsedSql, paramSource);
+		PreparedStatementCreatorFactory pscf = new PreparedStatementCreatorFactory(sqlToUse, declaredParameters);
+		if (keyColumnNames != null) {
+			pscf.setGeneratedKeysColumnNames(keyColumnNames);
+		} else {
+			pscf.setReturnGeneratedKeys(true);
+		}
+		Object[] params = NamedParameterUtils.buildValueArray(parsedSql, paramSource, null);
+		PreparedStatementCreator psc = pscf.newPreparedStatementCreator(params);
+		BatchPreparedStatementSetter bpss = new BatchPreparedStatementSetter() {
+			@Override
+			public void setValues(PreparedStatement ps, int i) throws SQLException {
+				Object[] values = NamedParameterUtils.buildValueArray(parsedSql, batchArgs[i], null);
+				pscf.newPreparedStatementSetter(values).setValues(ps);
+			}
+
+			@Override
+			public int getBatchSize() {
+				return batchArgs.length;
+			}
+		};
+		PreparedStatementCallback<int[]> preparedStatementCallback = ps -> {
+			int batchSize = bpss.getBatchSize();
+			generatedKeyHolder.getKeyList().clear();
+			if (JdbcUtils.supportsBatchUpdates(ps.getConnection())) {
+				for (int i = 0; i < batchSize; i++) {
+					bpss.setValues(ps, i);
+					ps.addBatch();
+				}
+				int[] results = ps.executeBatch();
+				storeGeneratedKeys(generatedKeyHolder, ps, batchSize);
+				return results;
+			} else {
+				List<Integer> rowsAffected = new ArrayList<>();
+				for (int i = 0; i < batchSize; i++) {
+					bpss.setValues(ps, i);
+					rowsAffected.add(ps.executeUpdate());
+					storeGeneratedKeys(generatedKeyHolder, ps, 1);
+				}
+				int[] rowsAffectedArray = new int[rowsAffected.size()];
+				for (int i = 0; i < rowsAffectedArray.length; i++) {
+					rowsAffectedArray[i] = rowsAffected.get(i);
+				}
+				return rowsAffectedArray;
+			}
+		};
+		int[] result = jdbcOperations.execute(psc, preparedStatementCallback);
+		Assert.state(result != null, "No result array");
+		return result;
+	}
+
+	private void storeGeneratedKeys(KeyHolder generatedKeyHolder, PreparedStatement ps, int rowsExpected) throws SQLException {
+
+		List<Map<String, Object>> generatedKeys = generatedKeyHolder.getKeyList();
+		ResultSet keys = ps.getGeneratedKeys();
+		if (keys != null) {
+			try {
+				RowMapperResultSetExtractor<Map<String, Object>> rse =
+						new RowMapperResultSetExtractor<>(new ColumnMapRowMapper(), rowsExpected);
+				//noinspection ConstantConditions
+				generatedKeys.addAll(rse.extractData(keys));
+			}
+			finally {
+				JdbcUtils.closeResultSet(keys);
+			}
+		}
+	}
+}