Skip to content

Commit ae7c9ad

Browse files
nosannosan
committed
Used DataSource.unwrap(...) method alongside with 'instance of' for determining AbstractRoutingDataSource
1 parent d3a2bf4 commit ae7c9ad

File tree

2 files changed

+175
-5
lines changed

2 files changed

+175
-5
lines changed

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

Lines changed: 30 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -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) -> !isAbstractRoutingDataSource(e.getValue()))
9293
.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));
9394
return createContributor(filteredDatasources);
9495
}
@@ -104,8 +105,9 @@ 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 (isAbstractRoutingDataSource(source)) {
109+
return new RoutingDataSourceHealthContributor(unwrapAbstractRoutingDataSource(source),
110+
this::createContributor);
109111
}
110112
return new DataSourceHealthIndicator(source, getValidationQuery(source));
111113
}
@@ -115,6 +117,31 @@ private String getValidationQuery(DataSource source) {
115117
return (poolMetadata != null) ? poolMetadata.getValidationQuery() : null;
116118
}
117119

120+
private static boolean isAbstractRoutingDataSource(DataSource dataSource) {
121+
if (dataSource instanceof AbstractRoutingDataSource) {
122+
return true;
123+
}
124+
try {
125+
return dataSource.isWrapperFor(AbstractRoutingDataSource.class);
126+
}
127+
catch (SQLException ex) {
128+
return false;
129+
}
130+
}
131+
132+
private static AbstractRoutingDataSource unwrapAbstractRoutingDataSource(DataSource dataSource) {
133+
if (dataSource instanceof AbstractRoutingDataSource routingDataSource) {
134+
return routingDataSource;
135+
}
136+
try {
137+
return dataSource.unwrap(AbstractRoutingDataSource.class);
138+
}
139+
catch (SQLException ex) {
140+
throw new IllegalStateException(
141+
"DataSource '%s' failed to unwrap '%s'".formatted(dataSource, AbstractRoutingDataSource.class), ex);
142+
}
143+
}
144+
118145
/**
119146
* {@link CompositeHealthContributor} used for {@link AbstractRoutingDataSource} beans
120147
* 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: 145 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -16,13 +16,22 @@
1616

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

19+
import java.io.PrintWriter;
20+
import java.sql.Connection;
21+
import java.sql.ConnectionBuilder;
22+
import java.sql.SQLException;
23+
import java.sql.SQLFeatureNotSupportedException;
24+
import java.sql.ShardingKeyBuilder;
1925
import java.util.HashMap;
2026
import java.util.Map;
27+
import java.util.logging.Logger;
2128

2229
import javax.sql.DataSource;
2330

2431
import org.junit.jupiter.api.Test;
2532

33+
import org.springframework.beans.BeansException;
34+
import org.springframework.beans.factory.config.BeanPostProcessor;
2635
import org.springframework.boot.actuate.autoconfigure.health.HealthContributorAutoConfiguration;
2736
import org.springframework.boot.actuate.autoconfigure.jdbc.DataSourceHealthContributorAutoConfiguration.RoutingDataSourceHealthContributor;
2837
import org.springframework.boot.actuate.health.CompositeHealthContributor;
@@ -87,6 +96,19 @@ void runWithRoutingAndEmbeddedDataSourceShouldIncludeRoutingDataSource() {
8796
});
8897
}
8998

99+
@Test
100+
void runWithProxyBeanPostProcessorRoutingAndEmbeddedDataSourceShouldIncludeRoutingDataSource() {
101+
this.contextRunner
102+
.withUserConfiguration(ProxyDataSourceBeanPostProcessor.class, EmbeddedDataSourceConfiguration.class,
103+
RoutingDataSourceConfig.class)
104+
.run((context) -> {
105+
CompositeHealthContributor composite = context.getBean(CompositeHealthContributor.class);
106+
assertThat(composite.getContributor("dataSource")).isInstanceOf(DataSourceHealthIndicator.class);
107+
assertThat(composite.getContributor("routingDataSource"))
108+
.isInstanceOf(RoutingDataSourceHealthContributor.class);
109+
});
110+
}
111+
90112
@Test
91113
void runWithRoutingAndEmbeddedDataSourceShouldNotIncludeRoutingDataSourceWhenIgnored() {
92114
this.contextRunner.withUserConfiguration(EmbeddedDataSourceConfiguration.class, RoutingDataSourceConfig.class)
@@ -98,6 +120,19 @@ void runWithRoutingAndEmbeddedDataSourceShouldNotIncludeRoutingDataSourceWhenIgn
98120
});
99121
}
100122

123+
@Test
124+
void runWithProxyBeanPostProcessorAndRoutingAndEmbeddedDataSourceShouldNotIncludeRoutingDataSourceWhenIgnored() {
125+
this.contextRunner
126+
.withUserConfiguration(ProxyDataSourceBeanPostProcessor.class, EmbeddedDataSourceConfiguration.class,
127+
RoutingDataSourceConfig.class)
128+
.withPropertyValues("management.health.db.ignore-routing-datasources:true")
129+
.run((context) -> {
130+
assertThat(context).doesNotHaveBean(CompositeHealthContributor.class);
131+
assertThat(context).hasSingleBean(DataSourceHealthIndicator.class);
132+
assertThat(context).doesNotHaveBean(RoutingDataSourceHealthContributor.class);
133+
});
134+
}
135+
101136
@Test
102137
void runWithOnlyRoutingDataSourceShouldIncludeRoutingDataSourceWithComposedIndicators() {
103138
this.contextRunner.withUserConfiguration(RoutingDataSourceConfig.class).run((context) -> {
@@ -112,6 +147,23 @@ void runWithOnlyRoutingDataSourceShouldIncludeRoutingDataSourceWithComposedIndic
112147
});
113148
}
114149

150+
@Test
151+
void runWithProxyBeanPostProcessorAndRoutingDataSourceShouldIncludeRoutingDataSourceWithComposedIndicators() {
152+
this.contextRunner.withUserConfiguration(ProxyDataSourceBeanPostProcessor.class, RoutingDataSourceConfig.class)
153+
.run((context) -> {
154+
assertThat(context).hasSingleBean(RoutingDataSourceHealthContributor.class);
155+
RoutingDataSourceHealthContributor routingHealthContributor = context
156+
.getBean(RoutingDataSourceHealthContributor.class);
157+
assertThat(routingHealthContributor.getContributor("one"))
158+
.isInstanceOf(DataSourceHealthIndicator.class);
159+
assertThat(routingHealthContributor.getContributor("two"))
160+
.isInstanceOf(DataSourceHealthIndicator.class);
161+
assertThat(routingHealthContributor.iterator()).toIterable()
162+
.extracting("name")
163+
.containsExactlyInAnyOrder("one", "two");
164+
});
165+
}
166+
115167
@Test
116168
void runWithOnlyRoutingDataSourceShouldCrashWhenIgnored() {
117169
this.contextRunner.withUserConfiguration(RoutingDataSourceConfig.class)
@@ -121,6 +173,15 @@ void runWithOnlyRoutingDataSourceShouldCrashWhenIgnored() {
121173
.hasRootCauseInstanceOf(IllegalArgumentException.class));
122174
}
123175

176+
@Test
177+
void runWithProxyBeanPostProcessorAndOnlyRoutingDataSourceShouldCrashWhenIgnored() {
178+
this.contextRunner.withUserConfiguration(ProxyDataSourceBeanPostProcessor.class, RoutingDataSourceConfig.class)
179+
.withPropertyValues("management.health.db.ignore-routing-datasources:true")
180+
.run((context) -> assertThat(context).hasFailed()
181+
.getFailure()
182+
.hasRootCauseInstanceOf(IllegalArgumentException.class));
183+
}
184+
124185
@Test
125186
void runWithValidationQueryPropertyShouldUseCustomQuery() {
126187
this.contextRunner
@@ -177,30 +238,112 @@ DataSource testDataSource() {
177238
static class RoutingDataSourceConfig {
178239

179240
@Bean
180-
AbstractRoutingDataSource routingDataSource() {
241+
AbstractRoutingDataSource routingDataSource() throws SQLException {
181242
Map<Object, DataSource> dataSources = new HashMap<>();
182243
dataSources.put("one", mock(DataSource.class));
183244
dataSources.put("two", mock(DataSource.class));
184245
AbstractRoutingDataSource routingDataSource = mock(AbstractRoutingDataSource.class);
246+
given(routingDataSource.isWrapperFor(AbstractRoutingDataSource.class)).willReturn(true);
247+
given(routingDataSource.unwrap(AbstractRoutingDataSource.class)).willReturn(routingDataSource);
185248
given(routingDataSource.getResolvedDataSources()).willReturn(dataSources);
186249
return routingDataSource;
187250
}
188251

189252
}
190253

254+
static class ProxyDataSourceBeanPostProcessor implements BeanPostProcessor {
255+
256+
@Override
257+
public Object postProcessBeforeInitialization(Object bean, String beanName) throws BeansException {
258+
if (bean instanceof DataSource dataSource) {
259+
return new ProxyDataSource(dataSource);
260+
}
261+
return bean;
262+
}
263+
264+
}
265+
191266
@Configuration(proxyBeanMethods = false)
192267
static class NullKeyRoutingDataSourceConfig {
193268

194269
@Bean
195-
AbstractRoutingDataSource routingDataSource() {
270+
AbstractRoutingDataSource routingDataSource() throws Exception {
196271
Map<Object, DataSource> dataSources = new HashMap<>();
197272
dataSources.put(null, mock(DataSource.class));
198273
dataSources.put("one", mock(DataSource.class));
199274
AbstractRoutingDataSource routingDataSource = mock(AbstractRoutingDataSource.class);
275+
given(routingDataSource.isWrapperFor(AbstractRoutingDataSource.class)).willReturn(true);
276+
given(routingDataSource.unwrap(AbstractRoutingDataSource.class)).willReturn(routingDataSource);
200277
given(routingDataSource.getResolvedDataSources()).willReturn(dataSources);
201278
return routingDataSource;
202279
}
203280

204281
}
205282

283+
static class ProxyDataSource implements DataSource {
284+
285+
private final DataSource dataSource;
286+
287+
ProxyDataSource(DataSource dataSource) {
288+
this.dataSource = dataSource;
289+
}
290+
291+
@Override
292+
public void setLogWriter(PrintWriter out) throws SQLException {
293+
this.dataSource.setLogWriter(out);
294+
}
295+
296+
@Override
297+
public Connection getConnection() throws SQLException {
298+
return this.dataSource.getConnection();
299+
}
300+
301+
@Override
302+
public Connection getConnection(String username, String password) throws SQLException {
303+
return this.dataSource.getConnection(username, password);
304+
}
305+
306+
@Override
307+
public PrintWriter getLogWriter() throws SQLException {
308+
return this.dataSource.getLogWriter();
309+
}
310+
311+
@Override
312+
public void setLoginTimeout(int seconds) throws SQLException {
313+
this.dataSource.setLoginTimeout(seconds);
314+
}
315+
316+
@Override
317+
public int getLoginTimeout() throws SQLException {
318+
return this.dataSource.getLoginTimeout();
319+
}
320+
321+
@Override
322+
public ConnectionBuilder createConnectionBuilder() throws SQLException {
323+
return this.dataSource.createConnectionBuilder();
324+
}
325+
326+
@Override
327+
public Logger getParentLogger() throws SQLFeatureNotSupportedException {
328+
return this.dataSource.getParentLogger();
329+
}
330+
331+
@Override
332+
public ShardingKeyBuilder createShardingKeyBuilder() throws SQLException {
333+
return this.dataSource.createShardingKeyBuilder();
334+
}
335+
336+
@Override
337+
@SuppressWarnings("unchecked")
338+
public <T> T unwrap(Class<T> iface) throws SQLException {
339+
return iface.isInstance(this) ? (T) this : this.dataSource.unwrap(iface);
340+
}
341+
342+
@Override
343+
public boolean isWrapperFor(Class<?> iface) throws SQLException {
344+
return (iface.isInstance(this) || this.dataSource.isWrapperFor(iface));
345+
}
346+
347+
}
348+
206349
}

0 commit comments

Comments
 (0)