Provide all counters in BatchUpdateException

This commit updates JbcTemplate#batchUpdate to provide additional
information when one batch fails. Previously, the raw
BatchUpdateException was thrown with no way to know what had completed
thus far.

This commit creates an AggregatedBatchUpdateException that wraps the
original BatchUpdateException, yet providing the counters of the batches
that ran prior to the exception. In essence, this represents the same
state as the return value of the method if no batch fails.

AggregateBatchUpdateException exposes the original BatchUpdateException
in advanced case, such as checking for a sub-class that may contain
additional information.

Closes gh-23867

Author: snicoll

spring-jdbc/src/main/java/org/springframework/jdbc/core/AggregatedBatchUpdateException.java

@@ -0,0 +1,73 @@
+/*
+ * Copyright 2002-2024 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.jdbc.core;
+
+import java.sql.BatchUpdateException;
+
+/**
+ * A {@link BatchUpdateException} that provides additional information about
+ * batches that were successful prior to one failing.
+ *
+ * @author Stephane Nicoll
+ * @since 6.2
+ */
+@SuppressWarnings("serial")
+public class AggregatedBatchUpdateException extends BatchUpdateException {
+
+	private final int[][] successfulUpdateCounts;
+
+	private final BatchUpdateException originalException;
+
+	/**
+	 * Create an aggregated exception with the batches that have completed prior
+	 * to the given {@code cause}.
+	 * @param successfulUpdateCounts the counts of the batches that run successfully
+	 * @param original the exception this instance aggregates
+	 */
+	public AggregatedBatchUpdateException(int[][] successfulUpdateCounts, BatchUpdateException original) {
+		super(original.getMessage(), original.getSQLState(), original.getErrorCode(),
+				original.getUpdateCounts(), original.getCause());
+		this.successfulUpdateCounts = successfulUpdateCounts;
+		this.originalException = original;
+		// Copy state of the original exception
+		setNextException(original.getNextException());
+		for (Throwable suppressed : original.getSuppressed()) {
+			addSuppressed(suppressed);
+		}
+	}
+
+	/**
+	 * Return the batches that have completed successfully, prior to this exception.
+	 * <p>Information about the batch that failed is available via
+	 * {@link #getUpdateCounts()}.
+	 * @return an array containing for each batch another array containing the numbers of
+	 * rows affected by each update in the batch
+	 * @see #getUpdateCounts()
+	 */
+	public int[][] getSuccessfulUpdateCounts() {
+		return this.successfulUpdateCounts;
+	}
+
+	/**
+	 * Return the original {@link BatchUpdateException} that this exception aggregates.
+	 * @return the original exception
+	 */
+	public BatchUpdateException getOriginalException() {
+		return this.originalException;
+	}
+
+}

spring-jdbc/src/main/java/org/springframework/jdbc/core/JdbcTemplate.java

@@ -1116,7 +1116,13 @@ public <T> int[][] batchUpdate(String sql, final Collection<T> batchArgs, final
 								int items = n - ((n % batchSize == 0) ? n / batchSize - 1 : (n / batchSize)) * batchSize;
 								logger.trace("Sending SQL batch update #" + batchIdx + " with " + items + " items");
 							}
-							rowsAffected.add(ps.executeBatch());
+							try {
+								int[] updateCounts = ps.executeBatch();
+								rowsAffected.add(updateCounts);
+							}
+							catch (BatchUpdateException ex) {
+								throw new AggregatedBatchUpdateException(rowsAffected.toArray(int[][]::new), ex);
+							}
 						}
 					}
 					else {

spring-jdbc/src/test/java/org/springframework/jdbc/core/JdbcTemplateTests.java

@@ -32,13 +32,16 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Map;
+import java.util.function.Consumer;
 
 import javax.sql.DataSource;
 
+import org.assertj.core.data.Index;
 import org.junit.jupiter.api.BeforeEach;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.dao.DataAccessException;
+import org.springframework.dao.DuplicateKeyException;
 import org.springframework.dao.InvalidDataAccessApiUsageException;
 import org.springframework.jdbc.BadSqlGrammarException;
 import org.springframework.jdbc.CannotGetJdbcConnectionException;
@@ -799,6 +802,55 @@ void testBatchUpdateWithCollectionOfObjects() throws Exception {
 		verify(this.connection, atLeastOnce()).close();
 	}
 
+	@Test
+	void testBatchUpdateWithBatchFailingHasUpdateCounts() throws Exception {
+		test3BatchesOf2ItemsFailing(exception -> assertThat(exception).cause()
+				.isInstanceOfSatisfying(AggregatedBatchUpdateException.class, ex -> {
+					assertThat(ex.getSuccessfulUpdateCounts()).hasDimensions(1, 2)
+							.contains(new int[] { 1, 1 }, Index.atIndex(0));
+					assertThat(ex.getUpdateCounts()).contains(-3, -3);
+				}));
+	}
+
+	@Test
+	void testBatchUpdateWithBatchFailingMatchesOriginalException() throws Exception {
+		test3BatchesOf2ItemsFailing(exception -> assertThat(exception).cause()
+				.isInstanceOfSatisfying(AggregatedBatchUpdateException.class, ex -> {
+					BatchUpdateException originalException = ex.getOriginalException();
+					assertThat(ex.getMessage()).isEqualTo(originalException.getMessage());
+					assertThat(ex.getCause()).isEqualTo(originalException.getCause());
+					assertThat(ex.getSQLState()).isEqualTo(originalException.getSQLState());
+					assertThat(ex.getErrorCode()).isEqualTo(originalException.getErrorCode());
+					assertThat((Exception) ex.getNextException()).isSameAs(originalException.getNextException());
+					assertThat(ex.getSuppressed()).isEqualTo(originalException.getSuppressed());
+				}));
+	}
+
+	void test3BatchesOf2ItemsFailing(Consumer<Exception> exception) throws Exception {
+		String sql = "INSERT INTO NOSUCHTABLE values (?)";
+		List<Integer> ids = Arrays.asList(1, 2, 3, 2, 4, 5);
+		int[] rowsAffected = new int[] {1, 1};
+
+		given(this.preparedStatement.executeBatch()).willReturn(rowsAffected).willThrow(new BatchUpdateException(
+				"duplicate key value violates unique constraint \"NOSUCHTABLE_pkey\" Detail: Key (id)=(2) already exists.",
+				"23505", 0, new int[] { -3, -3 }));
+		mockDatabaseMetaData(true);
+
+		ParameterizedPreparedStatementSetter<Integer> setter = (ps, argument) -> ps.setInt(1, argument);
+		JdbcTemplate template = new JdbcTemplate(this.dataSource, false);
+
+		assertThatExceptionOfType(DuplicateKeyException.class)
+				.isThrownBy(() -> template.batchUpdate(sql, ids, 2, setter))
+				.satisfies(exception);
+		verify(this.preparedStatement, times(4)).addBatch();
+		verify(this.preparedStatement).setInt(1, 1);
+		verify(this.preparedStatement, times(2)).setInt(1, 2);
+		verify(this.preparedStatement).setInt(1, 3);
+		verify(this.preparedStatement, times(2)).executeBatch();
+		verify(this.preparedStatement).close();
+		verify(this.connection, atLeastOnce()).close();
+	}
+
 	@Test
 	void testCouldNotGetConnectionForOperationOrExceptionTranslator() throws SQLException {
 		SQLException sqlException = new SQLException("foo", "07xxx");