From b0f28b19213ee1832a6651522ee00a947f2a8bf0 Mon Sep 17 00:00:00 2001 From: Ashish Kumar Jain Date: Wed, 28 May 2025 17:10:40 +0530 Subject: [PATCH] feat(health): enhance health check to report database connection status --- internal/api/handlers/v0/health.go | 12 +++++-- internal/api/handlers/v0/health_test.go | 48 +++++++++++++++++++++++-- internal/api/router/v0.go | 2 +- internal/docs/swagger.yaml | 13 +++++++ scripts/test_endpoints.sh | 7 +++- 5 files changed, 76 insertions(+), 6 deletions(-) diff --git a/internal/api/handlers/v0/health.go b/internal/api/handlers/v0/health.go index 3dd78924..c1b635f2 100644 --- a/internal/api/handlers/v0/health.go +++ b/internal/api/handlers/v0/health.go @@ -6,6 +6,7 @@ import ( "net/http" "github.com/modelcontextprotocol/registry/internal/config" + "github.com/modelcontextprotocol/registry/internal/service" ) type HealthResponse struct { @@ -14,11 +15,18 @@ type HealthResponse struct { } // HealthHandler returns a handler for health check endpoint -func HealthHandler(cfg *config.Config) http.HandlerFunc { +func HealthHandler(cfg *config.Config, registry service.RegistryService) http.HandlerFunc { return func(w http.ResponseWriter, _ *http.Request) { + status := "ok" + // Try a lightweight DB operation + _, _, err := registry.List("", 1) + if err != nil { + status = "db_error" + w.WriteHeader(http.StatusServiceUnavailable) + } w.Header().Set("Content-Type", "application/json") if err := json.NewEncoder(w).Encode(HealthResponse{ - Status: "ok", + Status: status, GitHubClientID: cfg.GithubClientID, }); err != nil { http.Error(w, "Failed to encode response", http.StatusInternalServerError) diff --git a/internal/api/handlers/v0/health_test.go b/internal/api/handlers/v0/health_test.go index baae604e..b916058c 100644 --- a/internal/api/handlers/v0/health_test.go +++ b/internal/api/handlers/v0/health_test.go @@ -9,14 +9,36 @@ import ( v0 "github.com/modelcontextprotocol/registry/internal/api/handlers/v0" "github.com/modelcontextprotocol/registry/internal/config" + "github.com/modelcontextprotocol/registry/internal/model" "github.com/stretchr/testify/assert" ) +// fakeDBRegistryService is a test double for service.RegistryService +type fakeDBRegistryService struct { + listErr error +} + +func (f *fakeDBRegistryService) List(cursor string, limit int) ([]model.Server, string, error) { + return nil, "", f.listErr +} + +// Implement other methods as no-ops or return zero values if needed for the interface +func (f *fakeDBRegistryService) GetByID(id string) (*model.ServerDetail, error) { + return nil, nil +} +func (f *fakeDBRegistryService) Publish(serverDetail *model.ServerDetail) error { + return nil +} +func (f *fakeDBRegistryService) Close() error { + return nil +} + func TestHealthHandler(t *testing.T) { // Test cases testCases := []struct { name string config *config.Config + registry *fakeDBRegistryService expectedStatus int expectedBody v0.HealthResponse }{ @@ -25,6 +47,9 @@ func TestHealthHandler(t *testing.T) { config: &config.Config{ GithubClientID: "test-github-client-id", }, + registry: &fakeDBRegistryService{ + listErr: nil, + }, expectedStatus: http.StatusOK, expectedBody: v0.HealthResponse{ Status: "ok", @@ -36,18 +61,35 @@ func TestHealthHandler(t *testing.T) { config: &config.Config{ GithubClientID: "", }, + registry: &fakeDBRegistryService{ + listErr: nil, + }, expectedStatus: http.StatusOK, expectedBody: v0.HealthResponse{ Status: "ok", GitHubClientID: "", }, }, + { + name: "unhealthy database", + config: &config.Config{ + GithubClientID: "test-github-client-id", + }, + registry: &fakeDBRegistryService{ + listErr: assert.AnError, + }, + expectedStatus: http.StatusServiceUnavailable, + expectedBody: v0.HealthResponse{ + Status: "db_error", + GitHubClientID: "test-github-client-id", + }, + }, } for _, tc := range testCases { t.Run(tc.name, func(t *testing.T) { // Create handler with the test config - handler := v0.HealthHandler(tc.config) + handler := v0.HealthHandler(tc.config, tc.registry) // Create request req, err := http.NewRequestWithContext(context.Background(), http.MethodGet, "/health", nil) @@ -84,8 +126,10 @@ func TestHealthHandlerIntegration(t *testing.T) { cfg := &config.Config{ GithubClientID: "integration-test-client-id", } + // Use a healthy fake registry service for integration + registry := &fakeDBRegistryService{listErr: nil} - server := httptest.NewServer(v0.HealthHandler(cfg)) + server := httptest.NewServer(v0.HealthHandler(cfg, registry)) defer server.Close() // Send request to the test server diff --git a/internal/api/router/v0.go b/internal/api/router/v0.go index 6d465f99..e287dde0 100644 --- a/internal/api/router/v0.go +++ b/internal/api/router/v0.go @@ -15,7 +15,7 @@ func RegisterV0Routes( mux *http.ServeMux, cfg *config.Config, registry service.RegistryService, authService auth.Service, ) { // Register v0 endpoints - mux.HandleFunc("/v0/health", v0.HealthHandler(cfg)) + mux.HandleFunc("/v0/health", v0.HealthHandler(cfg, registry)) mux.HandleFunc("/v0/servers", v0.ServersHandler(registry)) mux.HandleFunc("/v0/servers/{id}", v0.ServersDetailHandler(registry)) mux.HandleFunc("/v0/ping", v0.PingHandler(cfg)) diff --git a/internal/docs/swagger.yaml b/internal/docs/swagger.yaml index 17aea9fe..08079d93 100644 --- a/internal/docs/swagger.yaml +++ b/internal/docs/swagger.yaml @@ -36,6 +36,19 @@ paths: github_client_id: type: string example: "your_github_client_id" + '503': + description: Database connection is unhealthy + content: + application/json: + schema: + type: object + properties: + status: + type: string + example: db_error + github_client_id: + type: string + example: "your_github_client_id" /v0/ping: get: tags: diff --git a/scripts/test_endpoints.sh b/scripts/test_endpoints.sh index 2efd2f3b..b2d3686e 100755 --- a/scripts/test_endpoints.sh +++ b/scripts/test_endpoints.sh @@ -77,7 +77,12 @@ test_health() { else echo "Response:" echo "$http_response" | jq '.' 2>/dev/null || echo "$http_response" - echo "Health check failed" + # Check for db_error in the response + if echo "$http_response" | jq -e '.status == "db_error"' >/dev/null 2>&1; then + echo "Health check failed: Database connection is unhealthy!" + else + echo "Health check failed" + fi echo "-------------------------------------" return 1 fi