Polish "Perform best effort to retrieve DataSourceProxy"

Closes spring-projectsgh-15206

Author: snicoll

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

@@ -26,23 +26,20 @@
 import org.apache.commons.logging.LogFactory;
 import org.apache.tomcat.jdbc.pool.DataSourceProxy;
 
-import org.springframework.aop.framework.AopProxyUtils;
-import org.springframework.aop.support.AopUtils;
 import org.springframework.beans.factory.ObjectProvider;
 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.jdbc.datasource.DelegatingDataSource;
 import org.springframework.jmx.export.MBeanExporter;
 
 /**
  * Configures DataSource related MBeans.
  *
  * @author Stephane Nicoll
- * @author Tadaya Tsuyukubo
  */
 @Configuration
 @ConditionalOnProperty(prefix = "spring.jmx", name = "enabled", havingValue = "true", matchIfMissing = true)
@@ -93,7 +90,8 @@ static class TomcatDataSourceJmxConfiguration {
 		@Bean
 		@ConditionalOnMissingBean(name = "dataSourceMBean")
 		public Object dataSourceMBean(DataSource dataSource) {
-			DataSourceProxy dataSourceProxy = extractDataSourceProxy(dataSource);
+			DataSourceProxy dataSourceProxy = DataSourceUnwrapper.unwrap(dataSource,
+					DataSourceProxy.class);
 			if (dataSourceProxy != null) {
 				try {
 					return dataSourceProxy.createPool().getJmxPool();
@@ -105,36 +103,6 @@ public Object dataSourceMBean(DataSource dataSource) {
 			return null;
 		}
 
-		/**
-		 * Perform best effort to retrieve tomcat's {@link DataSourceProxy}.
-		 *
-		 * Since {@link DataSourceProxy#unwrap(Class)} always return {@code null}, it
-		 * cannot directly retrieve {@link DataSourceProxy}. This method tries best effort
-		 * to find {@link DataSourceProxy} if the given {@link DataSource} is wrapped or
-		 * proxied by spring.
-		 * @param dataSource candidate datasource
-		 * @return found DataSourceProxy or null
-		 */
-		private DataSourceProxy extractDataSourceProxy(DataSource dataSource) {
-			if (dataSource instanceof DataSourceProxy) {
-				return (DataSourceProxy) dataSource; // found
-			}
-			else if (dataSource instanceof DelegatingDataSource) {
-				// check delegating target
-				return extractDataSourceProxy(
-						((DelegatingDataSource) dataSource).getTargetDataSource());
-			}
-			else if (AopUtils.isAopProxy(dataSource)) {
-				// for proxy by spring, try target(advised) instance
-				Object target = AopProxyUtils.getSingletonTarget(dataSource);
-				if (target instanceof DataSource) {
-					return extractDataSourceProxy((DataSource) target);
-				}
-			}
-
-			return null;
-		}
-
 	}
 
 }

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

@@ -31,9 +31,7 @@
 import org.apache.tomcat.jdbc.pool.jmx.ConnectionPool;
 import org.junit.Test;
 
-import org.springframework.aop.framework.AopProxyUtils;
 import org.springframework.aop.framework.ProxyFactory;
-import org.springframework.aop.support.AopUtils;
 import org.springframework.beans.factory.config.BeanPostProcessor;
 import org.springframework.boot.autoconfigure.AutoConfigurations;
 import org.springframework.boot.autoconfigure.jmx.JmxAutoConfiguration;
@@ -159,6 +157,7 @@ public void tomcatAutoConfiguredCanExposeMBeanPool() {
 		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())
@@ -173,12 +172,9 @@ public void tomcatProxiedCanExposeMBeanPool() {
 						"spring.datasource.type=" + DataSource.class.getName(),
 						"spring.datasource.jmx-enabled=true")
 				.run((context) -> {
-					assertThat(context).hasSingleBean(ConnectionPool.class);
-					DataSourceProxy dataSourceProxy = (DataSourceProxy) AopProxyUtils
-							.getSingletonTarget(
-									context.getBean(javax.sql.DataSource.class));
-					assertThat(dataSourceProxy.createPool().getJmxPool())
-							.isSameAs(context.getBean(ConnectionPool.class));
+					assertThat(context).hasBean("dataSourceMBean");
+					assertThat(context).getBean("dataSourceMBean")
+							.isInstanceOf(ConnectionPool.class);
 				});
 	}
 
@@ -189,55 +185,12 @@ public void tomcatDelegateCanExposeMBeanPool() {
 						"spring.datasource.type=" + DataSource.class.getName(),
 						"spring.datasource.jmx-enabled=true")
 				.run((context) -> {
-					assertThat(context).hasSingleBean(ConnectionPool.class);
-					DataSourceProxy dataSourceProxy = (DataSourceProxy) context
-							.getBean(DelegatingDataSource.class).getTargetDataSource();
-					assertThat(dataSourceProxy.createPool().getJmxPool())
-							.isSameAs(context.getBean(ConnectionPool.class));
-				});
-	}
-
-	@Test
-	public void tomcatProxyAndDelegateCanExposeMBeanPool() {
-		this.contextRunner
-				.withUserConfiguration(DataSourceMixWrapAndProxyConfiguration.class)
-				.withPropertyValues(
-						"spring.datasource.type=" + DataSource.class.getName(),
-						"spring.datasource.jmx-enabled=true")
-				.run((context) -> {
-					assertThat(context).hasSingleBean(ConnectionPool.class);
-					DataSourceProxy dataSourceProxy = extractTomcatDataSource(
-							context.getBean(javax.sql.DataSource.class));
-					assertThat(dataSourceProxy.createPool().getJmxPool())
-							.isSameAs(context.getBean(ConnectionPool.class));
+					assertThat(context).hasBean("dataSourceMBean");
+					assertThat(context).getBean("dataSourceMBean")
+							.isInstanceOf(ConnectionPool.class);
 				});
 	}
 
