Merge pull request #15206 from ttddyy

* pr/15206:
  Polish "Perform best effort to retrieve DataSourceProxy"
  Perform best effort to retrieve DataSourceProxy

spring-boot-project/spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/jdbc/DataSourceJmxConfiguration.java

@@ -31,6 +31,7 @@ import org.springframework.boot.autoconfigure.condition.ConditionalOnClass;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnSingleCandidate;
+import org.springframework.boot.jdbc.DataSourceUnwrapper;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
 import org.springframework.jmx.export.MBeanExporter;
@@ -89,9 +90,11 @@ class DataSourceJmxConfiguration {
 		@Bean
 		@ConditionalOnMissingBean(name = "dataSourceMBean")
 		public Object dataSourceMBean(DataSource dataSource) {
-			if (dataSource instanceof DataSourceProxy) {
+			DataSourceProxy dataSourceProxy = DataSourceUnwrapper.unwrap(dataSource,
+					DataSourceProxy.class);
+			if (dataSourceProxy != null) {
 				try {
-					return ((DataSourceProxy) dataSource).createPool().getJmxPool();
+					return dataSourceProxy.createPool().getJmxPool();
 				}
 				catch (SQLException ex) {
 					logger.warn("Cannot expose DataSource to JMX (could not connect)");

spring-boot-project/spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/jdbc/DataSourceJmxConfigurationTests.java

@@ -38,6 +38,7 @@ import org.springframework.boot.autoconfigure.jmx.JmxAutoConfiguration;
 import org.springframework.boot.test.context.runner.ApplicationContextRunner;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.jdbc.datasource.DelegatingDataSource;
 
 import static org.assertj.core.api.Assertions.assertThat;
 
@@ -45,6 +46,7 @@ import static org.assertj.core.api.Assertions.assertThat;
  * Tests for {@link DataSourceJmxConfiguration}.
  *
  * @author Stephane Nicoll
+ * @author Tadaya Tsuyukubo
  */
 public class DataSourceJmxConfigurationTests {
 
@@ -155,6 +157,7 @@ public class DataSourceJmxConfigurationTests {
 		this.contextRunner.withPropertyValues(
 				"spring.datasource.type=" + DataSource.class.getName(),
 				"spring.datasource.jmx-enabled=true").run((context) -> {
+					assertThat(context).hasBean("dataSourceMBean");
 					assertThat(context).hasSingleBean(ConnectionPool.class);
 					assertThat(context.getBean(DataSourceProxy.class).createPool()
 							.getJmxPool())
@@ -162,6 +165,32 @@ public class DataSourceJmxConfigurationTests {
 				});
 	}
 
+	@Test
+	public void tomcatProxiedCanExposeMBeanPool() {
+		this.contextRunner.withUserConfiguration(DataSourceProxyConfiguration.class)
+				.withPropertyValues(
+						"spring.datasource.type=" + DataSource.class.getName(),
+						"spring.datasource.jmx-enabled=true")
+				.run((context) -> {
+					assertThat(context).hasBean("dataSourceMBean");
+					assertThat(context).getBean("dataSourceMBean")
+							.isInstanceOf(ConnectionPool.class);
+				});
+	}
+
+	@Test
+	public void tomcatDelegateCanExposeMBeanPool() {
+		this.contextRunner.withUserConfiguration(DataSourceDelegateConfiguration.class)
+				.withPropertyValues(
+						"spring.datasource.type=" + DataSource.class.getName(),
+						"spring.datasource.jmx-enabled=true")
+				.run((context) -> {
+					assertThat(context).hasBean("dataSourceMBean");
+					assertThat(context).getBean("dataSourceMBean")
+							.isInstanceOf(ConnectionPool.class);
+				});
+	}
+
 	@Configuration
 	static class DataSourceProxyConfiguration {
 
@@ -177,13 +206,28 @@ public class DataSourceJmxConfigurationTests {
 		@Override
 		public Object postProcessAfterInitialization(Object bean, String beanName) {
 			if (bean instanceof javax.sql.DataSource) {
-				return wrap((javax.sql.DataSource) bean);
+				return new ProxyFactory(bean).getProxy();
 			}
 			return bean;
 		}
 
-		private static javax.sql.DataSource wrap(javax.sql.DataSource dataSource) {
-			return (javax.sql.DataSource) new ProxyFactory(dataSource).getProxy();
+	}
+
+	@Configuration
+	static class DataSourceDelegateConfiguration {
+
+		@Bean
+		public static DataSourceBeanPostProcessor dataSourceBeanPostProcessor() {
+			return new DataSourceBeanPostProcessor() {
+				@Override
+				public Object postProcessAfterInitialization(Object bean,
+						String beanName) {
+					if (bean instanceof javax.sql.DataSource) {
+						return new DelegatingDataSource((javax.sql.DataSource) bean);
+					}
+					return bean;
+				}
+			};
 		}
 
 	}

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

@@ -0,0 +1,97 @@
+/*
+ * Copyright 2012-2018 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.boot.jdbc;
+
+import java.sql.Wrapper;
+
+import javax.sql.DataSource;
+
+import org.springframework.aop.framework.AopProxyUtils;
+import org.springframework.aop.support.AopUtils;
+import org.springframework.jdbc.datasource.DelegatingDataSource;
+import org.springframework.util.ClassUtils;
+
+/**
+ * Unwraps a {@link DataSource} that may have been proxied or wrapped in a custom
+ * {@link Wrapper} such as {@link DelegatingDataSource}.
+ *
+ * @author Tadaya Tsuyukubo
+ * @author Stephane Nicoll
+ * @since 2.0.7
+ */
+public final class DataSourceUnwrapper {
+
+	private static final boolean DELEGATING_DATA_SOURCE_PRESENT = ClassUtils.isPresent(
+			"org.springframework.jdbc.datasource.DelegatingDataSource",
+			DataSourceUnwrapper.class.getClassLoader());
+
+	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 target the type that the result 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) {
+		if (target.isInstance(dataSource)) {
+			return target.cast(dataSource);
+		}
+		T unwrapped = safeUnwrap(dataSource, target);
+		if (unwrapped != null) {
+			return unwrapped;
+		}
+		if (DELEGATING_DATA_SOURCE_PRESENT) {
+			DataSource targetDataSource = DelegatingDataSourceUnwrapper
+					.getTargetDataSource(dataSource);
+			if (targetDataSource != null) {
+				return unwrap(targetDataSource, target);
+			}
+		}
+		if (AopUtils.isAopProxy(dataSource)) {
+			Object proxyTarget = AopProxyUtils.getSingletonTarget(dataSource);
+			if (proxyTarget instanceof DataSource) {
+				return unwrap((DataSource) proxyTarget, target);
+			}
+		}
+		return null;
+	}
+
+	private static <S> S safeUnwrap(Wrapper wrapper, Class<S> target) {
+		try {
+			return wrapper.unwrap(target);
+		}
+		catch (Exception ex) {
+			return null;
+		}
+	}
+
+	private static class DelegatingDataSourceUnwrapper {
+
+		private static DataSource getTargetDataSource(DataSource dataSource) {
+			if (dataSource instanceof DelegatingDataSource) {
+				return ((DelegatingDataSource) dataSource).getTargetDataSource();
+			}
+			return null;
+		}
+
+	}
+
+}

spring-boot-project/spring-boot/src/test/java/org/springframework/boot/jdbc/DataSourceUnwrapperNoSpringJdbcTests.java

@@ -0,0 +1,61 @@
+/*
+ * Copyright 2012-2018 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.boot.jdbc;
+
+import javax.sql.DataSource;
+
+import com.zaxxer.hikari.HikariDataSource;
+import org.apache.tomcat.jdbc.pool.DataSourceProxy;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+import org.springframework.aop.framework.ProxyFactory;
+import org.springframework.boot.testsupport.runner.classpath.ClassPathExclusions;
+import org.springframework.boot.testsupport.runner.classpath.ModifiedClassPathRunner;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Integration tests for {@link DataSourceUnwrapper} when spring-jdbc is not available.
+ *
+ * @author Stephane Nicoll
+ */
+@RunWith(ModifiedClassPathRunner.class)
+@ClassPathExclusions("spring-jdbc-*.jar")
+public class DataSourceUnwrapperNoSpringJdbcTests {
+
+	@Test
+	public void unwrapWithProxy() {
+		DataSource dataSource = new HikariDataSource();
+		DataSource actual = wrapInProxy(wrapInProxy(dataSource));
+		assertThat(DataSourceUnwrapper.unwrap(actual, HikariDataSource.class))
+				.isSameAs(dataSource);
+	}
+
+	@Test
+	public 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);
+	}
+
+	private DataSource wrapInProxy(DataSource dataSource) {
+		return (DataSource) new ProxyFactory(dataSource).getProxy();
+	}
+
+}

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

@@ -0,0 +1,102 @@
+/*
+ * Copyright 2012-2018 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.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.springframework.boot.jdbc;
+
+import javax.sql.DataSource;
+
+import com.zaxxer.hikari.HikariDataSource;
+import org.apache.tomcat.jdbc.pool.DataSourceProxy;
+import org.junit.Test;
+
+import org.springframework.aop.framework.ProxyFactory;
+import org.springframework.jdbc.datasource.DelegatingDataSource;
+import org.springframework.jdbc.datasource.SingleConnectionDataSource;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+/**
+ * Tests for {@link DataSourceUnwrapper}.
+ *
+ * @author Stephane Nicoll
+ */
+public class DataSourceUnwrapperTests {
+
+	@Test
+	public void unwrapWithTarget() {
+		DataSource dataSource = new HikariDataSource();
+		assertThat(DataSourceUnwrapper.unwrap(dataSource, HikariDataSource.class))
+				.isSameAs(dataSource);
+	}
+
+	@Test
+	public void unwrapWithWrongTarget() {
+		DataSource dataSource = new HikariDataSource();
+		assertThat(
+				DataSourceUnwrapper.unwrap(dataSource, SingleConnectionDataSource.class))
+						.isNull();
+	}
+
+	@Test
+	public void unwrapWithDelegate() {
+		DataSource dataSource = new HikariDataSource();
+		DataSource actual = wrapInDelegate(wrapInDelegate(dataSource));
+		assertThat(DataSourceUnwrapper.unwrap(actual, HikariDataSource.class))
+				.isSameAs(dataSource);
+	}
+
+	@Test
+	public void unwrapWithProxy() {
+		DataSource dataSource = new HikariDataSource();
+		DataSource actual = wrapInProxy(wrapInProxy(dataSource));
+		assertThat(DataSourceUnwrapper.unwrap(actual, HikariDataSource.class))
+				.isSameAs(dataSource);
+	}
+
+	@Test
+	public void unwrapWithProxyAndDelegate() {
+		DataSource dataSource = new HikariDataSource();
+		DataSource actual = wrapInProxy(wrapInDelegate(dataSource));
+		assertThat(DataSourceUnwrapper.unwrap(actual, HikariDataSource.class))
+				.isSameAs(dataSource);
+	}
+
+	@Test
+	public void unwrapWithSeveralLevelOfWrapping() {
+		DataSource dataSource = new HikariDataSource();
+		DataSource actual = wrapInProxy(wrapInDelegate(
+				wrapInDelegate((wrapInProxy(wrapInDelegate(dataSource))))));
+		assertThat(DataSourceUnwrapper.unwrap(actual, HikariDataSource.class))
+				.isSameAs(dataSource);
+	}
+
+	@Test
+	public 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);
+	}
+
+	private DataSource wrapInProxy(DataSource dataSource) {
+		return (DataSource) new ProxyFactory(dataSource).getProxy();
+	}
+
+	private DataSource wrapInDelegate(DataSource dataSource) {
+		return new DelegatingDataSource(dataSource);
+	}
+
+}