Update BatchJdbcOperations to more closely mirror proprosed updates to NamedParameterJdbcOperations in spring-jdbc.

Author: ctailor2

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

@@ -18,6 +18,7 @@
 import java.sql.PreparedStatement;
 import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 
@@ -29,89 +30,71 @@ public BatchJdbcOperations(JdbcOperations jdbcOperations) {
 	}
 
 	@Nullable
-	int[] batchUpdate(String insertSql, SqlParameterSource[] sqlParameterSources,
-					  KeyHolder generatedKeyHolder) {
-//		TODO: This is largely duplicated from spring-jdbc and should be replaced with a call into
-//		NamedParameterJdbcTemplate#batchUpdate once a method taking KeyHolder is added there.
-		ParsedSql parsedSql = NamedParameterUtils.parseSqlStatement(insertSql);
-		SqlParameterSource paramSource = sqlParameterSources[0];
+	int[] batchUpdate(String sql, SqlParameterSource[] batchArgs, KeyHolder generatedKeyHolder) {
+		return batchUpdate(sql, batchArgs, generatedKeyHolder, null);
+	}
+
+	@Nullable
+	public int[] batchUpdate(String sql, SqlParameterSource[] batchArgs, KeyHolder generatedKeyHolder,
+			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);
-		pscf.setReturnGeneratedKeys(true);
+		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, sqlParameterSources[i], null);
+				Object[] values = NamedParameterUtils.buildValueArray(parsedSql, batchArgs[i], null);
 				pscf.newPreparedStatementSetter(values).setValues(ps);
 			}
 
 			@Override
 			public int getBatchSize() {
-				return sqlParameterSources.length;
+				return batchArgs.length;
 			}
 		};
 		PreparedStatementCallback<int[]> preparedStatementCallback = ps -> {
 			int batchSize = bpss.getBatchSize();
-			for (int i = 0; i < batchSize; i++) {
-				bpss.setValues(ps, i);
-				ps.addBatch();
-			}
-			int[] results = ps.executeBatch();
-			List<Map<String, Object>> generatedKeys = ((KeyHolder) generatedKeyHolder).getKeyList();
-			generatedKeys.clear();
-			ResultSet keys = ps.getGeneratedKeys();
-			if (keys != null) {
-				try {
-					RowMapperResultSetExtractor<Map<String, Object>> rse =
-							new RowMapperResultSetExtractor<>(new ColumnMapRowMapper(), 1);
-					//noinspection ConstantConditions
-					generatedKeys.addAll(rse.extractData(keys));
-				} finally {
-					JdbcUtils.closeResultSet(keys);
+			if (JdbcUtils.supportsBatchUpdates(ps.getConnection())) {
+				for (int i = 0; i < batchSize; i++) {
+					bpss.setValues(ps, i);
+					ps.addBatch();
+				}
+				int[] results = ps.executeBatch();
+				storeGeneratedKeys(generatedKeyHolder).doInPreparedStatement(ps);
+				return results;
+			} else {
+				List<Integer> rowsAffected = new ArrayList<>();
+				for (int i = 0; i < batchSize; i++) {
+					bpss.setValues(ps, i);
+					rowsAffected.add(ps.executeUpdate());
 				}
+				int[] rowsAffectedArray = new int[rowsAffected.size()];
+				storeGeneratedKeys(generatedKeyHolder).doInPreparedStatement(ps);
+				for (int i = 0; i < rowsAffectedArray.length; i++) {
+					rowsAffectedArray[i] = rowsAffected.get(i);
+				}
+				return rowsAffectedArray;
 			}
-			return results;
 		};
 		return jdbcOperations.execute(psc, preparedStatementCallback);
 	}
 
-	@Nullable
-	public int[] batchUpdate(String insertSql, SqlParameterSource[] sqlParameterSources,
-							 KeyHolder generatedKeyHolder, String[] keyColumnNames) {
-//		TODO: This is largely duplicated from spring-jdbc and should be replaced with a call into
-//		NamedParameterJdbcTemplate#batchUpdate once a method taking KeyHolder is added there.
-		ParsedSql parsedSql = NamedParameterUtils.parseSqlStatement(insertSql);
-		SqlParameterSource paramSource = sqlParameterSources[0];
-		String sqlToUse = NamedParameterUtils.substituteNamedParameters(parsedSql, paramSource);
-		List<SqlParameter> declaredParameters = NamedParameterUtils.buildSqlParameterList(parsedSql, paramSource);
-		PreparedStatementCreatorFactory pscf = new PreparedStatementCreatorFactory(sqlToUse, declaredParameters);
-		pscf.setReturnGeneratedKeys(true);
-		pscf.setGeneratedKeysColumnNames(keyColumnNames);
-		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, sqlParameterSources[i], null);
-				pscf.newPreparedStatementSetter(values).setValues(ps);
-			}
-
-			@Override
-			public int getBatchSize() {
-				return sqlParameterSources.length;
-			}
-		};
-		PreparedStatementCallback<int[]> preparedStatementCallback = ps -> {
-			int batchSize = bpss.getBatchSize();
-			for (int i = 0; i < batchSize; i++) {
-				bpss.setValues(ps, i);
-				ps.addBatch();
-			}
-			int[] results = ps.executeBatch();
-			List<Map<String, Object>> generatedKeys = ((KeyHolder) generatedKeyHolder).getKeyList();
+	private PreparedStatementCallback<Void> storeGeneratedKeys(KeyHolder generatedKeyHolder) {
+		return ps -> {
+			List<Map<String, Object>> generatedKeys = generatedKeyHolder.getKeyList();
 			generatedKeys.clear();
 			ResultSet keys = ps.getGeneratedKeys();
 			if (keys != null) {
@@ -120,12 +103,12 @@ public int getBatchSize() {
 							new RowMapperResultSetExtractor<>(new ColumnMapRowMapper(), 1);
 					//noinspection ConstantConditions
 					generatedKeys.addAll(rse.extractData(keys));
-				} finally {
+				}
+				finally {
 					JdbcUtils.closeResultSet(keys);
 				}
 			}
-			return results;
+			return null;
 		};
-		return jdbcOperations.execute(psc, preparedStatementCallback);
 	}
 }