Polish "Handle lazy connection proxy datasource in DataSourceBuilder"

See gh-50271

Author: snicoll

module/spring-boot-jdbc/src/main/java/org/springframework/boot/jdbc/DataSourceBuilder.java

@@ -41,7 +41,6 @@
 
 import org.springframework.beans.BeanUtils;
 import org.springframework.core.ResolvableType;
-import org.springframework.jdbc.datasource.LazyConnectionDataSourceProxy;
 import org.springframework.jdbc.datasource.SimpleDriverDataSource;
 import org.springframework.lang.Contract;
 import org.springframework.util.Assert;
@@ -238,27 +237,7 @@ public static DataSourceBuilder<?> create(@Nullable ClassLoader classLoader) {
 	 * @return a new {@link DataSource} builder
 	 */
 	public static DataSourceBuilder<?> derivedFrom(DataSource dataSource) {
-		return new DataSourceBuilder<>(unwrap(dataSource));
-	}
-
-	private static DataSource unwrap(DataSource dataSource) {
-		if ((dataSource instanceof LazyConnectionDataSourceProxy dataSourceProxy)
-				&& (dataSourceProxy.getTargetDataSource() != null)) {
-			dataSource = dataSourceProxy.getTargetDataSource();
-		}
-		try {
-			while (dataSource.isWrapperFor(DataSource.class)) {
-				DataSource unwrapped = dataSource.unwrap(DataSource.class);
-				if (unwrapped == dataSource) {
-					return unwrapped;
-				}
-				dataSource = unwrapped;
-			}
-		}
-		catch (SQLException ex) {
-			// Try to continue with the existing, potentially still wrapped, DataSource
-		}
-		return dataSource;
+		return new DataSourceBuilder<>(DataSourceUnwrapper.unwrapRoot(dataSource));
 	}
 
 	/**

module/spring-boot-jdbc/src/main/java/org/springframework/boot/jdbc/DataSourceUnwrapper.java

@@ -93,6 +93,37 @@ private DataSourceUnwrapper() {
 		return unwrap(dataSource, target, target);
 	}
 
+	/**
+	 * Return the root {@link DataSource} by recursively unwrapping all
+	 * {@link org.springframework.jdbc.datasource.DelegatingDataSource delegating}, proxy,
+	 * and {@link java.sql.Wrapper} layers until no further unwrapping is possible.
+	 * @param dataSource the datasource to unwrap
+	 * @return the root {@link DataSource}
+	 * @since 4.1.0
+	 */
+	public static DataSource unwrapRoot(DataSource dataSource) {
+		if (DELEGATING_DATA_SOURCE_PRESENT) {
+			DataSource targetDataSource = DelegatingDataSourceUnwrapper.getTargetDataSource(dataSource);
+			if (targetDataSource != null) {
+				return unwrapRoot(targetDataSource);
+			}
+		}
+		if (AopUtils.isAopProxy(dataSource)) {
+			Object proxyTarget = AopProxyUtils.getSingletonTarget(dataSource);
+			if (proxyTarget instanceof DataSource proxyDataSource) {
+				return unwrapRoot(proxyDataSource);
+			}
+		}
+		DataSource unwrapped = safeUnwrap(dataSource);
+		if (unwrapped != null) {
+			if (unwrapped == dataSource) {
+				return unwrapped;
+			}
+			return unwrapRoot(unwrapped);
+		}
+		return dataSource;
+	}
+
 	private static <S> @Nullable S safeUnwrap(Wrapper wrapper, Class<S> target) {
 		try {
 			if (target.isInterface() && wrapper.isWrapperFor(target)) {
@@ -105,6 +136,18 @@ private DataSourceUnwrapper() {
 		return null;
 	}
 
+	private static @Nullable DataSource safeUnwrap(DataSource dataSource) {
+		try {
+			if (dataSource.isWrapperFor(DataSource.class)) {
+				return dataSource.unwrap(DataSource.class);
+			}
+		}
+		catch (Exception ex) {
+			// continue
+		}
+		return null;
+	}
+
 	private static final class DelegatingDataSourceUnwrapper {
 
 		private static @Nullable DataSource getTargetDataSource(DataSource dataSource) {

module/spring-boot-jdbc/src/test/java/org/springframework/boot/jdbc/DataSourceUnwrapperTests.java

@@ -106,6 +106,40 @@ void unwrapDataSourceProxy() {
 			.isSameAs(dataSource);
 	}
 
+	@Test
+	void unwrapRootWithPlainDataSource() {
+		DataSource dataSource = new HikariDataSource();
+		assertThat(DataSourceUnwrapper.unwrapRoot(dataSource)).isSameAs(dataSource);
+	}
+
+	@Test
+	void unwrapRootWithDelegate() {
+		DataSource dataSource = new HikariDataSource();
+		DataSource actual = wrapInDelegate(wrapInDelegate(dataSource));
+		assertThat(DataSourceUnwrapper.unwrapRoot(actual)).isSameAs(dataSource);
+	}
+
+	@Test
+	void unwrapRootWithProxy() {
+		DataSource dataSource = new HikariDataSource();
+		DataSource actual = wrapInProxy(wrapInProxy(dataSource));
+		assertThat(DataSourceUnwrapper.unwrapRoot(actual)).isSameAs(dataSource);
+	}
+
+	@Test
+	void unwrapRootWithLazyConnectionDataSource() {
+		DataSource dataSource = new HikariDataSource();
+		DataSource actual = new LazyConnectionDataSourceProxy(dataSource);
+		assertThat(DataSourceUnwrapper.unwrapRoot(actual)).isSameAs(dataSource);
+	}
+
+	@Test
+	void unwrapRootWithSeveralLevelOfWrapping() {
+		DataSource dataSource = new HikariDataSource();
+		DataSource actual = wrapInProxy(wrapInDelegate(wrapInDelegate(wrapInProxy(wrapInDelegate(dataSource)))));
+		assertThat(DataSourceUnwrapper.unwrapRoot(actual)).isSameAs(dataSource);
+	}
+
 	@Test
 	void unwrappingIsNotAttemptedWhenTargetIsNotAnInterface() {
 		DataSource dataSource = mock(DataSource.class);