Tempo: Move health check requests to backend (#111668)

* move tempo health check requests to backend * fix ci and tests * add checkhealth test to backend
2025-09-30 17:57:45 +08:00
parent 3668d02650
commit 1ebfc3fea2
5 changed files with 182 additions and 106 deletions
@@ -13,6 +13,7 @@ import (
 )

 var (
+	_ backend.CheckHealthHandler  = (*Datasource)(nil)
 	_ backend.QueryDataHandler    = (*Datasource)(nil)
 	_ backend.StreamHandler       = (*Datasource)(nil)
 	_ backend.CallResourceHandler = (*Datasource)(nil)
@@ -28,6 +29,10 @@ func NewDatasource(c context.Context, b backend.DataSourceInstanceSettings) (ins
 	}, nil
 }

+func (d *Datasource) CheckHealth(ctx context.Context, req *backend.CheckHealthRequest) (*backend.CheckHealthResult, error) {
+	return d.Service.CheckHealth(ctx, req)
+}
+
 func (d *Datasource) QueryData(ctx context.Context, req *backend.QueryDataRequest) (*backend.QueryDataResponse, error) {
 	return d.Service.QueryData(ctx, req)
 }
@@ -2,6 +2,8 @@ package tempo

 import (
 	"context"
+	"encoding/json"
+	"errors"
 	"fmt"
 	"io"
 	"net/http"
@@ -156,6 +158,112 @@ func (s *Service) CallResource(ctx context.Context, req *backend.CallResourceReq
 	return s.resourceHandler.CallResource(ctx, req, sender)
 }

+func (s *Service) CheckHealth(ctx context.Context, req *backend.CheckHealthRequest) (*backend.CheckHealthResult, error) {
+	var streamingEnabled bool
+	var jsonData map[string]interface{}
+
+	pluginCtx := backend.PluginConfigFromContext(ctx)
+	dsInfo, err := s.getDSInfo(ctx, pluginCtx)
+	if err != nil {
+		return &backend.CheckHealthResult{
+			Status:  backend.HealthStatusError,
+			Message: err.Error(),
+		}, nil
+	}
+
+	if pluginCtx.DataSourceInstanceSettings != nil && pluginCtx.DataSourceInstanceSettings.JSONData != nil {
+		if err := json.Unmarshal(pluginCtx.DataSourceInstanceSettings.JSONData, &jsonData); err == nil {
+			if streaming, ok := jsonData["streamingEnabled"].(map[string]interface{}); ok {
+				if searchEnabled, ok := streaming["search"].(bool); ok && searchEnabled {
+					streamingEnabled = true
+				}
+			}
+		}
+	}
+
+	if streamingEnabled {
+		if dsInfo.StreamingClient == nil {
+			return &backend.CheckHealthResult{
+				Status:  backend.HealthStatusError,
+				Message: "Streaming client is not available",
+			}, nil
+		}
+
+		currentTime := time.Now()
+		queryStartTime := currentTime.Add(-15 * time.Minute)
+		searchRequest := &tempopb.SearchRequest{
+			Query: "{}",
+			Start: uint32(queryStartTime.Unix()),
+			End:   uint32(currentTime.Unix()),
+			Limit: 1,
+		}
+
+		streamingConnection, err := dsInfo.StreamingClient.Search(ctx, searchRequest)
+		if err != nil {
+			return &backend.CheckHealthResult{
+				Status:  backend.HealthStatusError,
+				Message: err.Error(),
+			}, nil
+		}
+
+		_, err = streamingConnection.Recv()
+		if err != nil && !errors.Is(err, io.EOF) {
+			return &backend.CheckHealthResult{
+				Status:  backend.HealthStatusError,
+				Message: err.Error(),
+			}, nil
+		}
+
+		return &backend.CheckHealthResult{
+			Status:  backend.HealthStatusOk,
+			Message: "Data source is working. Streaming test succeeded.",
+		}, nil
+	}
+
+	parsedURL, err := url.Parse(dsInfo.URL)
+	if err != nil {
+		return &backend.CheckHealthResult{
+			Status:  backend.HealthStatusError,
+			Message: err.Error(),
+		}, nil
+	}
+
+	parsedURL.Path = path.Join(parsedURL.Path, "api/echo")
+	httpReq, err := http.NewRequestWithContext(ctx, "GET", parsedURL.String(), nil)
+	if err != nil {
+		return &backend.CheckHealthResult{
+			Status:  backend.HealthStatusError,
+			Message: err.Error(),
+		}, nil
+	}
+
+	resp, err := dsInfo.HTTPClient.Do(httpReq)
+	if err != nil {
+		return &backend.CheckHealthResult{
+			Status:  backend.HealthStatusError,
+			Message: err.Error(),
+		}, nil
+	}
+
+	defer func() {
+		if err := resp.Body.Close(); err != nil {
+			s.logger.Warn("Failed to close response body", "error", err)
+		}
+	}()
+
+	if resp.StatusCode != 200 {
+		return &backend.CheckHealthResult{
+			Status:  backend.HealthStatusError,
+			Message: fmt.Sprintf("Tempo echo endpoint returned status %d", resp.StatusCode),
+		}, nil
+	}
+
+	return &backend.CheckHealthResult{
+		Status:  backend.HealthStatusOk,
+		Message: "Data source is working",
+	}, nil
+}
+
 // handleTags handles requests to /tags resource
 func (s *Service) handleTags(rw http.ResponseWriter, req *http.Request) {
 	s.proxyToTempo(rw, req, "api/v2/search/tags")
@@ -0,0 +1,67 @@
+package tempo
+
+import (
+	"context"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
+	"github.com/grafana/grafana-plugin-sdk-go/backend/datasource"
+	"github.com/grafana/grafana-plugin-sdk-go/backend/instancemgmt"
+	"github.com/stretchr/testify/assert"
+)
+
+func TestCheckHealth(t *testing.T) {
+	tests := []struct {
+		name            string
+		httpStatusCode  int
+		expectedStatus  backend.HealthStatus
+		expectedMessage string
+	}{
+		{
+			name:            "successful health check",
+			httpStatusCode:  200,
+			expectedStatus:  backend.HealthStatusOk,
+			expectedMessage: "Data source is working",
+		},
+		{
+			name:            "http error",
+			httpStatusCode:  500,
+			expectedStatus:  backend.HealthStatusError,
+			expectedMessage: "Tempo echo endpoint returned status 500",
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+				w.WriteHeader(tt.httpStatusCode)
+			}))
+			defer server.Close()
+
+			pluginCtx := backend.PluginContext{
+				DataSourceInstanceSettings: &backend.DataSourceInstanceSettings{
+					URL: server.URL,
+				},
+			}
+
+			im := datasource.NewInstanceManager(func(ctx context.Context, settings backend.DataSourceInstanceSettings) (instancemgmt.Instance, error) {
+				dsInfo := &DatasourceInfo{
+					URL:             server.URL,
+					HTTPClient:      server.Client(),
+					StreamingClient: nil,
+				}
+				return dsInfo, nil
+			})
+
+			service := &Service{im: im}
+			ctx := backend.WithPluginContext(context.Background(), pluginCtx)
+			result, err := service.CheckHealth(ctx, &backend.CheckHealthRequest{})
+
+			assert.NoError(t, err)
+			assert.Equal(t, tt.expectedStatus, result.Status)
+			assert.Contains(t, result.Message, tt.expectedMessage)
+		})
+	}
+}