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