From 041d60b5777ae931043a2eebb9318dc320f8cb20 Mon Sep 17 00:00:00 2001
From: aeneasr <3372410+aeneasr@users.noreply.github.com>
Date: Mon, 24 Jul 2023 15:12:39 +0200
Subject: [PATCH] feat: health checks need to consider DB health

Closes #831
---
 internal/driver/registry_default.go   | 33 +++++++++++++++++++++++++++
 internal/persistence/definitions.go   |  2 ++
 internal/persistence/sql/persister.go |  9 ++++++++
 3 files changed, 44 insertions(+)

diff --git a/internal/driver/registry_default.go b/internal/driver/registry_default.go
index 928ce0398..918899d1c 100644
--- a/internal/driver/registry_default.go
+++ b/internal/driver/registry_default.go
@@ -52,6 +52,7 @@ var (
 type (
 	RegistryDefault struct {
 		p               persistence.Persister
+		migrationStatus popx.MigrationStatuses
 		traverser       relationtuple.Traverser
 		mb              *popx.MigrationBox
 		extraMigrations []fs.FS
@@ -130,6 +131,38 @@ func (r *RegistryDefault) HealthHandler() *healthx.Handler {
 		if r.healthReadyCheckers == nil {
 			r.healthReadyCheckers = healthx.ReadyCheckers{}
 		}
+
+		if _, found := r.healthReadyCheckers["database"]; !found {
+			r.healthReadyCheckers["database"] = func(_ *http.Request) error {
+				return r.p.Ping()
+			}
+		}
+
+		if _, found := r.healthReadyCheckers["migrations"]; !found {
+			r.healthReadyCheckers["migrations"] = func(req *http.Request) error {
+				if r.migrationStatus != nil && !r.migrationStatus.HasPending() {
+					return nil
+				}
+
+				mb, err := r.MigrationBox(req.Context())
+				if err != nil {
+					return err
+				}
+
+				status, err := mb.Status(req.Context())
+				if err != nil {
+					return err
+				}
+
+				if status.HasPending() {
+					return errors.Errorf("migrations have not yet been fully applied")
+				}
+
+				r.migrationStatus = status
+				return nil
+			}
+		}
+
 		r.healthH = healthx.NewHandler(r.Writer(), config.Version, r.healthReadyCheckers)
 	}
 
diff --git a/internal/persistence/definitions.go b/internal/persistence/definitions.go
index 876504090..7fd79e980 100644
--- a/internal/persistence/definitions.go
+++ b/internal/persistence/definitions.go
@@ -24,6 +24,8 @@ type (
 		NetworkID(ctx context.Context) uuid.UUID
 		SetNetwork(nid uuid.UUID)
 		Transaction(ctx context.Context, f func(ctx context.Context) error) error
+
+		Ping() error
 	}
 	Migrator interface {
 		MigrationBox(ctx context.Context) (*popx.MigrationBox, error)
diff --git a/internal/persistence/sql/persister.go b/internal/persistence/sql/persister.go
index 7fd5d94f6..ad611bc03 100644
--- a/internal/persistence/sql/persister.go
+++ b/internal/persistence/sql/persister.go
@@ -70,6 +70,15 @@ func (p *Persister) Connection(ctx context.Context) *pop.Connection {
 	return popx.GetConnection(ctx, p.conn.WithContext(ctx))
 }
 
+func (p *Persister) Ping() error {
+	type pinger interface {
+		Ping() error
+	}
+
+	// This can not be contextualized because of some gobuffalo/pop limitations.
+	return errors.WithStack(p.conn.Store.(pinger).Ping())
+}
+
 func (p *Persister) createWithNetwork(ctx context.Context, v interface{}) (err error) {
 	ctx, span := p.d.Tracer(ctx).Tracer().Start(ctx, "persistence.sql.createWithNetwork")
 	defer otelx.End(span, &err)