Context: Add context to /api/health calls (#40031)

Signed-off-by: bergquist <carl.bergquist@gmail.com>

pkg/api/health.go

@@ -1,20 +1,21 @@
 package api
 
 import (
+	"context"
 	"time"
 
 	"github.com/grafana/grafana/pkg/bus"
 	"github.com/grafana/grafana/pkg/models"
 )
 
-func (hs *HTTPServer) databaseHealthy() bool {
+func (hs *HTTPServer) databaseHealthy(ctx context.Context) bool {
 	const cacheKey = "db-healthy"
 
 	if cached, found := hs.CacheService.Get(cacheKey); found {
 		return cached.(bool)
 	}
 
-	healthy := bus.Dispatch(&models.GetDBHealthQuery{}) == nil
+	healthy := bus.DispatchCtx(ctx, &models.GetDBHealthQuery{}) == nil
 
 	hs.CacheService.Set(cacheKey, healthy, time.Second*5)
 	return healthy

pkg/api/health_test.go

@@ -1,6 +1,7 @@
 package api
 
 import (
+	"context"
 	"errors"
 	"net/http"
 	"net/http/httptest"
@@ -21,7 +22,7 @@ func TestHealthAPI_Version(t *testing.T) {
 		cfg.BuildCommit = "59906ab1bf"
 	})
 
-	bus.AddHandler("test", func(query *models.GetDBHealthQuery) error {
+	bus.AddHandlerCtx("test", func(ctx context.Context, query *models.GetDBHealthQuery) error {
 		return nil
 	})
 
@@ -44,7 +45,7 @@ func TestHealthAPI_AnonymousHideVersion(t *testing.T) {
 	m, hs := setupHealthAPITestEnvironment(t)
 	hs.Cfg.AnonymousHideVersion = true
 
-	bus.AddHandler("test", func(query *models.GetDBHealthQuery) error {
+	bus.AddHandlerCtx("test", func(ctx context.Context, query *models.GetDBHealthQuery) error {
 		return nil
 	})
 
@@ -67,7 +68,7 @@ func TestHealthAPI_DatabaseHealthy(t *testing.T) {
 	m, hs := setupHealthAPITestEnvironment(t)
 	hs.Cfg.AnonymousHideVersion = true
 
-	bus.AddHandler("test", func(query *models.GetDBHealthQuery) error {
+	bus.AddHandlerCtx("test", func(ctx context.Context, query *models.GetDBHealthQuery) error {
 		return nil
 	})
 
@@ -98,7 +99,7 @@ func TestHealthAPI_DatabaseUnhealthy(t *testing.T) {
 	m, hs := setupHealthAPITestEnvironment(t)
 	hs.Cfg.AnonymousHideVersion = true
 
-	bus.AddHandler("test", func(query *models.GetDBHealthQuery) error {
+	bus.AddHandlerCtx("test", func(ctx context.Context, query *models.GetDBHealthQuery) error {
 		return errors.New("bad")
 	})
 
@@ -130,7 +131,7 @@ func TestHealthAPI_DatabaseHealthCached(t *testing.T) {
 	hs.Cfg.AnonymousHideVersion = true
 
 	// Database is healthy.
-	bus.AddHandler("test", func(query *models.GetDBHealthQuery) error {
+	bus.AddHandlerCtx("test", func(ctx context.Context, query *models.GetDBHealthQuery) error {
 		return nil
 	})
 

pkg/api/http_server.go

@@ -493,7 +493,7 @@ func (hs *HTTPServer) apiHealthHandler(ctx *web.Context) {
 		data.Set("commit", hs.Cfg.BuildCommit)
 	}
 
-	if !hs.databaseHealthy() {
+	if !hs.databaseHealthy(ctx.Req.Context()) {
 		data.Set("database", "failing")
 		ctx.Resp.Header().Set("Content-Type", "application/json; charset=UTF-8")
 		ctx.Resp.WriteHeader(503)

pkg/services/sqlstore/health.go

@@ -1,15 +1,21 @@
 package sqlstore
 
 import (
+	"context"
+
 	"github.com/grafana/grafana/pkg/bus"
 	"github.com/grafana/grafana/pkg/models"
 )
 
 func init() {
-	bus.AddHandler("sql", GetDBHealthQuery)
+	bus.AddHandlerCtx("sql", GetDBHealthQuery)
 }
 
-func GetDBHealthQuery(query *models.GetDBHealthQuery) error {
-	_, err := x.Exec("SELECT 1")
-	return err
+// GetDBHealthQuery executes a query to check
+// the availability of the database.
+func GetDBHealthQuery(ctx context.Context, query *models.GetDBHealthQuery) error {
+	return withDbSession(ctx, x, func(session *DBSession) error {
+		_, err := session.Exec("SELECT 1")
+		return err
+	})
 }

pkg/services/sqlstore/health_test.go

@@ -4,6 +4,7 @@
 package sqlstore
 
 import (
+	"context"
 	"testing"
 
 	"github.com/grafana/grafana/pkg/models"
@@ -14,6 +15,6 @@ func TestGetDBHealthQuery(t *testing.T) {
 	InitTestDB(t)
 
 	query := models.GetDBHealthQuery{}
-	err := GetDBHealthQuery(&query)
+	err := GetDBHealthQuery(context.Background(), &query)
 	require.NoError(t, err)
 }