Support user-defined key type in JDBC KeyHolder

Prior to this commit, the JDBC KeyHolder API only supported keys of
type Number. However, a generated key can be a UUID or something else,
and developers shouldn't have to go manually through complex
collections to access it.

This commit adds a new getKeyAs(Class<T> keyType) method to the
KeyHolder API that allows the user to specify the key type.

Closes gh-24655

spring-jdbc/src/main/java/org/springframework/jdbc/support/GeneratedKeyHolder.java

@@ -61,6 +61,12 @@ public class GeneratedKeyHolder implements KeyHolder {
 	@Override
 	@Nullable
 	public Number getKey() throws InvalidDataAccessApiUsageException, DataRetrievalFailureException {
+		return getKeyAs(Number.class);
+	}
+
+	@Override
+	@Nullable
+	public <T> T getKeyAs(Class<T> keyClass) throws InvalidDataAccessApiUsageException, DataRetrievalFailureException {
 		if (this.keyList.isEmpty()) {
 			return null;
 		}
@@ -72,13 +78,13 @@ public class GeneratedKeyHolder implements KeyHolder {
 		Iterator<Object> keyIter = this.keyList.get(0).values().iterator();
 		if (keyIter.hasNext()) {
 			Object key = keyIter.next();
-			if (!(key instanceof Number)) {
+			if (key == null || !(keyClass.isAssignableFrom(key.getClass()))) {
 				throw new DataRetrievalFailureException(
-						"The generated key is not of a supported numeric type. " +
+						"The generated key is not of a supported type. " +
 						"Unable to cast [" + (key != null ? key.getClass().getName() : null) +
-						"] to [" + Number.class.getName() + "]");
+						"] to [" + keyClass.getName() + "]");
 			}
-			return (Number) key;
+			return keyClass.cast(key);
 		}
 		else {
 			throw new DataRetrievalFailureException("Unable to retrieve the generated key. " +

spring-jdbc/src/main/java/org/springframework/jdbc/support/KeyHolder.java

@@ -58,6 +58,22 @@ public interface KeyHolder {
 	@Nullable
 	Number getKey() throws InvalidDataAccessApiUsageException;
 
+	/**
+	 * Retrieve the first item from the first map, assuming that there is just
+	 * one item and just one map, and that the item is an instance of provided class.
+	 * This is the typical case: a single generated key of desired class.
+	 * <p>Keys are held in a List of Maps, where each item in the list represents
+	 * the keys for each row. If there are multiple columns, then the Map will have
+	 * multiple entries as well. If this method encounters multiple entries in
+	 * either the map or the list meaning that multiple keys were returned,
+	 * then an InvalidDataAccessApiUsageException is thrown.
+	 * @param keyClass class of the requested key
+	 * @return the generated key as an instance of provided class
+	 * @throws InvalidDataAccessApiUsageException if multiple keys are encountered
+	 */
+	@Nullable
+	<T> T getKeyAs(Class<T> keyClass) throws InvalidDataAccessApiUsageException;
+
 	/**
 	 * Retrieve the first map of keys.
 	 * <p>If there are multiple entries in the list (meaning that multiple rows

spring-jdbc/src/test/java/org/springframework/jdbc/support/KeyHolderTests.java

@@ -51,13 +51,38 @@ public class KeyHolderTests {
 		assertThat(kh.getKey().intValue()).as("single key should be returned").isEqualTo(1);
 	}
 
+	@Test
+	public void singleKeyAsString() {
+		kh.getKeyList().addAll(singletonList(singletonMap("key", "1")));
+
+		assertThat(kh.getKeyAs(String.class)).as("single key should be returned").isEqualTo("1");
+	}
+
+	@Test
+	public void singleKeyAsWrongClass() {
+		kh.getKeyList().addAll(singletonList(singletonMap("key", "1")));
+
+		assertThatExceptionOfType(DataRetrievalFailureException.class).isThrownBy(() ->
+				kh.getKeyAs(Integer.class))
+				.withMessageStartingWith("The generated key is not of a supported type.");
+	}
+
+	@Test
+	public void singleKeyWithNullValue() {
+		kh.getKeyList().addAll(singletonList(singletonMap("key", null)));
+
+		assertThatExceptionOfType(DataRetrievalFailureException.class).isThrownBy(() ->
+				kh.getKeyAs(Integer.class))
+				.withMessageStartingWith("The generated key is not of a supported type.");
+	}
+
 	@Test
 	public void singleKeyNonNumeric() {
 		kh.getKeyList().addAll(singletonList(singletonMap("key", "1")));
 
 		assertThatExceptionOfType(DataRetrievalFailureException.class).isThrownBy(() ->
 				kh.getKey().intValue())
-			.withMessageStartingWith("The generated key is not of a supported numeric type.");
+			.withMessageStartingWith("The generated key is not of a supported type.");
 	}
 
 	@Test