commit
						04891746ff
					
				|  | @ -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) { | ||||
| 			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 | ||||
|  |  | |||
|  | @ -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; | ||||
| 		} | ||||
|  |  | |||
		Loading…
	
		Reference in New Issue