From e3c28a6c96b91779a22b36c1bf880dee3684b709 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Fr=C3=A9d=C3=A9ric=20Guillot?= <f@miniflux.net>
Date: Fri, 19 Feb 2021 18:47:50 -0800
Subject: [PATCH] Improve health check endpoint to test database connection

---
 cli/cli.go             | 10 +++++++---
 service/httpd/httpd.go |  5 +++++
 storage/feed.go        | 11 -----------
 storage/storage.go     | 17 +++++++++++++++++
 4 files changed, 29 insertions(+), 14 deletions(-)

diff --git a/cli/cli.go b/cli/cli.go
index 316e5164..5b100eb4 100644
--- a/cli/cli.go
+++ b/cli/cli.go
@@ -126,17 +126,21 @@ func Parse() {
 		config.Opts.DatabaseMaxConns(),
 	)
 	if err != nil {
-		logger.Fatal("Unable to connect to the database: %v", err)
+		logger.Fatal("Unable to initialize database connection pool: %v", err)
 	}
 	defer db.Close()
 
+	store := storage.NewStorage(db)
+
+	if err := store.Ping(); err != nil {
+		logger.Fatal("Unable to connect to the database: %v", err)
+	}
+
 	if flagMigrate {
 		database.Migrate(db)
 		return
 	}
 
-	store := storage.NewStorage(db)
-
 	if flagResetFeedErrors {
 		store.ResetFeedErrors()
 		return
diff --git a/service/httpd/httpd.go b/service/httpd/httpd.go
index b1ad10b3..09a0e0ef 100644
--- a/service/httpd/httpd.go
+++ b/service/httpd/httpd.go
@@ -184,6 +184,11 @@ func setupHandler(store *storage.Storage, pool *worker.Pool) *mux.Router {
 	ui.Serve(router, store, pool)
 
 	router.HandleFunc("/healthcheck", func(w http.ResponseWriter, r *http.Request) {
+		if err := store.Ping(); err != nil {
+			http.Error(w, "Database Connection Error", http.StatusInternalServerError)
+			return
+		}
+
 		w.Write([]byte("OK"))
 	}).Name("healthcheck")
 
diff --git a/storage/feed.go b/storage/feed.go
index 1e21858a..e9d51d1b 100644
--- a/storage/feed.go
+++ b/storage/feed.go
@@ -79,17 +79,6 @@ func (s *Storage) CountFeeds(userID int64) int {
 	return result
 }
 
-// DatabaseVersion returns the version of the database which is in use
-func (s *Storage) DatabaseVersion() string {
-	var dbVersion string
-	err := s.db.QueryRow(`SELECT current_setting('server_version')`).Scan(&dbVersion)
-	if err != nil {
-		return err.Error()
-	}
-
-	return dbVersion
-}
-
 // CountUserFeedsWithErrors returns the number of feeds with parsing errors that belong to the given user.
 func (s *Storage) CountUserFeedsWithErrors(userID int64) int {
 	pollingParsingErrorLimit := config.Opts.PollingParsingErrorLimit()
diff --git a/storage/storage.go b/storage/storage.go
index c2aef77e..ef73e569 100644
--- a/storage/storage.go
+++ b/storage/storage.go
@@ -17,3 +17,20 @@ type Storage struct {
 func NewStorage(db *sql.DB) *Storage {
 	return &Storage{db}
 }
+
+// DatabaseVersion returns the version of the database which is in use.
+func (s *Storage) DatabaseVersion() string {
+	var dbVersion string
+	err := s.db.QueryRow(`SELECT current_setting('server_version')`).Scan(&dbVersion)
+	if err != nil {
+		return err.Error()
+	}
+
+	return dbVersion
+}
+
+// Ping checks if the database connection works.
+func (s *Storage) Ping() error {
+	_, err := s.db.Exec(`SELECT true`)
+	return err
+}