From 1d026f2365578607984c7b0f56fbe93dbc9c7a3e Mon Sep 17 00:00:00 2001
From: Tristan Stenner <ts@ppi.de>
Date: Wed, 18 Oct 2023 10:17:04 +0200
Subject: [PATCH] Add health check; pings the DB

---
 cmd/main.go                        |  5 ++++-
 internal/postgres/mock/postgres.go | 15 +++++++++++++++
 internal/postgres/postgres.go      |  6 ++++++
 3 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/cmd/main.go b/cmd/main.go
index 9f9bab2a..2d4802e4 100644
--- a/cmd/main.go
+++ b/cmd/main.go
@@ -17,6 +17,7 @@ package main
 
 import (
 	"flag"
+	"net/http"
 	"os"
 
 	// Import all Kubernetes client auth plugins (e.g. Azure, GCP, OIDC, etc.)
@@ -125,7 +126,9 @@ func main() {
 	}
 	//+kubebuilder:scaffold:builder
 
-	if err := mgr.AddHealthzCheck("healthz", healthz.Ping); err != nil {
+	if err := mgr.AddHealthzCheck("healthz", func(req *http.Request) error {
+		return pg.Ping(req.Context())
+	}); err != nil {
 		setupLog.Error(err, "unable to set up health check")
 		os.Exit(1)
 	}
diff --git a/internal/postgres/mock/postgres.go b/internal/postgres/mock/postgres.go
index 09c2f74f..373363d8 100644
--- a/internal/postgres/mock/postgres.go
+++ b/internal/postgres/mock/postgres.go
@@ -5,6 +5,7 @@
 package mock_postgres
 
 import (
+	context "context"
 	reflect "reflect"
 
 	logr "github.com/go-logr/logr"
@@ -203,6 +204,20 @@ func (mr *MockPGMockRecorder) GrantRole(role, grantee interface{}) *gomock.Call
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "GrantRole", reflect.TypeOf((*MockPG)(nil).GrantRole), role, grantee)
 }
 
+// Ping mocks base method.
+func (m *MockPG) Ping(ctx context.Context) error {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Ping", ctx)
+	ret0, _ := ret[0].(error)
+	return ret0
+}
+
+// Ping indicates an expected call of Ping.
+func (mr *MockPGMockRecorder) Ping(ctx interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Ping", reflect.TypeOf((*MockPG)(nil).Ping), ctx)
+}
+
 // RevokeRole mocks base method.
 func (m *MockPG) RevokeRole(role, revoked string) error {
 	m.ctrl.T.Helper()
diff --git a/internal/postgres/postgres.go b/internal/postgres/postgres.go
index 430679f8..19fa6b35 100644
--- a/internal/postgres/postgres.go
+++ b/internal/postgres/postgres.go
@@ -1,6 +1,7 @@
 package postgres
 
 import (
+	"context"
 	"database/sql"
 	"fmt"
 	"log"
@@ -25,6 +26,7 @@ type PG interface {
 	GetUser() string
 	GetHost() string
 	GetDefaultDatabase() string
+	Ping(ctx context.Context) error
 }
 
 type pg struct {
@@ -81,6 +83,10 @@ func (c *pg) GetDefaultDatabase() string {
 	return c.default_database
 }
 
+func (c *pg) Ping(ctx context.Context) error {
+	return c.db.PingContext(ctx)
+}
+
 func GetConnection(user, password, host, database, uri_args string) (*sql.DB, error) {
 	db, err := sql.Open("postgres", fmt.Sprintf("postgresql://%s:%s@%s/%s?%s", user, password, host, database, uri_args))
 	if err != nil {