diff --git a/spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/jdbc/DataSourceHealthContributorAutoConfiguration.java b/spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/jdbc/DataSourceHealthContributorAutoConfiguration.java index 966296cd00f..a94172671e9 100644 --- a/spring-boot-project/spring-boot-actuator-autoconfigure/src/main/java/org/springframework/boot/actuate/autoconfigure/jdbc/DataSourceHealthContributorAutoConfiguration.java +++ b/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 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) { - return new RoutingDataSourceHealthContributor(routingDataSource, this::createContributor); + if (isRoutingDataSource(source)) { + 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 diff --git a/spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/jdbc/DataSourceHealthContributorAutoConfigurationTests.java b/spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/jdbc/DataSourceHealthContributorAutoConfigurationTests.java index 6945cb71bbb..bcbf8dc752b 100644 --- a/spring-boot-project/spring-boot-actuator-autoconfigure/src/test/java/org/springframework/boot/actuate/autoconfigure/jdbc/DataSourceHealthContributorAutoConfigurationTests.java +++ b/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 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 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; }