Provide health for an AbstractRoutingDataSource's resolved targets

See gh-25708

Author: onobc

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

@@ -17,19 +17,19 @@
 package org.springframework.boot.actuate.autoconfigure.jdbc;
 
 import java.util.Collection;
+import java.util.Iterator;
 import java.util.Map;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 
 import javax.sql.DataSource;
 
 import org.springframework.beans.factory.InitializingBean;
 import org.springframework.beans.factory.ObjectProvider;
-import org.springframework.boot.actuate.autoconfigure.health.CompositeHealthContributorConfiguration;
 import org.springframework.boot.actuate.autoconfigure.health.ConditionalOnEnabledHealthIndicator;
-import org.springframework.boot.actuate.health.AbstractHealthIndicator;
-import org.springframework.boot.actuate.health.Health.Builder;
+import org.springframework.boot.actuate.health.CompositeHealthContributor;
 import org.springframework.boot.actuate.health.HealthContributor;
-import org.springframework.boot.actuate.health.HealthIndicator;
+import org.springframework.boot.actuate.health.NamedContributor;
 import org.springframework.boot.actuate.jdbc.DataSourceHealthIndicator;
 import org.springframework.boot.autoconfigure.AutoConfigureAfter;
 import org.springframework.boot.autoconfigure.EnableAutoConfiguration;
@@ -45,6 +45,7 @@
 import org.springframework.context.annotation.Configuration;
 import org.springframework.jdbc.core.JdbcTemplate;
 import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
+import org.springframework.util.Assert;
 
 /**
  * {@link EnableAutoConfiguration Auto-configuration} for
@@ -64,8 +65,7 @@
 @ConditionalOnEnabledHealthIndicator("db")
 @AutoConfigureAfter(DataSourceAutoConfiguration.class)
 @EnableConfigurationProperties(DataSourceHealthIndicatorProperties.class)
-public class DataSourceHealthContributorAutoConfiguration extends
-		CompositeHealthContributorConfiguration<AbstractHealthIndicator, DataSource> implements InitializingBean {
+public class DataSourceHealthContributorAutoConfiguration implements InitializingBean {
 
 	private final Collection<DataSourcePoolMetadataProvider> metadataProviders;
 
@@ -94,10 +94,18 @@ public HealthContributor dbHealthContributor(Map<String, DataSource> dataSources
 		return createContributor(dataSources);
 	}
 
-	@Override
-	protected AbstractHealthIndicator createIndicator(DataSource source) {
+	private HealthContributor createContributor(Map<String, DataSource> beans) {
+		Assert.notEmpty(beans, "Beans must not be empty");
+		if (beans.size() == 1) {
+			return createIndicator(beans.values().iterator().next());
+		}
+		return CompositeHealthContributor.fromMap(beans, this::createIndicator);
+	}
+
+	private HealthContributor createIndicator(DataSource source) {
 		if (source instanceof AbstractRoutingDataSource) {
-			return new RoutingDataSourceHealthIndicator();
+			AbstractRoutingDataSource routingDataSource = (AbstractRoutingDataSource) source;
+			return new RoutingDataSourceHealthIndicator(routingDataSource, this::createIndicator);
 		}
 		return new DataSourceHealthIndicator(source, getValidationQuery(source));
 	}
@@ -108,14 +116,29 @@ private String getValidationQuery(DataSource source) {
 	}
 
 	/**
-	 * {@link HealthIndicator} used for {@link AbstractRoutingDataSource} beans where we
-	 * can't actually query for the status.
+	 * {@link CompositeHealthContributor} used for {@link AbstractRoutingDataSource} beans
+	 * where the overall health is composed of a {@link DataSourceHealthIndicator} for
+	 * each routed datasource.
 	 */
-	static class RoutingDataSourceHealthIndicator extends AbstractHealthIndicator {
+	static class RoutingDataSourceHealthIndicator implements CompositeHealthContributor {
+
+		private CompositeHealthContributor delegate;
+
+		RoutingDataSourceHealthIndicator(AbstractRoutingDataSource routingDataSource,
+				Function<DataSource, HealthContributor> indicatorFunction) {
+			Map<String, DataSource> routedDataSources = routingDataSource.getResolvedDataSources().entrySet().stream()
+					.collect(Collectors.toMap((e) -> e.getKey().toString(), Map.Entry::getValue));
+			this.delegate = CompositeHealthContributor.fromMap(routedDataSources, indicatorFunction);
+		}
+
+		@Override
+		public HealthContributor getContributor(String name) {
+			return this.delegate.getContributor(name);
+		}
 
 		@Override
-		protected void doHealthCheck(Builder builder) throws Exception {
-			builder.unknown().withDetail("routing", true);
+		public Iterator<NamedContributor<HealthContributor>> iterator() {
+			return this.delegate.iterator();
 		}
 
 	}

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

@@ -16,6 +16,9 @@
 
 package org.springframework.boot.actuate.autoconfigure.jdbc;
 
+import java.util.HashMap;
+import java.util.Map;
+
 import javax.sql.DataSource;
 
 import org.junit.jupiter.api.Test;
@@ -38,6 +41,7 @@
 import org.springframework.jdbc.datasource.lookup.AbstractRoutingDataSource;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.BDDMockito.given;
 import static org.mockito.Mockito.mock;
 
 /**
@@ -93,9 +97,16 @@ void runWithRoutingAndEmbeddedDataSourceShouldNotIncludeRoutingDataSourceWhenIgn
 	}
 
 	@Test
-	void runWithOnlyRoutingDataSourceShouldIncludeRoutingDataSource() {
-		this.contextRunner.withUserConfiguration(RoutingDataSourceConfig.class)
-				.run((context) -> assertThat(context).hasSingleBean(RoutingDataSourceHealthIndicator.class));
+	void runWithOnlyRoutingDataSourceShouldIncludeRoutingDataSourceWithComposedIndicators() {
+		this.contextRunner.withUserConfiguration(RoutingDataSourceConfig.class).run((context) -> {
+			assertThat(context).hasSingleBean(RoutingDataSourceHealthIndicator.class);
+			RoutingDataSourceHealthIndicator routingHealthContributor = context
+					.getBean(RoutingDataSourceHealthIndicator.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
@@ -143,7 +154,12 @@ static class RoutingDataSourceConfig {
 
 		@Bean
 		AbstractRoutingDataSource routingDataSource() {
-			return mock(AbstractRoutingDataSource.class);
+			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.getResolvedDataSources()).willReturn(dataSources);
+			return routingDataSource;
 		}
 
 	}