Merge pull request #2763 from cloudfoundry/feat/healthzDatabaseCheckB…

…ackwardsCompatible

feat: Use Database availability as indicator for /healthz response

server/src/main/java/org/cloudfoundry/identity/uaa/health/HealthzEndpoint.java

@@ -1,14 +1,19 @@
 package org.cloudfoundry.identity.uaa.health;
 
+import javax.servlet.http.HttpServletResponse;
+import javax.sql.DataSource;
+
+import java.sql.Connection;
+import java.sql.Statement;
+
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 import org.springframework.beans.factory.annotation.Value;
+import org.springframework.scheduling.annotation.Scheduled;
 import org.springframework.stereotype.Controller;
 import org.springframework.web.bind.annotation.GetMapping;
 import org.springframework.web.bind.annotation.ResponseBody;
 
-import javax.servlet.http.HttpServletResponse;
-
 /**
  * Simple controller that just returns "ok" in a request body for the purposes
  * of monitoring health of the application. It also registers a shutdown hook
@@ -18,10 +23,13 @@
 public class HealthzEndpoint {
     private static Logger logger = LoggerFactory.getLogger(HealthzEndpoint.class);
     private volatile boolean stopping = false;
+    private volatile Boolean wasLastConnectionSuccessful = null;
+    private DataSource dataSource;
 
     public HealthzEndpoint(
             @Value("${uaa.shutdown.sleep:10000}") final long sleepTime,
-            final Runtime runtime) {
+            final Runtime runtime,
+            final DataSource dataSource) {
         Thread shutdownHook = new Thread(() -> {
             stopping = true;
             logger.warn("Shutdown hook received, future requests to this endpoint will return 503");
@@ -35,6 +43,7 @@ public HealthzEndpoint(
             }
         });
         runtime.addShutdownHook(shutdownHook);
+        this.dataSource = dataSource;
     }
 
     @GetMapping("/healthz")
@@ -45,8 +54,28 @@ public String getHealthz(HttpServletResponse response) {
             response.setStatus(503);
             return "stopping\n";
         } else {
-            return "ok\n";
+            if (wasLastConnectionSuccessful == null) {
+                return "UAA running. Database status unknown.\n";
+            }
+
+            if (wasLastConnectionSuccessful) {
+                return "ok\n";
+            } else {
+                response.setStatus(503);
+                return "Database Connection failed.\n";
+            }
         }
     }
 
-}
+    @Scheduled(fixedRateString = "${uaa.health.db.rate:10000}")
+    void isDataSourceConnectionAvailable() {
+        try (Connection c = dataSource.getConnection(); Statement statement = c.createStatement()) {
+            statement.execute("SELECT 1 from identity_zone;"); //"SELECT 1;" Not supported by HSQLDB
+            wasLastConnectionSuccessful = true;
+            return;
+        } catch (Exception ex) {
+            logger.error("Could not establish connection to DB - " + ex.getMessage());
+        }
+        wasLastConnectionSuccessful = false;
+    }
+}

server/src/test/java/org/cloudfoundry/identity/uaa/web/HealthzEndpointTests.java → server/src/test/java/org/cloudfoundry/identity/uaa/health/HealthzEndpointTests.java

@@ -1,18 +1,25 @@
-package org.cloudfoundry.identity.uaa.web;
-
-import org.cloudfoundry.identity.uaa.health.HealthzEndpoint;
-import org.junit.jupiter.api.BeforeEach;
-import org.junit.jupiter.api.Nested;
-import org.junit.jupiter.api.Test;
-import org.mockito.ArgumentCaptor;
-import org.springframework.mock.web.MockHttpServletResponse;
+package org.cloudfoundry.identity.uaa.health;
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.greaterThanOrEqualTo;
 import static org.hamcrest.Matchers.lessThanOrEqualTo;
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.mockito.ArgumentMatchers.anyString;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.when;
+
+import javax.sql.DataSource;
+
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.sql.Statement;
+
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Nested;
+import org.junit.jupiter.api.Test;
+import org.mockito.ArgumentCaptor;
+import org.springframework.mock.web.MockHttpServletResponse;
 
 class HealthzEndpointTests {
 
@@ -21,11 +28,19 @@ class HealthzEndpointTests {
     private HealthzEndpoint endpoint;
     private MockHttpServletResponse response;
     private Thread shutdownHook;
+    private DataSource dataSource;
+    private Connection connection;
+    private Statement statement;
 
     @BeforeEach
-    void setUp() {
+    void setUp() throws SQLException {
         Runtime mockRuntime = mock(Runtime.class);
-        endpoint = new HealthzEndpoint(SLEEP_UPON_SHUTDOWN, mockRuntime);
+        dataSource = mock(DataSource.class);
+        connection = mock(Connection.class);
+        statement = mock(Statement.class);
+        when(dataSource.getConnection()).thenReturn(connection);
+        when(connection.createStatement()).thenReturn(statement);
+        endpoint = new HealthzEndpoint(SLEEP_UPON_SHUTDOWN, mockRuntime, dataSource);
         response = new MockHttpServletResponse();
 
         ArgumentCaptor<Thread> threadArgumentCaptor = ArgumentCaptor.forClass(Thread.class);
@@ -35,8 +50,21 @@ void setUp() {
 
     @Test
     void getHealthz() {
+        assertEquals("UAA running. Database status unknown.\n", endpoint.getHealthz(response));
+    }
+
+    @Test
+    void getHealthz_connectionSuccess() {
+        endpoint.isDataSourceConnectionAvailable();
         assertEquals("ok\n", endpoint.getHealthz(response));
     }
+    @Test
+    void getHealthz_connectionFailed() throws SQLException {
+        when(statement.execute(anyString())).thenThrow(new SQLException());
+        endpoint.isDataSourceConnectionAvailable();
+        assertEquals("Database Connection failed.\n", endpoint.getHealthz(response));
+        assertEquals(503, response.getStatus());
+    }
 
     @Test
     void shutdownSendsStopping() throws InterruptedException {
@@ -54,7 +82,8 @@ class WithoutSleeping {
         @BeforeEach
         void setUp() {
             Runtime mockRuntime = mock(Runtime.class);
-            endpoint = new HealthzEndpoint(-1, mockRuntime);
+            DataSource dataSource = mock(DataSource.class);
+            endpoint = new HealthzEndpoint(-1, mockRuntime, dataSource);
             response = new MockHttpServletResponse();
 
             ArgumentCaptor<Thread> threadArgumentCaptor = ArgumentCaptor.forClass(Thread.class);

uaa/src/test/java/org/cloudfoundry/identity/uaa/integration/HealthzEndpointIntegrationTests.java

@@ -42,7 +42,6 @@ public void testHappyDay() {
 
         String body = response.getBody();
         assertTrue(body.contains("ok"));
-
     }
 
 }