Skip to content

feat(health): enhance health check to report database connection status #109

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 10 additions & 2 deletions internal/api/handlers/v0/health.go
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@ import (
"net/http"

"github.com/modelcontextprotocol/registry/internal/config"
"github.com/modelcontextprotocol/registry/internal/service"
)

type HealthResponse struct {
Expand All @@ -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)
Expand Down
48 changes: 46 additions & 2 deletions internal/api/handlers/v0/health_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}{
Expand All @@ -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",
Expand All @@ -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)
Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion internal/api/router/v0.go
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
13 changes: 13 additions & 0 deletions internal/docs/swagger.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down
7 changes: 6 additions & 1 deletion scripts/test_endpoints.sh
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down