Reuse data source validation query in health endpoint

This commit improves DataSourceMetadata to expose the validation
query. This can be used by DataSourceHealthIndicator as the query
to use instead of "guessing" which query could be applied according
to the database type.

Fixes gh-1282

spring-boot-actuator/src/main/java/org/springframework/boot/actuate/autoconfigure/HealthIndicatorAutoConfiguration.java

@@ -16,6 +16,8 @@
 
 package org.springframework.boot.actuate.autoconfigure;
 
+import java.util.Collection;
+import java.util.Collections;
 import java.util.List;
 import java.util.Map;
 
@@ -42,7 +44,10 @@ import org.springframework.boot.autoconfigure.amqp.RabbitAutoConfiguration;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnBean;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnExpression;
 import org.springframework.boot.autoconfigure.condition.ConditionalOnMissingBean;
+import org.springframework.boot.autoconfigure.jdbc.CompositeDataSourceMetadataProvider;
 import org.springframework.boot.autoconfigure.jdbc.DataSourceAutoConfiguration;
+import org.springframework.boot.autoconfigure.jdbc.DataSourceMetadata;
+import org.springframework.boot.autoconfigure.jdbc.DataSourceMetadataProvider;
 import org.springframework.boot.autoconfigure.mongo.MongoAutoConfiguration;
 import org.springframework.boot.autoconfigure.mongo.MongoDataAutoConfiguration;
 import org.springframework.boot.autoconfigure.redis.RedisAutoConfiguration;
@@ -57,6 +62,7 @@ import org.springframework.data.redis.connection.RedisConnectionFactory;
  *
  * @author Christian Dupuis
  * @author Andy Wilkinson
+ * @author Stephane Nicoll
  * @since 1.1.0
  */
 @Configuration
@@ -96,21 +102,37 @@ public class HealthIndicatorAutoConfiguration {
 		@Autowired(required = false)
 		private Map<String, DataSource> dataSources;
 
+		@Autowired(required = false)
+		private Collection<DataSourceMetadataProvider> metadataProviders = Collections.emptyList();
+
 		@Bean
 		@ConditionalOnMissingBean(name = "dbHealthIndicator")
 		public HealthIndicator dbHealthIndicator() {
+			DataSourceMetadataProvider metadataProvider =
+					new CompositeDataSourceMetadataProvider(this.metadataProviders);
 			if (this.dataSources.size() == 1) {
-				return new DataSourceHealthIndicator(this.dataSources.values().iterator()
-						.next());
+				return createDataSourceHealthIndicator(metadataProvider,
+						this.dataSources.values().iterator().next());
 			}
 			CompositeHealthIndicator composite = new CompositeHealthIndicator(
 					this.healthAggregator);
 			for (Map.Entry<String, DataSource> entry : this.dataSources.entrySet()) {
 				composite.addHealthIndicator(entry.getKey(),
-						new DataSourceHealthIndicator(entry.getValue()));
+						createDataSourceHealthIndicator(metadataProvider, entry.getValue()));
 			}
 			return composite;
 		}
+
+		private DataSourceHealthIndicator createDataSourceHealthIndicator(DataSourceMetadataProvider provider,
+				DataSource dataSource) {
+
+			String validationQuery = null;
+			DataSourceMetadata dataSourceMetadata = provider.getDataSourceMetadata(dataSource);
+			if (dataSourceMetadata != null) {
+				validationQuery = dataSourceMetadata.getValidationQuery();
+			}
+			return new DataSourceHealthIndicator(dataSource, validationQuery);
+		}
 	}
 
 	@Configuration

spring-boot-actuator/src/main/java/org/springframework/boot/actuate/health/DataSourceHealthIndicator.java

@@ -42,6 +42,7 @@ import org.springframework.util.StringUtils;
  * @author Dave Syer
  * @author Christian Dupuis
  * @author Andy Wilkinson
