diff --git a/pkg/tsdb/tempo/standalone/datasource.go b/pkg/tsdb/tempo/standalone/datasource.go index f449fb98fc0..d41d68f2424 100644 --- a/pkg/tsdb/tempo/standalone/datasource.go +++ b/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) } diff --git a/pkg/tsdb/tempo/tempo.go b/pkg/tsdb/tempo/tempo.go index 0c9d904e7b1..ccb6173489b 100644 --- a/pkg/tsdb/tempo/tempo.go +++ b/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") diff --git a/pkg/tsdb/tempo/tempo_test.go b/pkg/tsdb/tempo/tempo_test.go new file mode 100644 index 00000000000..02e6c56050e --- /dev/null +++ b/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) + }) + } +} diff --git a/public/app/plugins/datasource/tempo/datasource.test.ts b/public/app/plugins/datasource/tempo/datasource.test.ts index 05c5399d8f9..737feaa4851 100644 --- a/public/app/plugins/datasource/tempo/datasource.test.ts +++ b/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; -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); diff --git a/public/app/plugins/datasource/tempo/datasource.ts b/public/app/plugins/datasource/tempo/datasource.ts index 193c6f053ff..2ded44441f6 100644 --- a/public/app/plugins/datasource/tempo/datasource.ts +++ b/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 { - 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) {