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

* move tempo health check requests to backend

* fix ci and tests

* add checkhealth test to backend

pkg/tsdb/tempo/standalone/datasource.go

@@ -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)
 }

pkg/tsdb/tempo/tempo.go

@@ -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")

pkg/tsdb/tempo/tempo_test.go

@@ -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)
+		})
+	}
+}

public/app/plugins/datasource/tempo/datasource.test.ts

@@ -52,17 +52,6 @@ import { createTempoDatasource } from './test/mocks';
 import { initTemplateSrv } from './test/test_utils';
 import { TempoJsonData, TempoQuery } from './types';
 
-let mockObservable: () => Observable<unknown>;
-jest.mock('@grafana/runtime', () => {
-  return {
-    ...jest.requireActual('@grafana/runtime'),
-    getBackendSrv: () => ({
-      fetch: mockObservable,
-      _request: mockObservable,
-    }),
-  };
-});
-
 describe('Tempo data source', () => {
   // Mock the console error so that running the test suite doesnt throw the error
   const origError = console.error;
@@ -339,20 +328,6 @@ describe('Tempo data source', () => {
     expect(edgesFrame.meta?.preferredVisualisationType).toBe('nodeGraph');
   });
 
-  describe('test the testDatasource function', () => {
-    it('should return a success msg if response.ok is true', async () => {
-      mockObservable = () => of({ ok: true });
-      const handleStreamingQuery = jest
-        .spyOn(TempoDatasource.prototype, 'handleStreamingQuery')
-        .mockImplementation(() => of({ data: [] }));
-
-      const ds = new TempoDatasource(defaultSettings);
-      const response = await ds.testDatasource();
-      expect(response.status).toBe('success');
-      expect(handleStreamingQuery).toHaveBeenCalled();
-    });
-  });
-
   describe('test the metadataRequest function', () => {
     it('should return the data from getResource', async () => {
       const ds = new TempoDatasource(defaultSettings);

public/app/plugins/datasource/tempo/datasource.ts

@@ -1,5 +1,5 @@
 import { groupBy } from 'lodash';
-import { EMPTY, forkJoin, from, lastValueFrom, merge, Observable, of } from 'rxjs';
+import { EMPTY, from, merge, Observable, of } from 'rxjs';
 import { catchError, concatMap, finalize, map, mergeMap, toArray } from 'rxjs/operators';
 
 import {
@@ -931,86 +931,7 @@ export class TempoDatasource extends DataSourceWithBackend<TempoQuery, TempoJson
   }
 
   async testDatasource(): Promise<TestDataSourceResponse> {
-    const observables = [];
-
-    const options: BackendSrvRequest = {
-      headers: {},
-      method: 'GET',
-      url: `${this.instanceSettings.url}/api/echo`,
-    };
-    observables.push(
-      getBackendSrv()
-        .fetch(options)
-        .pipe(
-          mergeMap(() => {
-            return of({ status: 'success', message: 'Health check succeeded' });
-          }),
-          catchError((err) => {
-            return of({
-              status: 'error',
-              message: getErrorMessage(err?.data?.message, 'Unable to connect with Tempo'),
-            });
-          })
-        )
-    );
-
-    if (this.streamingEnabled?.search) {
-      const now = new Date();
-      const from = new Date(now);
-      from.setMinutes(from.getMinutes() - 15);
-      observables.push(
-        this.handleStreamingQuery(
-          {
-            range: {
-              from: dateTime(from),
-              to: dateTime(now),
-              raw: { from: 'now-15m', to: 'now' },
-            },
-            requestId: '',
-            interval: '',
-            intervalMs: 0,
-            scopedVars: {},
-            targets: [],
-            timezone: '',
-            app: '',
-            startTime: 0,
-          },
-          [
-            {
-              datasource: this.instanceSettings,
-              limit: 1,
-              query: '{}',
-              queryType: 'traceql',
-              refId: 'A',
-              tableType: SearchTableType.Traces,
-              filters: [],
-            },
-          ],
-          '{}'
-        ).pipe(
-          mergeMap(() => {
-            return of({ status: 'success', message: 'Streaming test succeeded.' });
-          }),
-          catchError((err) => {
-            return of({
-              status: 'error',
-              message: getErrorMessage(err?.data?.message, 'Test for streaming failed, consider disabling streaming'),
-            });
-          })
-        )
-      );
-    }
-
-    return await lastValueFrom(
-      forkJoin(observables).pipe(
-        mergeMap((observableResults) => {
-          const erroredResult = observableResults.find((result) => result.status !== 'success');
-          return erroredResult
-            ? of(erroredResult)
-            : of({ status: 'success', message: 'Successfully connected to Tempo data source.' });
-        })
-      )
-    );
+    return await super.testDatasource();
   }
 
   getQueryDisplayText(query: TempoQuery) {