From 3f9f0490a63b19456a614a435c171995549e597d Mon Sep 17 00:00:00 2001 From: Dmytro Nosan Date: Mon, 9 Sep 2024 12:43:19 +0300 Subject: [PATCH 1/2] Use DataSource.unwrap to get routing data source This commit uses DataSource.isWrapperFor and DataSource.unwrap to detect if a DataSource is an AbstractRoutingDataSource. Previously, it relied on instanceof which does not account for cases where the datasource has been proxied. See gh-42313 --- ...rceHealthContributorAutoConfiguration.java | 35 +++- ...althContributorAutoConfigurationTests.java | 149 +++++++++++++++++- 2 files changed, 177 insertions(+), 7 deletions(-) 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..02bd94f88a5 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) -> !isAbstractRoutingDataSource(e.getValue())) .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); return createContributor(filteredDatasources); } @@ -104,8 +105,9 @@ public class DataSourceHealthContributorAutoConfiguration implements Initializin } private HealthContributor createContributor(DataSource source) { - if (source instanceof AbstractRoutingDataSource routingDataSource) { - return new RoutingDataSourceHealthContributor(routingDataSource, this::createContributor); + if (isAbstractRoutingDataSource(source)) { + return new RoutingDataSourceHealthContributor(unwrapAbstractRoutingDataSource(source), + this::createContributor); } return new DataSourceHealthIndicator(source, getValidationQuery(source)); } @@ -115,6 +117,31 @@ public class DataSourceHealthContributorAutoConfiguration implements Initializin return (poolMetadata != null) ? poolMetadata.getValidationQuery() : null; } + private static boolean isAbstractRoutingDataSource(DataSource dataSource) { + if (dataSource instanceof AbstractRoutingDataSource) { + return true; + } + try { + return dataSource.isWrapperFor(AbstractRoutingDataSource.class); + } + catch (SQLException ex) { + return false; + } + } + + private static AbstractRoutingDataSource unwrapAbstractRoutingDataSource(DataSource dataSource) { + if (dataSource instanceof AbstractRoutingDataSource routingDataSource) { + return routingDataSource; + } + try { + return dataSource.unwrap(AbstractRoutingDataSource.class); + } + catch (SQLException ex) { + throw new IllegalStateException( + "DataSource '%s' failed to unwrap '%s'".formatted(dataSource, AbstractRoutingDataSource.class), 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..3badcb91fdf 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,13 +16,22 @@ package org.springframework.boot.actuate.autoconfigure.jdbc; +import java.io.PrintWriter; +import java.sql.Connection; +import java.sql.ConnectionBuilder; +import java.sql.SQLException; +import java.sql.SQLFeatureNotSupportedException; +import java.sql.ShardingKeyBuilder; import java.util.HashMap; import java.util.Map; +import java.util.logging.Logger; 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 +96,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 +120,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 +147,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 +173,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,30 +238,112 @@ 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 new ProxyDataSource(dataSource); + } + return bean; + } + + } + @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; } } + static class ProxyDataSource implements DataSource { + + private final DataSource dataSource; + + ProxyDataSource(DataSource dataSource) { + this.dataSource = dataSource; + } + + @Override + public void setLogWriter(PrintWriter out) throws SQLException { + this.dataSource.setLogWriter(out); + } + + @Override + public Connection getConnection() throws SQLException { + return this.dataSource.getConnection(); + } + + @Override + public Connection getConnection(String username, String password) throws SQLException { + return this.dataSource.getConnection(username, password); + } + + @Override + public PrintWriter getLogWriter() throws SQLException { + return this.dataSource.getLogWriter(); + } + + @Override + public void setLoginTimeout(int seconds) throws SQLException { + this.dataSource.setLoginTimeout(seconds); + } + + @Override + public int getLoginTimeout() throws SQLException { + return this.dataSource.getLoginTimeout(); + } + + @Override + public ConnectionBuilder createConnectionBuilder() throws SQLException { + return this.dataSource.createConnectionBuilder(); + } + + @Override + public Logger getParentLogger() throws SQLFeatureNotSupportedException { + return this.dataSource.getParentLogger(); + } + + @Override + public ShardingKeyBuilder createShardingKeyBuilder() throws SQLException { + return this.dataSource.createShardingKeyBuilder(); + } + + @Override + @SuppressWarnings("unchecked") + public T unwrap(Class iface) throws SQLException { + return iface.isInstance(this) ? (T) this : this.dataSource.unwrap(iface); + } + + @Override + public boolean isWrapperFor(Class iface) throws SQLException { + return (iface.isInstance(this) || this.dataSource.isWrapperFor(iface)); + } + + } + } From 78a140ae258a2d2e499a2c1965ad45b83295d067 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?St=C3=A9phane=20Nicoll?= Date: Mon, 16 Sep 2024 09:29:04 +0200 Subject: [PATCH 2/2] Polish "Use DataSource.unwrap to get routing data source" See gh-42313 --- ...rceHealthContributorAutoConfiguration.java | 14 ++- ...althContributorAutoConfigurationTests.java | 87 +++---------------- 2 files changed, 20 insertions(+), 81 deletions(-) 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 02bd94f88a5..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 @@ -89,7 +89,7 @@ public class DataSourceHealthContributorAutoConfiguration implements Initializin if (dataSourceHealthIndicatorProperties.isIgnoreRoutingDataSources()) { Map filteredDatasources = dataSources.entrySet() .stream() - .filter((e) -> !isAbstractRoutingDataSource(e.getValue())) + .filter((e) -> !isRoutingDataSource(e.getValue())) .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)); return createContributor(filteredDatasources); } @@ -105,9 +105,8 @@ public class DataSourceHealthContributorAutoConfiguration implements Initializin } private HealthContributor createContributor(DataSource source) { - if (isAbstractRoutingDataSource(source)) { - return new RoutingDataSourceHealthContributor(unwrapAbstractRoutingDataSource(source), - this::createContributor); + if (isRoutingDataSource(source)) { + return new RoutingDataSourceHealthContributor(extractRoutingDataSource(source), this::createContributor); } return new DataSourceHealthIndicator(source, getValidationQuery(source)); } @@ -117,7 +116,7 @@ public class DataSourceHealthContributorAutoConfiguration implements Initializin return (poolMetadata != null) ? poolMetadata.getValidationQuery() : null; } - private static boolean isAbstractRoutingDataSource(DataSource dataSource) { + private static boolean isRoutingDataSource(DataSource dataSource) { if (dataSource instanceof AbstractRoutingDataSource) { return true; } @@ -129,7 +128,7 @@ public class DataSourceHealthContributorAutoConfiguration implements Initializin } } - private static AbstractRoutingDataSource unwrapAbstractRoutingDataSource(DataSource dataSource) { + private static AbstractRoutingDataSource extractRoutingDataSource(DataSource dataSource) { if (dataSource instanceof AbstractRoutingDataSource routingDataSource) { return routingDataSource; } @@ -137,8 +136,7 @@ public class DataSourceHealthContributorAutoConfiguration implements Initializin return dataSource.unwrap(AbstractRoutingDataSource.class); } catch (SQLException ex) { - throw new IllegalStateException( - "DataSource '%s' failed to unwrap '%s'".formatted(dataSource, AbstractRoutingDataSource.class), ex); + throw new IllegalStateException("Failed to unwrap AbstractRoutingDataSource from " + dataSource, ex); } } 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 3badcb91fdf..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 @@ -16,15 +16,9 @@ package org.springframework.boot.actuate.autoconfigure.jdbc; -import java.io.PrintWriter; -import java.sql.Connection; -import java.sql.ConnectionBuilder; import java.sql.SQLException; -import java.sql.SQLFeatureNotSupportedException; -import java.sql.ShardingKeyBuilder; import java.util.HashMap; import java.util.Map; -import java.util.logging.Logger; import javax.sql.DataSource; @@ -256,11 +250,24 @@ class DataSourceHealthContributorAutoConfigurationTests { @Override public Object postProcessBeforeInitialization(Object bean, String beanName) throws BeansException { if (bean instanceof DataSource dataSource) { - return new ProxyDataSource(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) @@ -280,70 +287,4 @@ class DataSourceHealthContributorAutoConfigurationTests { } - static class ProxyDataSource implements DataSource { - - private final DataSource dataSource; - - ProxyDataSource(DataSource dataSource) { - this.dataSource = dataSource; - } - - @Override - public void setLogWriter(PrintWriter out) throws SQLException { - this.dataSource.setLogWriter(out); - } - - @Override - public Connection getConnection() throws SQLException { - return this.dataSource.getConnection(); - } - - @Override - public Connection getConnection(String username, String password) throws SQLException { - return this.dataSource.getConnection(username, password); - } - - @Override - public PrintWriter getLogWriter() throws SQLException { - return this.dataSource.getLogWriter(); - } - - @Override - public void setLoginTimeout(int seconds) throws SQLException { - this.dataSource.setLoginTimeout(seconds); - } - - @Override - public int getLoginTimeout() throws SQLException { - return this.dataSource.getLoginTimeout(); - } - - @Override - public ConnectionBuilder createConnectionBuilder() throws SQLException { - return this.dataSource.createConnectionBuilder(); - } - - @Override - public Logger getParentLogger() throws SQLFeatureNotSupportedException { - return this.dataSource.getParentLogger(); - } - - @Override - public ShardingKeyBuilder createShardingKeyBuilder() throws SQLException { - return this.dataSource.createShardingKeyBuilder(); - } - - @Override - @SuppressWarnings("unchecked") - public T unwrap(Class iface) throws SQLException { - return iface.isInstance(this) ? (T) this : this.dataSource.unwrap(iface); - } - - @Override - public boolean isWrapperFor(Class iface) throws SQLException { - return (iface.isInstance(this) || this.dataSource.isWrapperFor(iface)); - } - - } - }