Merge pull request #42313 from nosan

* pr/42313:
  Polish "Use DataSource.unwrap to get routing data source"
  Use DataSource.unwrap to get routing data source

Closes gh-42313

spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/jdbc/DataSourceHealthContributorAutoConfiguration.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2023 the original author or authors.
+ * Copyright 2012-2024 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,6 +16,7 @@
 
 package org.springframework.boot.actuate.autoconfigure.jdbc;
 
+import java.sql.SQLException;
 import java.util.Collection;
 import java.util.Iterator;
 import java.util.Map;
@@ -88,7 +89,7 @@ public class DataSourceHealthContributorAutoConfiguration implements Initializin
 		if (dataSourceHealthIndicatorProperties.isIgnoreRoutingDataSources()) {
 			Map<String, DataSource> filteredDatasources = dataSources.entrySet()
 				.stream()
-				.filter((e) -> !(e.getValue() instanceof AbstractRoutingDataSource))
+				.filter((e) -> !isRoutingDataSource(e.getValue()))
 				.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
 			return createContributor(filteredDatasources);
 		}
@@ -104,8 +105,8 @@ public class DataSourceHealthContributorAutoConfiguration implements Initializin
 	}
 
 	private HealthContributor createContributor(DataSource source) {
-		if (source instanceof AbstractRoutingDataSource routingDataSource) {
+		if (isRoutingDataSource(source)) {
-			return new RoutingDataSourceHealthContributor(routingDataSource, this::createContributor);
+			return new RoutingDataSourceHealthContributor(extractRoutingDataSource(source), this::createContributor);
 		}
 		return new DataSourceHealthIndicator(source, getValidationQuery(source));
 	}
@@ -115,6 +116,30 @@ public class DataSourceHealthContributorAutoConfiguration implements Initializin
 		return (poolMetadata != null) ? poolMetadata.getValidationQuery() : null;
 	}
 
