From 43bd42f0f5a132bc61af77fc4e82118df61043f5 Mon Sep 17 00:00:00 2001
From: Tyler Frederick <tyler@3cinteractive.com>
Date: Mon, 5 May 2014 11:27:04 +0100
Subject: [PATCH] Add multi-datasource health indicator support

Update EndpointAutoConfiguration to support multiple datasources.

Fixes gh-783
---
 .../EndpointAutoConfiguration.java            | 31 +++++++++++++------
 .../actuate/health/SimpleHealthIndicator.java | 15 +++++++++
 2 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/autoconfigure/EndpointAutoConfiguration.java b/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/autoconfigure/EndpointAutoConfiguration.java
index f2f3744260..c06da53e27 100644
--- a/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/autoconfigure/EndpointAutoConfiguration.java
+++ b/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/autoconfigure/EndpointAutoConfiguration.java
@@ -38,6 +38,7 @@ import org.springframework.boot.actuate.endpoint.RequestMappingEndpoint;
 import org.springframework.boot.actuate.endpoint.ShutdownEndpoint;
 import org.springframework.boot.actuate.endpoint.TraceEndpoint;
 import org.springframework.boot.actuate.endpoint.VanillaPublicMetrics;
+import org.springframework.boot.actuate.health.CompositeHealthIndicator;
 import org.springframework.boot.actuate.health.HealthIndicator;
 import org.springframework.boot.actuate.health.SimpleHealthIndicator;
 import org.springframework.boot.actuate.health.VanillaHealthIndicator;
@@ -75,7 +76,7 @@ public class EndpointAutoConfiguration {
 	private HealthIndicator<? extends Object> healthIndicator;
 
 	@Autowired(required = false)
-	private DataSource dataSource;
+	private Map<String, DataSource> dataSources;
 
 	@Autowired
 	private InfoPropertiesConfiguration properties;
@@ -97,20 +98,30 @@ public class EndpointAutoConfiguration {
 
 	@Bean
 	@ConditionalOnMissingBean
-	public HealthEndpoint<Object> healthEndpoint() {
+	public HealthEndpoint<?> healthEndpoint() {
 		if (this.healthIndicator == null) {
-			if (this.dataSource == null) {
-				this.healthIndicator = new VanillaHealthIndicator();
-			}
-			else {
-				SimpleHealthIndicator healthIndicator = new SimpleHealthIndicator();
-				healthIndicator.setDataSource(this.dataSource);
-				this.healthIndicator = healthIndicator;
-			}
+			this.healthIndicator = createHealthIndicator();
 		}
 		return new HealthEndpoint<Object>(this.healthIndicator);
 	}
 
+	private HealthIndicator<? extends Object> createHealthIndicator() {
+		if (this.dataSources == null || this.dataSources.isEmpty()) {
+			return new VanillaHealthIndicator();
+		}
+
+		if (this.dataSources.size() == 1) {
+			return new SimpleHealthIndicator(this.dataSources.values().iterator().next());
+		}
+
+		CompositeHealthIndicator composite = new CompositeHealthIndicator();
+		for (Map.Entry<String, DataSource> entry : this.dataSources.entrySet()) {
+			composite.addHealthIndicator(entry.getKey(),
+					new SimpleHealthIndicator(entry.getValue()));
+		}
+		return composite;
+	}
+
 	@Bean
 	@ConditionalOnMissingBean
 	public BeansEndpoint beansEndpoint() {
diff --git a/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/health/SimpleHealthIndicator.java b/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/health/SimpleHealthIndicator.java
index 61fcc73580..4c185c4a5b 100644
--- a/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/health/SimpleHealthIndicator.java
+++ b/spring-boot-actuator/src/main/java/org/springframework/boot/actuate/health/SimpleHealthIndicator.java
@@ -54,6 +54,21 @@ public class SimpleHealthIndicator implements HealthIndicator<Map<String, Object
 
 	private String query = null;
 
+	/**
+	 * Create a new {@link SimpleHealthIndicator} instance.
+	 */
+	public SimpleHealthIndicator() {
+	}
+
+	/**
+	 * Create a new {@link SimpleHealthIndicator} using the specified datasource.
+	 * @param dataSource the data source
+	 */
+	public SimpleHealthIndicator(DataSource dataSource) {
+		this.dataSource = dataSource;
+		this.jdbcTemplate = new JdbcTemplate(dataSource);
+	}
+
 	@Override
 	public Map<String, Object> health() {
 		LinkedHashMap<String, Object> health = new LinkedHashMap<String, Object>();