Unwrap Datasource against an actual interface

This commit updates DataSourceUnwrapper to take a separate interface type argument if the target datasource has to be unwrapped, given that the target type is usually not an interface. Closes spring-projectsgh-24697
snicoll · Jan 14, 2021 · dacf158 · dacf158
1 parent 3ecd100
commit dacf158
Show file tree

Hide file tree

Showing 6 changed files with 75 additions and 30 deletions.
diff --git a/...ework/boot/actuate/autoconfigure/metrics/jdbc/DataSourcePoolMetricsAutoConfiguration.java b/...ework/boot/actuate/autoconfigure/metrics/jdbc/DataSourcePoolMetricsAutoConfiguration.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2019 the original author or authors.
+ * Copyright 2012-2021 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.
@@ -24,6 +24,7 @@
 
 import javax.sql.DataSource;
 
+import com.zaxxer.hikari.HikariConfigMXBean;
 import com.zaxxer.hikari.HikariDataSource;
 import com.zaxxer.hikari.metrics.micrometer.MicrometerMetricsTrackerFactory;
 import io.micrometer.core.instrument.MeterRegistry;
@@ -112,7 +113,8 @@ static class HikariDataSourceMetricsConfiguration {
 		@Autowired
 		void bindMetricsRegistryToHikariDataSources(Collection<DataSource> dataSources) {
 			for (DataSource dataSource : dataSources) {
-				HikariDataSource hikariDataSource = DataSourceUnwrapper.unwrap(dataSource, HikariDataSource.class);
+				HikariDataSource hikariDataSource = DataSourceUnwrapper.unwrap(dataSource, HikariConfigMXBean.class,
+						HikariDataSource.class);
 				if (hikariDataSource != null) {
 					bindMetricsRegistryToHikariDataSource(hikariDataSource);
 				}

diff --git a/...src/main/java/org/springframework/boot/autoconfigure/jdbc/DataSourceJmxConfiguration.java b/...src/main/java/org/springframework/boot/autoconfigure/jdbc/DataSourceJmxConfiguration.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2020 the original author or authors.
+ * Copyright 2012-2021 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.
@@ -20,10 +20,12 @@
 
 import javax.sql.DataSource;
 
+import com.zaxxer.hikari.HikariConfigMXBean;
 import com.zaxxer.hikari.HikariDataSource;
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 import org.apache.tomcat.jdbc.pool.DataSourceProxy;
+import org.apache.tomcat.jdbc.pool.PoolConfiguration;
 
 import org.springframework.beans.factory.ObjectProvider;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
@@ -62,7 +64,8 @@ static class Hikari {
 		}
 
 		private void validateMBeans() {
-			HikariDataSource hikariDataSource = DataSourceUnwrapper.unwrap(this.dataSource, HikariDataSource.class);
+			HikariDataSource hikariDataSource = DataSourceUnwrapper.unwrap(this.dataSource, HikariConfigMXBean.class,
+					HikariDataSource.class);
 			if (hikariDataSource != null && hikariDataSource.isRegisterMbeans()) {
 				this.mBeanExporter.ifUnique((exporter) -> exporter.addExcludedBean("dataSource"));
 			}
@@ -79,7 +82,8 @@ static class TomcatDataSourceJmxConfiguration {
 		@Bean
 		@ConditionalOnMissingBean(name = "dataSourceMBean")
 		Object dataSourceMBean(DataSource dataSource) {
-			DataSourceProxy dataSourceProxy = DataSourceUnwrapper.unwrap(dataSource, DataSourceProxy.class);
+			DataSourceProxy dataSourceProxy = DataSourceUnwrapper.unwrap(dataSource, PoolConfiguration.class,
+					DataSourceProxy.class);
 			if (dataSourceProxy != null) {
 				try {
 					return dataSourceProxy.createPool().getJmxPool();

diff --git a/...mework/boot/autoconfigure/jdbc/metadata/DataSourcePoolMetadataProvidersConfiguration.java b/...mework/boot/autoconfigure/jdbc/metadata/DataSourcePoolMetadataProvidersConfiguration.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2019 the original author or authors.
+ * Copyright 2012-2021 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.
@@ -16,8 +16,11 @@
 
 package org.springframework.boot.autoconfigure.jdbc.metadata;
 
+import com.zaxxer.hikari.HikariConfigMXBean;
 import com.zaxxer.hikari.HikariDataSource;
 import org.apache.commons.dbcp2.BasicDataSource;
+import org.apache.commons.dbcp2.BasicDataSourceMXBean;
+import org.apache.tomcat.jdbc.pool.jmx.ConnectionPoolMBean;
 
 import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.jdbc.DataSourceUnwrapper;
@@ -46,7 +49,7 @@ static class TomcatDataSourcePoolMetadataProviderConfiguration {
 		DataSourcePoolMetadataProvider tomcatPoolDataSourceMetadataProvider() {
 			return (dataSource) -> {
 				org.apache.tomcat.jdbc.pool.DataSource tomcatDataSource = DataSourceUnwrapper.unwrap(dataSource,
-						org.apache.tomcat.jdbc.pool.DataSource.class);
+						ConnectionPoolMBean.class, org.apache.tomcat.jdbc.pool.DataSource.class);
 				if (tomcatDataSource != null) {
 					return new TomcatDataSourcePoolMetadata(tomcatDataSource);
 				}
@@ -63,7 +66,8 @@ static class HikariPoolDataSourceMetadataProviderConfiguration {
 		@Bean
 		DataSourcePoolMetadataProvider hikariPoolDataSourceMetadataProvider() {
 			return (dataSource) -> {
-				HikariDataSource hikariDataSource = DataSourceUnwrapper.unwrap(dataSource, HikariDataSource.class);
+				HikariDataSource hikariDataSource = DataSourceUnwrapper.unwrap(dataSource, HikariConfigMXBean.class,
+						HikariDataSource.class);
 				if (hikariDataSource != null) {
 					return new HikariDataSourcePoolMetadata(hikariDataSource);
 				}
@@ -80,7 +84,8 @@ static class CommonsDbcp2PoolDataSourceMetadataProviderConfiguration {
 		@Bean
 		DataSourcePoolMetadataProvider commonsDbcp2PoolDataSourceMetadataProvider() {
 			return (dataSource) -> {
-				BasicDataSource dbcpDataSource = DataSourceUnwrapper.unwrap(dataSource, BasicDataSource.class);
+				BasicDataSource dbcpDataSource = DataSourceUnwrapper.unwrap(dataSource, BasicDataSourceMXBean.class,
+						BasicDataSource.class);
 				if (dbcpDataSource != null) {
 					return new CommonsDbcp2DataSourcePoolMetadata(dbcpDataSource);
 				}

diff --git a/...-project/spring-boot/src/main/java/org/springframework/boot/jdbc/DataSourceUnwrapper.java b/...-project/spring-boot/src/main/java/org/springframework/boot/jdbc/DataSourceUnwrapper.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2020 the original author or authors.
+ * Copyright 2012-2021 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.
@@ -45,33 +45,49 @@ private DataSourceUnwrapper() {
 	 * Return an object that implements the given {@code target} type, unwrapping delegate
 	 * or proxy if necessary.
 	 * @param dataSource the datasource to handle
+	 * @param unwrapInterface the interface that the target type must implement
 	 * @param target the type that the result must implement
+	 * @param <I> the interface that the target type must implement
 	 * @param <T> the target type
 	 * @return an object that implements the target type or {@code null}
 	 */
-	public static <T> T unwrap(DataSource dataSource, Class<T> target) {
+	public static <I, T extends I> T unwrap(DataSource dataSource, Class<I> unwrapInterface, Class<T> target) {
 		if (target.isInstance(dataSource)) {
 			return target.cast(dataSource);
 		}
-		T unwrapped = safeUnwrap(dataSource, target);
-		if (unwrapped != null) {
-			return unwrapped;
+		I unwrapped = safeUnwrap(dataSource, unwrapInterface);
+		if (unwrapped != null && unwrapInterface.isAssignableFrom(target)) {
+			return target.cast(unwrapped);
 		}
 		if (DELEGATING_DATA_SOURCE_PRESENT) {
 			DataSource targetDataSource = DelegatingDataSourceUnwrapper.getTargetDataSource(dataSource);
 			if (targetDataSource != null) {
-				return unwrap(targetDataSource, target);
+				return unwrap(targetDataSource, unwrapInterface, target);
 			}
 		}
 		if (AopUtils.isAopProxy(dataSource)) {
 			Object proxyTarget = AopProxyUtils.getSingletonTarget(dataSource);
 			if (proxyTarget instanceof DataSource) {
-				return unwrap((DataSource) proxyTarget, target);
+				return unwrap((DataSource) proxyTarget, unwrapInterface, target);
 			}
 		}
 		return null;
 	}
 
+	/**
+	 * Return an object that implements the given {@code target} type, unwrapping delegate
+	 * or proxy if necessary.
+	 * @param dataSource the datasource to handle
+	 * @param target the type that the result must implement
+	 * @param <T> the target type
+	 * @return an object that implements the target type or {@code null}
+	 * @deprecated as of 2.3.8 in favour of {@link #unwrap(DataSource, Class, Class)}
+	 */
+	@Deprecated
+	public static <T> T unwrap(DataSource dataSource, Class<T> target) {
+		return unwrap(dataSource, target, target);
+	}
+
 	private static <S> S safeUnwrap(Wrapper wrapper, Class<S> target) {
 		try {
 			if (target.isInterface() && wrapper.isWrapperFor(target)) {
@@ -84,6 +100,10 @@ private static <S> S safeUnwrap(Wrapper wrapper, Class<S> target) {
 		return null;
 	}
 
+	private static <S> S safeCast(Object target, Class<S> type) {
+		return (type.isInstance(target)) ? type.cast(target) : null;
+	}
+
 	private static class DelegatingDataSourceUnwrapper {
 
 		private static DataSource getTargetDataSource(DataSource dataSource) {

diff --git a/...oot/src/test/java/org/springframework/boot/jdbc/DataSourceUnwrapperNoSpringJdbcTests.java b/...oot/src/test/java/org/springframework/boot/jdbc/DataSourceUnwrapperNoSpringJdbcTests.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2019 the original author or authors.
+ * Copyright 2012-2021 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.
@@ -18,8 +18,10 @@
 
 import javax.sql.DataSource;
 
+import com.zaxxer.hikari.HikariConfigMXBean;
 import com.zaxxer.hikari.HikariDataSource;
 import org.apache.tomcat.jdbc.pool.DataSourceProxy;
+import org.apache.tomcat.jdbc.pool.PoolConfiguration;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.aop.framework.ProxyFactory;
@@ -39,14 +41,16 @@ class DataSourceUnwrapperNoSpringJdbcTests {
 	void unwrapWithProxy() {
 		DataSource dataSource = new HikariDataSource();
 		DataSource actual = wrapInProxy(wrapInProxy(dataSource));
-		assertThat(DataSourceUnwrapper.unwrap(actual, HikariDataSource.class)).isSameAs(dataSource);
+		assertThat(DataSourceUnwrapper.unwrap(actual, HikariConfigMXBean.class, HikariDataSource.class))
+				.isSameAs(dataSource);
 	}
 
 	@Test
 	void unwrapDataSourceProxy() {
 		org.apache.tomcat.jdbc.pool.DataSource dataSource = new org.apache.tomcat.jdbc.pool.DataSource();
 		DataSource actual = wrapInProxy(wrapInProxy(dataSource));
-		assertThat(DataSourceUnwrapper.unwrap(actual, DataSourceProxy.class)).isSameAs(dataSource);
+		assertThat(DataSourceUnwrapper.unwrap(actual, PoolConfiguration.class, DataSourceProxy.class))
+				.isSameAs(dataSource);
 	}
 
 	private DataSource wrapInProxy(DataSource dataSource) {

diff --git a/...ect/spring-boot/src/test/java/org/springframework/boot/jdbc/DataSourceUnwrapperTests.java b/...ect/spring-boot/src/test/java/org/springframework/boot/jdbc/DataSourceUnwrapperTests.java
@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2020 the original author or authors.
+ * Copyright 2012-2021 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.
@@ -21,13 +21,16 @@
 
 import javax.sql.DataSource;
 
+import com.zaxxer.hikari.HikariConfigMXBean;
 import com.zaxxer.hikari.HikariDataSource;
 import org.apache.tomcat.jdbc.pool.DataSourceProxy;
+import org.apache.tomcat.jdbc.pool.PoolConfiguration;
 import org.junit.jupiter.api.Test;
 
 import org.springframework.aop.framework.ProxyFactory;
 import org.springframework.jdbc.datasource.DelegatingDataSource;
 import org.springframework.jdbc.datasource.SingleConnectionDataSource;
+import org.springframework.jdbc.datasource.SmartDataSource;
 
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.mockito.Mockito.mock;
@@ -44,61 +47,68 @@ class DataSourceUnwrapperTests {
 	@Test
 	void unwrapWithTarget() {
 		DataSource dataSource = new HikariDataSource();
-		assertThat(DataSourceUnwrapper.unwrap(dataSource, HikariDataSource.class)).isSameAs(dataSource);
+		assertThat(DataSourceUnwrapper.unwrap(dataSource, HikariConfigMXBean.class, HikariDataSource.class))
+				.isSameAs(dataSource);
 	}
 
 	@Test
 	void unwrapWithWrongTarget() {
 		DataSource dataSource = new HikariDataSource();
-		assertThat(DataSourceUnwrapper.unwrap(dataSource, SingleConnectionDataSource.class)).isNull();
+		assertThat(DataSourceUnwrapper.unwrap(dataSource, SmartDataSource.class, SingleConnectionDataSource.class))
+				.isNull();
 	}
 
 	@Test
 	void unwrapWithDelegate() {
 		DataSource dataSource = new HikariDataSource();
 		DataSource actual = wrapInDelegate(wrapInDelegate(dataSource));
-		assertThat(DataSourceUnwrapper.unwrap(actual, HikariDataSource.class)).isSameAs(dataSource);
+		assertThat(DataSourceUnwrapper.unwrap(actual, HikariConfigMXBean.class, HikariDataSource.class))
+				.isSameAs(dataSource);
 	}
 
 	@Test
 	void unwrapWithProxy() {
 		DataSource dataSource = new HikariDataSource();
 		DataSource actual = wrapInProxy(wrapInProxy(dataSource));
-		assertThat(DataSourceUnwrapper.unwrap(actual, HikariDataSource.class)).isSameAs(dataSource);
+		assertThat(DataSourceUnwrapper.unwrap(actual, HikariConfigMXBean.class, HikariDataSource.class))
+				.isSameAs(dataSource);
 	}
 
 	@Test
 	void unwrapWithProxyAndDelegate() {
 		DataSource dataSource = new HikariDataSource();
 		DataSource actual = wrapInProxy(wrapInDelegate(dataSource));
-		assertThat(DataSourceUnwrapper.unwrap(actual, HikariDataSource.class)).isSameAs(dataSource);
+		assertThat(DataSourceUnwrapper.unwrap(actual, HikariConfigMXBean.class, HikariDataSource.class))
+				.isSameAs(dataSource);
 	}
 
 	@Test
 	void unwrapWithSeveralLevelOfWrapping() {
 		DataSource dataSource = new HikariDataSource();
 		DataSource actual = wrapInProxy(wrapInDelegate(wrapInDelegate(wrapInProxy(wrapInDelegate(dataSource)))));
-		assertThat(DataSourceUnwrapper.unwrap(actual, HikariDataSource.class)).isSameAs(dataSource);
+		assertThat(DataSourceUnwrapper.unwrap(actual, HikariConfigMXBean.class, HikariDataSource.class))
+				.isSameAs(dataSource);
 	}
 
 	@Test
 	void unwrapDataSourceProxy() {
 		org.apache.tomcat.jdbc.pool.DataSource dataSource = new org.apache.tomcat.jdbc.pool.DataSource();
 		DataSource actual = wrapInDelegate(wrapInProxy(dataSource));
-		assertThat(DataSourceUnwrapper.unwrap(actual, DataSourceProxy.class)).isSameAs(dataSource);
+		assertThat(DataSourceUnwrapper.unwrap(actual, PoolConfiguration.class, DataSourceProxy.class))
+				.isSameAs(dataSource);
 	}
 
 	@Test
-	void unwrappingIsNotAttemptedWhenTargetIsNotAnInterface() throws SQLException {
+	void unwrappingIsNotAttemptedWhenTargetIsNotAnInterface() {
 		DataSource dataSource = mock(DataSource.class);
-		assertThat(DataSourceUnwrapper.unwrap(dataSource, HikariDataSource.class)).isNull();
+		assertThat(DataSourceUnwrapper.unwrap(dataSource, HikariDataSource.class, HikariDataSource.class)).isNull();
 		verifyNoMoreInteractions(dataSource);
 	}
 
 	@Test
 	void unwrappingIsNotAttemptedWhenDataSourceIsNotWrapperForTarget() throws SQLException {
 		DataSource dataSource = mock(DataSource.class);
-		assertThat(DataSourceUnwrapper.unwrap(dataSource, Consumer.class)).isNull();
+		assertThat(DataSourceUnwrapper.unwrap(dataSource, Consumer.class, Consumer.class)).isNull();
 		verify(dataSource).isWrapperFor(Consumer.class);
 		verifyNoMoreInteractions(dataSource);
 	}