+ * @author Stephane Nicoll
  * @since 1.1.0
  */
 public class DataSourceHealthIndicator extends AbstractHealthIndicator {
@@ -50,7 +51,7 @@ public class DataSourceHealthIndicator extends AbstractHealthIndicator {
 
 	private JdbcTemplate jdbcTemplate;
 
-	private static Map<String, String> queries = new HashMap<String, String>();
+	private static final Map<String, String> queries = new HashMap<String, String>();
 
 	static {
 		queries.put("HSQL Database Engine", "SELECT COUNT(*) FROM "
@@ -59,7 +60,7 @@ public class DataSourceHealthIndicator extends AbstractHealthIndicator {
 		queries.put("Apache Derby", "SELECT 1 FROM SYSIBM.SYSDUMMY1");
 	}
 
-	private static String DEFAULT_QUERY = "SELECT 1";
+	private static final String DEFAULT_QUERY = "SELECT 1";
 
 	private String query = null;
 
@@ -72,12 +73,22 @@ public class DataSourceHealthIndicator extends AbstractHealthIndicator {
 	/**
 	 * Create a new {@link DataSourceHealthIndicator} using the specified datasource.
 	 * @param dataSource the data source
+	 * @param query the validation query to use (can be {@code null})
 	 */
-	public DataSourceHealthIndicator(DataSource dataSource) {
+	public DataSourceHealthIndicator(DataSource dataSource, String query) {
 		this.dataSource = dataSource;
+		this.query = query;
 		this.jdbcTemplate = new JdbcTemplate(dataSource);
 	}
 
+	/**
+	 * Create a new {@link DataSourceHealthIndicator} using the specified datasource.
+	 * @param dataSource the data source
+	 */
+	public DataSourceHealthIndicator(DataSource dataSource) {
+		this(dataSource, null);
+	}
+
 	@Override
 	protected void doHealthCheck(Health.Builder builder) throws Exception {
 		if (this.dataSource == null) {
@@ -127,15 +138,29 @@ public class DataSourceHealthIndicator extends AbstractHealthIndicator {
 		return query;
 	}
 
+	/**
+	 * Set the {@link DataSource} to use.
+	 */
 	public void setDataSource(DataSource dataSource) {
 		this.dataSource = dataSource;
 		this.jdbcTemplate = new JdbcTemplate(dataSource);
 	}
 
+	/**
+	 * Set a specific validation query to use to validate a connection. If
+	 * none is set, a default validation query is used.
+	 */
 	public void setQuery(String query) {
 		this.query = query;
 	}
 
+	/**
+	 * Return the specific validation query, if any.
+	 */
+	public String getQuery() {
+		return query;
+	}
+
 	/**
 	 * {@link RowMapper} that expects and returns results from a single column.
 	 */

spring-boot-actuator/src/test/java/org/springframework/boot/actuate/autoconfigure/HealthIndicatorAutoConfigurationTests.java

@@ -18,6 +18,8 @@ package org.springframework.boot.actuate.autoconfigure;
 
 import java.util.Map;
 
+import javax.sql.DataSource;
+
 import org.junit.After;
 import org.junit.Before;
 import org.junit.Test;
@@ -28,14 +30,22 @@ import org.springframework.boot.actuate.health.RabbitHealthIndicator;
 import org.springframework.boot.actuate.health.RedisHealthIndicator;
 import org.springframework.boot.actuate.health.SolrHealthIndicator;
 import org.springframework.boot.actuate.health.ApplicationHealthIndicator;
+import org.springframework.boot.autoconfigure.PropertyPlaceholderAutoConfiguration;
 import org.springframework.boot.autoconfigure.amqp.RabbitAutoConfiguration;
+import org.springframework.boot.autoconfigure.jdbc.DataSourceBuilder;
+import org.springframework.boot.autoconfigure.jdbc.DataSourceMetadataProvidersConfiguration;
+import org.springframework.boot.autoconfigure.jdbc.DataSourceProperties;
 import org.springframework.boot.autoconfigure.jdbc.EmbeddedDataSourceConfiguration;
 import org.springframework.boot.autoconfigure.mongo.MongoAutoConfiguration;
 import org.springframework.boot.autoconfigure.mongo.MongoDataAutoConfiguration;
 import org.springframework.boot.autoconfigure.redis.RedisAutoConfiguration;
 import org.springframework.boot.autoconfigure.solr.SolrAutoConfiguration;
+import org.springframework.boot.context.properties.ConfigurationProperties;
+import org.springframework.boot.context.properties.EnableConfigurationProperties;
 import org.springframework.boot.test.EnvironmentTestUtils;
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
+import org.springframework.context.annotation.Bean;
+import org.springframework.context.annotation.Configuration;
 
 import static org.junit.Assert.assertEquals;
 
@@ -153,6 +163,22 @@ public class HealthIndicatorAutoConfigurationTests {
 				.getClass());
 	}
 
+	@Test
+	public void dataSourceHealthIndicatorWithCustomValidationQuery() {
+		this.context = new AnnotationConfigApplicationContext();
+		this.context.register(PropertyPlaceholderAutoConfiguration.class, DataSourceProperties.class,
+				DataSourceConfig.class, DataSourceMetadataProvidersConfiguration.class,
+				HealthIndicatorAutoConfiguration.class);
+		EnvironmentTestUtils.addEnvironment(this.context, "spring.datasource.validation-query:SELECT from FOOBAR");
+		this.context.refresh();
+		Map<String, HealthIndicator> beans = this.context.getBeansOfType(HealthIndicator.class);
+		assertEquals(1, beans.size());
+		HealthIndicator healthIndicator = beans.values().iterator().next();
+		assertEquals(DataSourceHealthIndicator.class, healthIndicator.getClass());
+		DataSourceHealthIndicator dataSourceHealthIndicator = (DataSourceHealthIndicator) healthIndicator;
+		assertEquals("SELECT from FOOBAR", dataSourceHealthIndicator.getQuery());
+	}
+
 	@Test
 	public void notDataSourceHealthIndicator() {
 		this.context = new AnnotationConfigApplicationContext();
@@ -220,4 +246,19 @@ public class HealthIndicatorAutoConfigurationTests {
 		assertEquals(ApplicationHealthIndicator.class, beans.values().iterator().next()
 				.getClass());
 	}
+
+
+	@Configuration
+	@EnableConfigurationProperties
+	protected static class DataSourceConfig {
+
+		@Bean
+		@ConfigurationProperties(prefix = DataSourceProperties.PREFIX)
+		public DataSource dataSource() {
+			return DataSourceBuilder.create()
+					.driverClassName("org.hsqldb.jdbc.JDBCDriver")
+					.url("jdbc:hsqldb:mem:test")
+					.username("sa").build();
+		}
+	}
 }

spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/jdbc/CommonsDbcpDataSourceMetadata.java

@@ -46,4 +46,8 @@ public class CommonsDbcpDataSourceMetadata extends AbstractDataSourceMetadata<Ba
 		return getDataSource().getMinIdle();
 	}
 
+	@Override
+	public String getValidationQuery() {
+		return getDataSource().getValidationQuery();
+	}
 }

spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/jdbc/DataSourceMetadata.java

@@ -64,4 +64,10 @@ public interface DataSourceMetadata {
 	 */
 	Integer getMinPoolSize();
 
+	/**
+	 * Return the query to use to validate that a connection is
+	 * valid or {@code null} if that information is not available.
+	 */
+	String getValidationQuery();
+
 }

spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/jdbc/HikariDataSourceMetadata.java

@@ -57,6 +57,11 @@ public class HikariDataSourceMetadata extends AbstractDataSourceMetadata<HikariD
 		return getDataSource().getMinimumIdle();
 	}
 
+	@Override
+	public String getValidationQuery() {
+		return getDataSource().getConnectionTestQuery();
+	}
+
 	/**
 	 * Provide the {@link HikariPool} instance managed internally by
 	 * the {@link HikariDataSource} as there is no other way to retrieve

spring-boot-autoconfigure/src/main/java/org/springframework/boot/autoconfigure/jdbc/TomcatDataSourceMetadata.java

@@ -48,4 +48,8 @@ public class TomcatDataSourceMetadata extends AbstractDataSourceMetadata<DataSou
 		return getDataSource().getMinIdle();
 	}
 
+	@Override
+	public String getValidationQuery() {
+		return getDataSource().getValidationQuery();
+	}
 }

spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/jdbc/AbstractDataSourceMetadataTests.java

@@ -94,6 +94,9 @@ public abstract class AbstractDataSourceMetadataTests<D extends AbstractDataSour
 		});
 	}
 
+	@Test
+	public abstract void getValidationQuery();
+
 	protected DataSourceBuilder initializeBuilder() {
 		return DataSourceBuilder.create()
 				.driverClassName("org.hsqldb.jdbc.JDBCDriver")

spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/jdbc/CommonsDbcpDataSourceMetadataTests.java

@@ -68,6 +68,13 @@ public class CommonsDbcpDataSourceMetadataTests extends AbstractDataSourceMetada
 		assertEquals(Float.valueOf(-1F), unlimitedDataSource.getPoolUsage());
 	}
 
+	@Override
+	public void getValidationQuery() {
+		BasicDataSource dataSource = createDataSource();
+		dataSource.setValidationQuery("SELECT FROM FOO");
+		assertEquals("SELECT FROM FOO", new CommonsDbcpDataSourceMetadata(dataSource).getValidationQuery());
+	}
+
 	private CommonsDbcpDataSourceMetadata createDataSourceMetadata(int minSize, int maxSize) {
 		BasicDataSource dataSource = createDataSource();
 		dataSource.setMinIdle(minSize);

spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/jdbc/HikariDataSourceMetadataTests.java

@@ -19,8 +19,9 @@ package org.springframework.boot.autoconfigure.jdbc;
 import com.zaxxer.hikari.HikariDataSource;
 import org.junit.Before;
 
+import static org.junit.Assert.assertEquals;
+
 /**
- *
  * @author Stephane Nicoll
  */
 public class HikariDataSourceMetadataTests extends AbstractDataSourceMetadataTests<HikariDataSourceMetadata> {
@@ -29,7 +30,7 @@ public class HikariDataSourceMetadataTests extends AbstractDataSourceMetadataTes
 
 	@Before
 	public void setup() {
-		this.dataSourceMetadata = createDataSourceMetadata(0, 2);
+		this.dataSourceMetadata =  new HikariDataSourceMetadata(createDataSource(0, 2));
 	}
 
 	@Override
@@ -37,11 +38,17 @@ public class HikariDataSourceMetadataTests extends AbstractDataSourceMetadataTes
 		return this.dataSourceMetadata;
 	}
 
-	private HikariDataSourceMetadata createDataSourceMetadata(int minSize, int maxSize) {
+	@Override
+	public void getValidationQuery() {
+		HikariDataSource dataSource = createDataSource(0, 4);
+		dataSource.setConnectionTestQuery("SELECT FROM FOO");
+		assertEquals("SELECT FROM FOO", new HikariDataSourceMetadata(dataSource).getValidationQuery());
+	}
+
+	private HikariDataSource createDataSource(int minSize, int maxSize) {
 		HikariDataSource dataSource = (HikariDataSource) initializeBuilder().type(HikariDataSource.class).build();
 		dataSource.setMinimumIdle(minSize);
 		dataSource.setMaximumPoolSize(maxSize);
-
-		return new HikariDataSourceMetadata(dataSource);
+		return dataSource;
 	}
 }

spring-boot-autoconfigure/src/test/java/org/springframework/boot/autoconfigure/jdbc/TomcatDataSourceMetadataTests.java

@@ -19,6 +19,8 @@ package org.springframework.boot.autoconfigure.jdbc;
 import org.apache.tomcat.jdbc.pool.DataSource;
 import org.junit.Before;
 
+import static org.junit.Assert.assertEquals;
+
 /**
  *
  * @author Stephane Nicoll
@@ -29,7 +31,7 @@ public class TomcatDataSourceMetadataTests extends AbstractDataSourceMetadataTes
 
 	@Before
 	public void setup() {
-		this.dataSourceMetadata = createDataSourceMetadata(0, 2);
+		this.dataSourceMetadata = new TomcatDataSourceMetadata(createDataSource(0, 2));
 	}
 
 	@Override
@@ -37,7 +39,14 @@ public class TomcatDataSourceMetadataTests extends AbstractDataSourceMetadataTes
 		return this.dataSourceMetadata;
 	}
 
-	private TomcatDataSourceMetadata createDataSourceMetadata(int minSize, int maxSize) {
+	@Override
+	public void getValidationQuery() {
+		DataSource dataSource = createDataSource(0, 4);
+		dataSource.setValidationQuery("SELECT FROM FOO");
+		assertEquals("SELECT FROM FOO", new TomcatDataSourceMetadata(dataSource).getValidationQuery());
+	}
+
+	private DataSource createDataSource(int minSize, int maxSize) {
 		DataSource dataSource = (DataSource) initializeBuilder().type(DataSource.class).build();
 		dataSource.setMinIdle(minSize);
 		dataSource.setMaxActive(maxSize);
@@ -45,7 +54,7 @@ public class TomcatDataSourceMetadataTests extends AbstractDataSourceMetadataTes
 		// Avoid warnings
 		dataSource.setInitialSize(minSize);
 		dataSource.setMaxIdle(maxSize);
-		return new TomcatDataSourceMetadata(dataSource);
+		return dataSource;
 	}
 
 }

spring-boot-docs/src/main/asciidoc/production-ready-features.adoc

@@ -164,8 +164,8 @@ To provide custom health information you can register a Spring bean that impleme
 
 Spring Boot provides a
 {sc-spring-boot-actuator}/health/DataSourceHealthIndicator.{sc-ext}[`DataSourceHealthIndicator`]
-implementation that attempts a simple database test as well as implementations for
-Redis, MongoDB and RabbitMQ.
+implementation that attempts a simple database test (reusing the validation query set on the data
+source, if any) as well as implementations for Redis, MongoDB and RabbitMQ.
 
 Spring Boot adds the `HealthIndicator` instances automatically if beans of type `DataSource`,
 `MongoTemplate`, `RedisConnectionFactory`, `RabbitTemplate` are present in the `ApplicationContext`.