-	private static javax.sql.DataSource wrap(javax.sql.DataSource dataSource) {
-		return (javax.sql.DataSource) new ProxyFactory(dataSource).getProxy();
-	}
-
-	private static javax.sql.DataSource delegate(javax.sql.DataSource dataSource) {
-		return new DelegatingDataSource(dataSource);
-	}
-
-	private static DataSource extractTomcatDataSource(javax.sql.DataSource dataSource) {
-		if (dataSource instanceof DataSource) {
-			return (DataSource) dataSource;
-		}
-		else if (dataSource instanceof DelegatingDataSource) {
-			return extractTomcatDataSource(
-					((DelegatingDataSource) dataSource).getTargetDataSource());
-		}
-		else if (AopUtils.isAopProxy(dataSource)) {
-			return extractTomcatDataSource(
-					(javax.sql.DataSource) AopProxyUtils.getSingletonTarget(dataSource));
-		}
-
-		throw new RuntimeException(
-				"Not proxied or delegated tomcat DataSource: " + dataSource);
-	}
-
 	@Configuration
 	static class DataSourceProxyConfiguration {
 
@@ -253,7 +206,7 @@ private static class DataSourceBeanPostProcessor implements BeanPostProcessor {
 		@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;
 		}
@@ -279,34 +232,4 @@ public Object postProcessAfterInitialization(Object bean,
 
 	}
 
-	@Configuration
-	static class DataSourceMixWrapAndProxyConfiguration {
-
-		@Bean
-		public static DataSourceBeanPostProcessor dataSourceBeanPostProcessor() {
-			return new DataSourceBeanPostProcessor() {
-				@Override
-				public Object postProcessAfterInitialization(Object bean,
-						String beanName) {
-					if (bean instanceof javax.sql.DataSource) {
-						javax.sql.DataSource dataSource = (javax.sql.DataSource) bean;
-						// delegate/wrap multiple times
-						for (int i = 0; i < 10; i++) {
-							if (i % 2 == 0) {
-								dataSource = wrap(dataSource);
-							}
-							else {
-								dataSource = delegate(dataSource);
-							}
-						}
-
-						return dataSource;
-					}
-					return bean;
-				}
-			};
-		}
-
-	}
-
 }

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

@@ -0,0 +1,55 @@
+/*
+ * 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 org.springframework.aop.framework.AopProxyUtils;
+import org.springframework.aop.support.AopUtils;
+import org.springframework.jdbc.datasource.DelegatingDataSource;
+
+/**
+ * Unwraps a {@link DataSource} that may have been proxied or wrapped in a
+ * {@link DelegatingDataSource}.
+ *
+ * @author Tadaya Tsuyukubo
+ * @author Stephane Nicoll
+ * @since 2.0.7
+ */
+public final class DataSourceUnwrapper {
+
+	private DataSourceUnwrapper() {
+	}
+
+	public static <T> T unwrap(DataSource dataSource, java.lang.Class<T> target) {
+		if (target.isInstance(dataSource)) {
+			return target.cast(dataSource);
+		}
+		if (dataSource instanceof DelegatingDataSource) {
+			return unwrap(((DelegatingDataSource) dataSource).getTargetDataSource(),
+					target);
+		}
+		else if (AopUtils.isAopProxy(dataSource)) {
+			Object proxyTarget = AopProxyUtils.getSingletonTarget(dataSource);
+			if (proxyTarget instanceof DataSource) {
+				return unwrap((DataSource) proxyTarget, target);
+			}
+		}
+		return null;
+	}
+
+}

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

@@ -0,0 +1,84 @@
+/*
+ * 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.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);
+	}
+
+	private DataSource wrapInProxy(DataSource dataSource) {
+		return (DataSource) new ProxyFactory(dataSource).getProxy();
+	}
+
+	private DataSource wrapInDelegate(DataSource dataSource) {
+		return new DelegatingDataSource(dataSource);
+	}
+
+}