Merge pull request #51410 from jjh75607

Closes gh-51410

* fix/unwrap-root-honours-wrapper-unwrap-41x:
  Polish "Honor a wrapper's unwrap when resolving the root DataSource"
  Honor a wrapper's unwrap when resolving the root DataSource

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

@@ -94,14 +94,19 @@ public final class DataSourceUnwrapper {
 	}
 
 	/**
-	 * 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.
+	 * Return the root {@link DataSource} by recursively unwrapping
+	 * {@link java.sql.Wrapper},
+	 * {@link org.springframework.jdbc.datasource.DelegatingDataSource delegating}, and
+	 * proxy 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) {
+		DataSource unwrapped = safeUnwrap(dataSource);
+		if (unwrapped != null && unwrapped != dataSource) {
+			return unwrapRoot(unwrapped);
+		}
 		if (DELEGATING_DATA_SOURCE_PRESENT) {
 			DataSource targetDataSource = DelegatingDataSourceUnwrapper.getTargetDataSource(dataSource);
 			if (targetDataSource != null) {
@@ -114,13 +119,6 @@ public final class DataSourceUnwrapper {
 				return unwrapRoot(proxyDataSource);
 			}
 		}
-		DataSource unwrapped = safeUnwrap(dataSource);
-		if (unwrapped != null) {
-			if (unwrapped == dataSource) {
-				return unwrapped;
-			}
-			return unwrapRoot(unwrapped);
-		}
 		return dataSource;
 	}
 

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

@@ -16,6 +16,7 @@
 
 package org.springframework.boot.jdbc;
 
+import java.sql.Connection;
 import java.sql.SQLException;
 import java.util.function.Consumer;
 
@@ -28,6 +29,7 @@ import org.apache.tomcat.jdbc.pool.PoolConfiguration;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.aop.framework.ProxyFactory;
+import org.springframework.jdbc.datasource.AbstractDataSource;
 import org.springframework.jdbc.datasource.DelegatingDataSource;
 import org.springframework.jdbc.datasource.LazyConnectionDataSourceProxy;
 import org.springframework.jdbc.datasource.SingleConnectionDataSource;
@@ -140,6 +142,19 @@ class DataSourceUnwrapperTests {
 		assertThat(DataSourceUnwrapper.unwrapRoot(actual)).isSameAs(dataSource);
 	}
 
+	@Test
+	void unwrapRootWithSelfReturningWrapper() {
+		DataSource dataSource = new SelfReturningDataSource();
+		assertThat(DataSourceUnwrapper.unwrapRoot(dataSource)).isSameAs(dataSource);
+	}
+
+	@Test
+	void unwrapRootWithDelegateWrappingSelfReturningWrapper() {
+		DataSource dataSource = new HikariDataSource();
+		DataSource actual = new RootAwareDelegatingDataSource(new SelfReturningDataSource(), dataSource);
+		assertThat(DataSourceUnwrapper.unwrapRoot(actual)).isSameAs(dataSource);
+	}
+
 	@Test
 	void unwrappingIsNotAttemptedWhenTargetIsNotAnInterface() {
 		DataSource dataSource = mock(DataSource.class);
@@ -163,4 +178,34 @@ class DataSourceUnwrapperTests {
 		return new DelegatingDataSource(dataSource);
 	}
 
+	private static final class SelfReturningDataSource extends AbstractDataSource {
+
+		@Override
+		public Connection getConnection() {
+			throw new UnsupportedOperationException();
+		}
+
+		@Override
+		public Connection getConnection(String username, String password) {
+			throw new UnsupportedOperationException();
+		}
+
+	}
+
+	private static final class RootAwareDelegatingDataSource extends DelegatingDataSource {
+
+		private final DataSource root;
+
+		private RootAwareDelegatingDataSource(DataSource target, DataSource root) {
+			super(target);
+			this.root = root;
+		}
+
+		@Override
+		public <T> T unwrap(Class<T> iface) throws SQLException {
+			return (iface.isInstance(this.root)) ? iface.cast(this.root) : super.unwrap(iface);
+		}
+
+	}
+
 }