+	private static boolean isRoutingDataSource(DataSource dataSource) {
+		if (dataSource instanceof AbstractRoutingDataSource) {
+			return true;
+		}
+		try {
+			return dataSource.isWrapperFor(AbstractRoutingDataSource.class);
+		}
+		catch (SQLException ex) {
+			return false;
+		}
+	}
+
+	private static AbstractRoutingDataSource extractRoutingDataSource(DataSource dataSource) {
+		if (dataSource instanceof AbstractRoutingDataSource routingDataSource) {
+			return routingDataSource;
+		}
+		try {
+			return dataSource.unwrap(AbstractRoutingDataSource.class);
+		}
+		catch (SQLException ex) {
+			throw new IllegalStateException("Failed to unwrap AbstractRoutingDataSource from " + dataSource, ex);
+		}
+	}
+
 	/**
 	 * {@link CompositeHealthContributor} used for {@link AbstractRoutingDataSource} beans
 	 * where the overall health is composed of a {@link DataSourceHealthIndicator} for

spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/jdbc/DataSourceHealthContributorAutoConfigurationTests.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2012-2023 the original author or authors.
+ * Copyright 2012-2024 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,6 +16,7 @@
 
 package org.springframework.boot.actuate.autoconfigure.jdbc;
 
+import java.sql.SQLException;
 import java.util.HashMap;
 import java.util.Map;
 
@@ -23,6 +24,8 @@ import javax.sql.DataSource;
 
 import org.junit.jupiter.api.Test;
 
+import org.springframework.beans.BeansException;
+import org.springframework.beans.factory.config.BeanPostProcessor;
 import org.springframework.boot.actuate.autoconfigure.health.HealthContributorAutoConfiguration;
 import org.springframework.boot.actuate.autoconfigure.jdbc.DataSourceHealthContributorAutoConfiguration.RoutingDataSourceHealthContributor;
 import org.springframework.boot.actuate.health.CompositeHealthContributor;
@@ -87,6 +90,19 @@ class DataSourceHealthContributorAutoConfigurationTests {
 			});
 	}
 
+	@Test
+	void runWithProxyBeanPostProcessorRoutingAndEmbeddedDataSourceShouldIncludeRoutingDataSource() {
+		this.contextRunner
+			.withUserConfiguration(ProxyDataSourceBeanPostProcessor.class, EmbeddedDataSourceConfiguration.class,
+					RoutingDataSourceConfig.class)
+			.run((context) -> {
+				CompositeHealthContributor composite = context.getBean(CompositeHealthContributor.class);
+				assertThat(composite.getContributor("dataSource")).isInstanceOf(DataSourceHealthIndicator.class);
+				assertThat(composite.getContributor("routingDataSource"))
+					.isInstanceOf(RoutingDataSourceHealthContributor.class);
+			});
+	}
+
 	@Test
 	void runWithRoutingAndEmbeddedDataSourceShouldNotIncludeRoutingDataSourceWhenIgnored() {
 		this.contextRunner.withUserConfiguration(EmbeddedDataSourceConfiguration.class, RoutingDataSourceConfig.class)
@@ -98,6 +114,19 @@ class DataSourceHealthContributorAutoConfigurationTests {
 			});
 	}
 
+	@Test
+	void runWithProxyBeanPostProcessorAndRoutingAndEmbeddedDataSourceShouldNotIncludeRoutingDataSourceWhenIgnored() {
+		this.contextRunner
+			.withUserConfiguration(ProxyDataSourceBeanPostProcessor.class, EmbeddedDataSourceConfiguration.class,
+					RoutingDataSourceConfig.class)
+			.withPropertyValues("management.health.db.ignore-routing-datasources:true")
+			.run((context) -> {
+				assertThat(context).doesNotHaveBean(CompositeHealthContributor.class);
+				assertThat(context).hasSingleBean(DataSourceHealthIndicator.class);
+				assertThat(context).doesNotHaveBean(RoutingDataSourceHealthContributor.class);
+			});
+	}
+
 	@Test
 	void runWithOnlyRoutingDataSourceShouldIncludeRoutingDataSourceWithComposedIndicators() {
 		this.contextRunner.withUserConfiguration(RoutingDataSourceConfig.class).run((context) -> {
@@ -112,6 +141,23 @@ class DataSourceHealthContributorAutoConfigurationTests {
 		});
 	}
 
+	@Test
+	void runWithProxyBeanPostProcessorAndRoutingDataSourceShouldIncludeRoutingDataSourceWithComposedIndicators() {
+		this.contextRunner.withUserConfiguration(ProxyDataSourceBeanPostProcessor.class, RoutingDataSourceConfig.class)
+			.run((context) -> {
+				assertThat(context).hasSingleBean(RoutingDataSourceHealthContributor.class);
+				RoutingDataSourceHealthContributor routingHealthContributor = context
+					.getBean(RoutingDataSourceHealthContributor.class);
+				assertThat(routingHealthContributor.getContributor("one"))
+					.isInstanceOf(DataSourceHealthIndicator.class);
+				assertThat(routingHealthContributor.getContributor("two"))
+					.isInstanceOf(DataSourceHealthIndicator.class);
+				assertThat(routingHealthContributor.iterator()).toIterable()
+					.extracting("name")
+					.containsExactlyInAnyOrder("one", "two");
+			});
+	}
+
 	@Test
 	void runWithOnlyRoutingDataSourceShouldCrashWhenIgnored() {
 		this.contextRunner.withUserConfiguration(RoutingDataSourceConfig.class)
@@ -121,6 +167,15 @@ class DataSourceHealthContributorAutoConfigurationTests {
 				.hasRootCauseInstanceOf(IllegalArgumentException.class));
 	}
 
+	@Test
+	void runWithProxyBeanPostProcessorAndOnlyRoutingDataSourceShouldCrashWhenIgnored() {
+		this.contextRunner.withUserConfiguration(ProxyDataSourceBeanPostProcessor.class, RoutingDataSourceConfig.class)
+			.withPropertyValues("management.health.db.ignore-routing-datasources:true")
+			.run((context) -> assertThat(context).hasFailed()
+				.getFailure()
+				.hasRootCauseInstanceOf(IllegalArgumentException.class));
+	}
+
 	@Test
 	void runWithValidationQueryPropertyShouldUseCustomQuery() {
 		this.contextRunner
@@ -177,26 +232,55 @@ class DataSourceHealthContributorAutoConfigurationTests {
 	static class RoutingDataSourceConfig {
 
 		@Bean
-		AbstractRoutingDataSource routingDataSource() {
+		AbstractRoutingDataSource routingDataSource() throws SQLException {
 			Map<Object, DataSource> dataSources = new HashMap<>();
 			dataSources.put("one", mock(DataSource.class));
 			dataSources.put("two", mock(DataSource.class));
 			AbstractRoutingDataSource routingDataSource = mock(AbstractRoutingDataSource.class);
+			given(routingDataSource.isWrapperFor(AbstractRoutingDataSource.class)).willReturn(true);
+			given(routingDataSource.unwrap(AbstractRoutingDataSource.class)).willReturn(routingDataSource);
 			given(routingDataSource.getResolvedDataSources()).willReturn(dataSources);
 			return routingDataSource;
 		}
 
 	}
 
+	static class ProxyDataSourceBeanPostProcessor implements BeanPostProcessor {
+
+		@Override
+		public Object postProcessBeforeInitialization(Object bean, String beanName) throws BeansException {
+			if (bean instanceof DataSource dataSource) {
+				return proxyDataSource(dataSource);
+			}
+			return bean;
+		}
+
+		private static DataSource proxyDataSource(DataSource dataSource) {
+			try {
+				DataSource mock = mock(DataSource.class);
+				given(mock.isWrapperFor(AbstractRoutingDataSource.class))
+					.willReturn(dataSource instanceof AbstractRoutingDataSource);
+				given(mock.unwrap(AbstractRoutingDataSource.class)).willAnswer((invocation) -> dataSource);
+				return mock;
+			}
+			catch (SQLException ex) {
+				throw new IllegalStateException(ex);
+			}
+		}
+
+	}
+
 	@Configuration(proxyBeanMethods = false)
 	static class NullKeyRoutingDataSourceConfig {
 
 		@Bean
-		AbstractRoutingDataSource routingDataSource() {
+		AbstractRoutingDataSource routingDataSource() throws Exception {
 			Map<Object, DataSource> dataSources = new HashMap<>();
 			dataSources.put(null, mock(DataSource.class));
 			dataSources.put("one", mock(DataSource.class));
 			AbstractRoutingDataSource routingDataSource = mock(AbstractRoutingDataSource.class);
+			given(routingDataSource.isWrapperFor(AbstractRoutingDataSource.class)).willReturn(true);
+			given(routingDataSource.unwrap(AbstractRoutingDataSource.class)).willReturn(routingDataSource);
 			given(routingDataSource.getResolvedDataSources()).willReturn(dataSources);
 			return routingDataSource;
 		}