Skip to content

Commit ece9cb3

Browse files
snicollsnicoll
committed
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
2 parents 99d0805 + 78a140a commit ece9cb3

File tree

2 files changed

+116
-7
lines changed

2 files changed

+116
-7
lines changed

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

Lines changed: 29 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2012-2023 the original author or authors.
2+
* Copyright 2012-2024 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -16,6 +16,7 @@
1616

1717
package org.springframework.boot.actuate.autoconfigure.jdbc;
1818

19+
import java.sql.SQLException;
1920
import java.util.Collection;
2021
import java.util.Iterator;
2122
import java.util.Map;
@@ -88,7 +89,7 @@ public HealthContributor dbHealthContributor(Map<String, DataSource> dataSources
8889
if (dataSourceHealthIndicatorProperties.isIgnoreRoutingDataSources()) {
8990
Map<String, DataSource> filteredDatasources = dataSources.entrySet()
9091
.stream()
91-
.filter((e) -> !(e.getValue() instanceof AbstractRoutingDataSource))
92+
.filter((e) -> !isRoutingDataSource(e.getValue()))
9293
.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
9394
return createContributor(filteredDatasources);
9495
}
@@ -104,8 +105,8 @@ private HealthContributor createContributor(Map<String, DataSource> beans) {
104105
}
105106

106107
private HealthContributor createContributor(DataSource source) {
107-
if (source instanceof AbstractRoutingDataSource routingDataSource) {
108-
return new RoutingDataSourceHealthContributor(routingDataSource, this::createContributor);
108+
if (isRoutingDataSource(source)) {
109+
return new RoutingDataSourceHealthContributor(extractRoutingDataSource(source), this::createContributor);
109110
}
110111
return new DataSourceHealthIndicator(source, getValidationQuery(source));
111112
}
@@ -115,6 +116,30 @@ private String getValidationQuery(DataSource source) {
115116
return (poolMetadata != null) ? poolMetadata.getValidationQuery() : null;
116117
}
117118

119+
private static boolean isRoutingDataSource(DataSource dataSource) {
120+
if (dataSource instanceof AbstractRoutingDataSource) {
121+
return true;
122+
}
123+
try {
124+
return dataSource.isWrapperFor(AbstractRoutingDataSource.class);
125+
}
126+
catch (SQLException ex) {
127+
return false;
128+
}
129+
}
130+
131+
private static AbstractRoutingDataSource extractRoutingDataSource(DataSource dataSource) {
132+
if (dataSource instanceof AbstractRoutingDataSource routingDataSource) {
133+
return routingDataSource;
134+
}
135+
try {
136+
return dataSource.unwrap(AbstractRoutingDataSource.class);
137+
}
138+
catch (SQLException ex) {
139+
throw new IllegalStateException("Failed to unwrap AbstractRoutingDataSource from " + dataSource, ex);
140+
}
141+
}
142+
118143
/**
119144
* {@link CompositeHealthContributor} used for {@link AbstractRoutingDataSource} beans
120145
* 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‎

Lines changed: 87 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2012-2023 the original author or authors.
2+
* Copyright 2012-2024 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -16,13 +16,16 @@
1616

1717
package org.springframework.boot.actuate.autoconfigure.jdbc;
1818

19+
import java.sql.SQLException;
1920
import java.util.HashMap;
2021
import java.util.Map;
2122

2223
import javax.sql.DataSource;
2324

2425
import org.junit.jupiter.api.Test;
2526

27+
import org.springframework.beans.BeansException;
28+
import org.springframework.beans.factory.config.BeanPostProcessor;
2629
import org.springframework.boot.actuate.autoconfigure.health.HealthContributorAutoConfiguration;
2730
import org.springframework.boot.actuate.autoconfigure.jdbc.DataSourceHealthContributorAutoConfiguration.RoutingDataSourceHealthContributor;
2831
import org.springframework.boot.actuate.health.CompositeHealthContributor;
@@ -87,6 +90,19 @@ void runWithRoutingAndEmbeddedDataSourceShouldIncludeRoutingDataSource() {
8790
});
8891
}
8992

93+
@Test
94+
void runWithProxyBeanPostProcessorRoutingAndEmbeddedDataSourceShouldIncludeRoutingDataSource() {
95+
this.contextRunner
96+
.withUserConfiguration(ProxyDataSourceBeanPostProcessor.class, EmbeddedDataSourceConfiguration.class,
97+
RoutingDataSourceConfig.class)
98+
.run((context) -> {
99+
CompositeHealthContributor composite = context.getBean(CompositeHealthContributor.class);
100+
assertThat(composite.getContributor("dataSource")).isInstanceOf(DataSourceHealthIndicator.class);
101+
assertThat(composite.getContributor("routingDataSource"))
102+
.isInstanceOf(RoutingDataSourceHealthContributor.class);
103+
});
104+
}
105+
90106
@Test
91107
void runWithRoutingAndEmbeddedDataSourceShouldNotIncludeRoutingDataSourceWhenIgnored() {
92108
this.contextRunner.withUserConfiguration(EmbeddedDataSourceConfiguration.class, RoutingDataSourceConfig.class)
@@ -98,6 +114,19 @@ void runWithRoutingAndEmbeddedDataSourceShouldNotIncludeRoutingDataSourceWhenIgn
98114
});
99115
}
100116

117+
@Test
118+
void runWithProxyBeanPostProcessorAndRoutingAndEmbeddedDataSourceShouldNotIncludeRoutingDataSourceWhenIgnored() {
119+
this.contextRunner
120+
.withUserConfiguration(ProxyDataSourceBeanPostProcessor.class, EmbeddedDataSourceConfiguration.class,
121+
RoutingDataSourceConfig.class)
122+
.withPropertyValues("management.health.db.ignore-routing-datasources:true")
123+
.run((context) -> {
124+
assertThat(context).doesNotHaveBean(CompositeHealthContributor.class);
125+
assertThat(context).hasSingleBean(DataSourceHealthIndicator.class);
126+
assertThat(context).doesNotHaveBean(RoutingDataSourceHealthContributor.class);
127+
});
128+
}
129+
101130
@Test
102131
void runWithOnlyRoutingDataSourceShouldIncludeRoutingDataSourceWithComposedIndicators() {
103132
this.contextRunner.withUserConfiguration(RoutingDataSourceConfig.class).run((context) -> {
@@ -112,6 +141,23 @@ void runWithOnlyRoutingDataSourceShouldIncludeRoutingDataSourceWithComposedIndic
112141
});
113142
}
114143

144+
@Test
145+
void runWithProxyBeanPostProcessorAndRoutingDataSourceShouldIncludeRoutingDataSourceWithComposedIndicators() {
146+
this.contextRunner.withUserConfiguration(ProxyDataSourceBeanPostProcessor.class, RoutingDataSourceConfig.class)
147+
.run((context) -> {
148+
assertThat(context).hasSingleBean(RoutingDataSourceHealthContributor.class);
149+
RoutingDataSourceHealthContributor routingHealthContributor = context
150+
.getBean(RoutingDataSourceHealthContributor.class);
151+
assertThat(routingHealthContributor.getContributor("one"))
152+
.isInstanceOf(DataSourceHealthIndicator.class);
153+
assertThat(routingHealthContributor.getContributor("two"))
154+
.isInstanceOf(DataSourceHealthIndicator.class);
155+
assertThat(routingHealthContributor.iterator()).toIterable()
156+
.extracting("name")
157+
.containsExactlyInAnyOrder("one", "two");
158+
});
159+
}
160+
115161
@Test
116162
void runWithOnlyRoutingDataSourceShouldCrashWhenIgnored() {
117163
this.contextRunner.withUserConfiguration(RoutingDataSourceConfig.class)
@@ -121,6 +167,15 @@ void runWithOnlyRoutingDataSourceShouldCrashWhenIgnored() {
121167
.hasRootCauseInstanceOf(IllegalArgumentException.class));
122168
}
123169

170+
@Test
171+
void runWithProxyBeanPostProcessorAndOnlyRoutingDataSourceShouldCrashWhenIgnored() {
172+
this.contextRunner.withUserConfiguration(ProxyDataSourceBeanPostProcessor.class, RoutingDataSourceConfig.class)
173+
.withPropertyValues("management.health.db.ignore-routing-datasources:true")
174+
.run((context) -> assertThat(context).hasFailed()
175+
.getFailure()
176+
.hasRootCauseInstanceOf(IllegalArgumentException.class));
177+
}
178+
124179
@Test
125180
void runWithValidationQueryPropertyShouldUseCustomQuery() {
126181
this.contextRunner
@@ -177,26 +232,55 @@ DataSource testDataSource() {
177232
static class RoutingDataSourceConfig {
178233

179234
@Bean
180-
AbstractRoutingDataSource routingDataSource() {
235+
AbstractRoutingDataSource routingDataSource() throws SQLException {
181236
Map<Object, DataSource> dataSources = new HashMap<>();
182237
dataSources.put("one", mock(DataSource.class));
183238
dataSources.put("two", mock(DataSource.class));
184239
AbstractRoutingDataSource routingDataSource = mock(AbstractRoutingDataSource.class);
240+
given(routingDataSource.isWrapperFor(AbstractRoutingDataSource.class)).willReturn(true);
241+
given(routingDataSource.unwrap(AbstractRoutingDataSource.class)).willReturn(routingDataSource);
185242
given(routingDataSource.getResolvedDataSources()).willReturn(dataSources);
186243
return routingDataSource;
187244
}
188245

189246
}
190247

248+
static class ProxyDataSourceBeanPostProcessor implements BeanPostProcessor {
249+
250+
@Override
251+
public Object postProcessBeforeInitialization(Object bean, String beanName) throws BeansException {
252+
if (bean instanceof DataSource dataSource) {
253+
return proxyDataSource(dataSource);
254+
}
255+
return bean;
256+
}
257+
258+
private static DataSource proxyDataSource(DataSource dataSource) {
259+
try {
260+
DataSource mock = mock(DataSource.class);
261+
given(mock.isWrapperFor(AbstractRoutingDataSource.class))
262+
.willReturn(dataSource instanceof AbstractRoutingDataSource);
263+
given(mock.unwrap(AbstractRoutingDataSource.class)).willAnswer((invocation) -> dataSource);
264+
return mock;
265+
}
266+
catch (SQLException ex) {
267+
throw new IllegalStateException(ex);
268+
}
269+
}
270+
271+
}
272+
191273
@Configuration(proxyBeanMethods = false)
192274
static class NullKeyRoutingDataSourceConfig {
193275

194276
@Bean
195-
AbstractRoutingDataSource routingDataSource() {
277+
AbstractRoutingDataSource routingDataSource() throws Exception {
196278
Map<Object, DataSource> dataSources = new HashMap<>();
197279
dataSources.put(null, mock(DataSource.class));
198280
dataSources.put("one", mock(DataSource.class));
199281
AbstractRoutingDataSource routingDataSource = mock(AbstractRoutingDataSource.class);
282+
given(routingDataSource.isWrapperFor(AbstractRoutingDataSource.class)).willReturn(true);
283+
given(routingDataSource.unwrap(AbstractRoutingDataSource.class)).willReturn(routingDataSource);
200284
given(routingDataSource.getResolvedDataSources()).willReturn(dataSources);
201285
return routingDataSource;
202286
}

0 commit comments

Comments
 (0)