handlers: check DB access from heartbeat but don't fail when inaccess…

…ible

handlers.go

@@ -268,6 +268,19 @@ func (a *autographer) handleHeartbeat(w http.ResponseWriter, r *http.Request) {
 		}
 	}
 
+	// check the database connection and return its status, but
+	// don't fail the heartbeat since we only care about DB
+	// connectivity on server start
+	if a.db != nil {
+		err := a.db.CheckConnection()
+		if err == nil {
+			result["dbAccessible"] = true
+		} else {
+			log.Errorf("error checking DB connection: %s", err)
+			result["dbAccessible"] = false
+		}
+	}
+
 	respdata, err := json.Marshal(result)
 	if err != nil {
 		log.Errorf("heartbeat failed to marshal JSON with error: %s", err)

handlers_test.go

@@ -24,6 +24,7 @@ import (
 	"strings"
 	"testing"
 
+	"go.mozilla.org/autograph/database"
 	"go.mozilla.org/autograph/signer/apk"
 	"go.mozilla.org/autograph/signer/contentsignature"
 	"go.mozilla.org/autograph/signer/mar"
@@ -416,7 +417,7 @@ func TestHeartbeat(t *testing.T) {
 }
 
 func TestHeartbeatChecksHSMStatusFails(t *testing.T) {
-	// NB: do not run in parallel with TestHeartbeat
+	// NB: do not run in parallel with TestHeartbeat*
 	ag.hsmHeartbeatSignerConf = &ag.signers[0].(*contentsignature.ContentSigner).Configuration
 
 	expectedStatus := http.StatusInternalServerError
@@ -439,6 +440,40 @@ func TestHeartbeatChecksHSMStatusFails(t *testing.T) {
 	ag.hsmHeartbeatSignerConf = nil
 }
 
+func TestHeartbeatChecksDBStatusOK(t *testing.T) {
+	// NB: do not run in parallel with TestHeartbeat* or DB tests
+	db, err := database.Connect(database.Config{
+		Name:     "autograph",
+		User:     "myautographdbuser",
+		Password: "myautographdbpassword",
+		Host:     "127.0.0.1:5432",
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	ag.db = db
+
+	expectedStatus := http.StatusOK
+	expectedBody := []byte("{\"dbAccessible\":true}")
+
+	req, err := http.NewRequest(`GET`, "http://foo.bar/__heartbeat__", nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	w := httptest.NewRecorder()
+	ag.handleHeartbeat(w, req)
+
+	if w.Code != expectedStatus {
+		t.Fatalf("failed with code %d but %d was expected", w.Code, expectedStatus)
+	}
+	if !bytes.Equal(w.Body.Bytes(), expectedBody) {
+		t.Fatalf("got unexpected heartbeat body %s expected %s", w.Body.Bytes(), expectedBody)
+	}
+
+	db.Close()
+	ag.db = nil
+}
+
 func TestVersion(t *testing.T) {
 	t.Parallel()