Tempo: convert to backend data source (#31618)

* Tempo: Support opentelemetry response

* Tempo: convert Tempo to backend data source

* Update data source test

* Fix lint issues

* Apply suggestions from code review

Co-authored-by: Arve Knudsen <arve.knudsen@gmail.com>

* Return error when trace not found

Co-authored-by: Arve Knudsen <arve.knudsen@gmail.com>

go.mod

@@ -52,6 +52,7 @@ require (
 	github.com/hashicorp/go-version v1.2.1
 	github.com/inconshreveable/log15 v0.0.0-20180818164646-67afb5ed74ec
 	github.com/influxdata/influxdb-client-go/v2 v2.2.0
+	github.com/jaegertracing/jaeger v1.22.0
 	github.com/jmespath/go-jmespath v0.4.0
 	github.com/jonboulle/clockwork v0.2.2 // indirect
 	github.com/json-iterator/go v1.1.10
@@ -67,7 +68,6 @@ require (
 	github.com/prometheus/client_golang v1.9.0
 	github.com/prometheus/client_model v0.2.0
 	github.com/prometheus/common v0.18.0
-	github.com/quasilyte/go-ruleguard/dsl/fluent v0.0.0-20201222093424-5d7e62a465d3 // indirect
 	github.com/robfig/cron v0.0.0-20180505203441-b41be1df6967
 	github.com/robfig/cron/v3 v3.0.1
 	github.com/russellhaering/goxmldsig v1.1.0
@@ -82,8 +82,9 @@ require (
 	github.com/weaveworks/common v0.0.0-20201119133501-0619918236ec
 	github.com/xorcare/pointer v1.1.0
 	github.com/yudai/gojsondiff v1.0.0
-	golang.org/x/crypto v0.0.0-20201208171446-5f87f3452ae9
-	golang.org/x/net v0.0.0-20201224014010-6772e930b67b
+	go.opentelemetry.io/collector v0.21.0
+	golang.org/x/crypto v0.0.0-20201221181555-eec23a3978ad
+	golang.org/x/net v0.0.0-20210119194325-5f4716e94777
 	golang.org/x/oauth2 v0.0.0-20210113205817-d3ed898aa8a3
 	golang.org/x/sync v0.0.0-20201207232520-09787c993a3a
 	golang.org/x/time v0.0.0-20200630173020-3af7569d3a1e

pkg/cmd/grafana-server/main.go

@@ -32,6 +32,7 @@ import (
 	_ "github.com/grafana/grafana/pkg/tsdb/opentsdb"
 	_ "github.com/grafana/grafana/pkg/tsdb/postgres"
 	_ "github.com/grafana/grafana/pkg/tsdb/prometheus"
+	_ "github.com/grafana/grafana/pkg/tsdb/tempo"
 	_ "github.com/grafana/grafana/pkg/tsdb/testdatasource"
 )
 

pkg/tsdb/tempo/tempo.go

@@ -0,0 +1,127 @@
+package tempo
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"io/ioutil"
+	"net/http"
+
+	"github.com/grafana/grafana-plugin-sdk-go/data"
+	"github.com/grafana/grafana/pkg/infra/log"
+	"github.com/grafana/grafana/pkg/models"
+	"github.com/grafana/grafana/pkg/tsdb"
+
+	jaeger "github.com/jaegertracing/jaeger/model"
+	jaeger_json "github.com/jaegertracing/jaeger/model/converter/json"
+
+	ot_pdata "go.opentelemetry.io/collector/consumer/pdata"
+	ot_jaeger "go.opentelemetry.io/collector/translator/trace/jaeger"
+)
+
+type tempoExecutor struct {
+	httpClient *http.Client
+}
+
+func newTempoExecutor(dsInfo *models.DataSource) (tsdb.TsdbQueryEndpoint, error) {
+	httpClient, err := dsInfo.GetHttpClient()
+	if err != nil {
+		return nil, err
+	}
+
+	return &tempoExecutor{
+		httpClient: httpClient,
+	}, nil
+}
+
+var (
+	plog log.Logger
+)
+
+func init() {
+	plog = log.New("tsdb.tempo")
+	tsdb.RegisterTsdbQueryEndpoint("tempo", newTempoExecutor)
+}
+
+func (e *tempoExecutor) Query(ctx context.Context, dsInfo *models.DataSource, tsdbQuery *tsdb.TsdbQuery) (*tsdb.Response, error) {
+	result := &tsdb.Response{
+		Results: map[string]*tsdb.QueryResult{},
+	}
+	refID := tsdbQuery.Queries[0].RefId
+	queryResult := &tsdb.QueryResult{}
+	result.Results[refID] = queryResult
+
+	traceID := tsdbQuery.Queries[0].Model.Get("query").MustString("")
+
+	plog.Debug("Querying tempo with traceID", "traceID", traceID)
+
+	req, err := http.NewRequestWithContext(ctx, "GET", dsInfo.Url+"/api/traces/"+traceID, nil)
+	if err != nil {
+		return nil, err
+	}
+
+	req.Header.Set("Accept", "application/protobuf")
+
+	resp, err := e.httpClient.Do(req)
+	if err != nil {
+		return nil, fmt.Errorf("failed get to tempo: %w", err)
+	}
+
+	defer func() {
+		if err := resp.Body.Close(); err != nil {
+			plog.Warn("failed to close response body", "err", err)
+		}
+	}()
+
+	if resp.StatusCode == http.StatusNotFound {
+		queryResult.Error = fmt.Errorf("failed to get trace: %s", traceID)
+		return result, nil
+	}
+
+	body, err := ioutil.ReadAll(resp.Body)
+	if err != nil {
+		return nil, err
+	}
+
+	otTrace := ot_pdata.NewTraces()
+	err = otTrace.FromOtlpProtoBytes(body)
+	if err != nil {
+		return nil, fmt.Errorf("failed to convert tempo response to Otlp: %w", err)
+	}
+
+	jaegerBatches, err := ot_jaeger.InternalTracesToJaegerProto(otTrace)
+	if err != nil {
+		return nil, fmt.Errorf("failed to translate to jaegerBatches %v: %w", traceID, err)
+	}
+
+	jaegerTrace := &jaeger.Trace{
+		Spans:      []*jaeger.Span{},
+		ProcessMap: []jaeger.Trace_ProcessMapping{},
+	}
+
+	// otel proto conversion doesn't set jaeger processes
+	for _, batch := range jaegerBatches {
+		for _, s := range batch.Spans {
+			s.Process = batch.Process
+		}
+
+		jaegerTrace.Spans = append(jaegerTrace.Spans, batch.Spans...)
+		jaegerTrace.ProcessMap = append(jaegerTrace.ProcessMap, jaeger.Trace_ProcessMapping{
+			Process:   *batch.Process,
+			ProcessID: batch.Process.ServiceName,
+		})
+	}
+	jsonTrace := jaeger_json.FromDomain(jaegerTrace)
+
+	traceBytes, err := json.Marshal(jsonTrace)
+	if err != nil {
+		return nil, fmt.Errorf("failed to json.Marshal trace \"%s\" :%w", traceID, err)
+	}
+
+	frames := []*data.Frame{
+		{Name: "Traces", RefID: refID, Fields: []*data.Field{data.NewField("trace", nil, []string{string(traceBytes)})}},
+	}
+	queryResult.Dataframes = tsdb.NewDecodedDataFrames(frames)
+
+	return result, nil
+}

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

@@ -1,82 +1,39 @@
-import { TempoDatasource, TempoQuery } from './datasource';
-import { DataQueryRequest, DataSourceInstanceSettings, FieldType, PluginType, dateTime } from '@grafana/data';
-import { BackendSrv, BackendSrvRequest, getBackendSrv, setBackendSrv } from '@grafana/runtime';
+import { DataSourceInstanceSettings, FieldType, MutableDataFrame, PluginType } from '@grafana/data';
+import { backendSrv } from 'app/core/services/backend_srv';
+import { of } from 'rxjs';
+import { createFetchResponse } from 'test/helpers/createFetchResponse';
+import { TempoDatasource } from './datasource';
+
+jest.mock('../../../../../packages/grafana-runtime/src/services/backendSrv.ts', () => ({
+  getBackendSrv: () => backendSrv,
+}));
+
+jest.mock('../../../../../packages/grafana-runtime/src/utils/queryResponse.ts', () => ({
+  toDataQueryResponse: (resp: any) => resp,
+}));
+
+describe('Tempo data source', () => {
+  beforeEach(() => {
+    jest.clearAllMocks();
+  });
 
-describe('JaegerDatasource', () => {
   it('returns trace when queried', async () => {
-    await withMockedBackendSrv(makeBackendSrvMock('12345'), async () => {
-      const ds = new TempoDatasource(defaultSettings);
-      const response = await ds.query(defaultQuery).toPromise();
-      const field = response.data[0].fields[0];
-      expect(field.name).toBe('trace');
-      expect(field.type).toBe(FieldType.trace);
-      expect(field.values.get(0)).toEqual({
-        traceId: '12345',
-      });
-    });
-  });
-
-  it('returns trace when traceId with special characters is queried', async () => {
-    await withMockedBackendSrv(makeBackendSrvMock('a/b'), async () => {
-      const ds = new TempoDatasource(defaultSettings);
-      const query = {
-        ...defaultQuery,
-        targets: [
-          {
-            query: 'a/b',
-            refId: '1',
-          },
-        ],
-      };
-      const response = await ds.query(query).toPromise();
-      const field = response.data[0].fields[0];
-      expect(field.name).toBe('trace');
-      expect(field.type).toBe(FieldType.trace);
-      expect(field.values.get(0)).toEqual({
-        traceId: 'a/b',
-      });
-    });
-  });
-
-  it('returns empty response if trace id is not specified', async () => {
+    const responseDataFrame = new MutableDataFrame({ fields: [{ name: 'trace', values: ['{}'] }] });
+    setupBackendSrv([responseDataFrame]);
     const ds = new TempoDatasource(defaultSettings);
-    const response = await ds
-      .query({
-        ...defaultQuery,
-        targets: [],
-      })
-      .toPromise();
-    const field = response.data[0].fields[0];
-    expect(field.name).toBe('trace');
-    expect(field.type).toBe(FieldType.trace);
-    expect(field.values.length).toBe(0);
+    await expect(ds.query({ targets: [{ query: '12345' }] } as any)).toEmitValuesWith((response) => {
+      const field = response[0].data[0].fields[0];
+      expect(field.name).toBe('trace');
+      expect(field.type).toBe(FieldType.trace);
+    });
   });
 });
 
-function makeBackendSrvMock(traceId: string) {
-  return {
-    datasourceRequest(options: BackendSrvRequest): Promise<any> {
-      expect(options.url.substr(options.url.length - 17, options.url.length)).toBe(
-        `/api/traces/${encodeURIComponent(traceId)}`
-      );
-      return Promise.resolve({
-        data: {
-          data: [
-            {
-              traceId,
-            },
-          ],
-        },
-      });
-    },
-  } as any;
-}
+function setupBackendSrv(response: any) {
+  const defaultMock = () => of(createFetchResponse(response));
 
-async function withMockedBackendSrv(srv: BackendSrv, fn: () => Promise<void>) {
-  const oldSrv = getBackendSrv();
-  setBackendSrv(srv);
-  await fn();
-  setBackendSrv(oldSrv);
+  const fetchMock = jest.spyOn(backendSrv, 'fetch');
+  fetchMock.mockImplementation(defaultMock);
 }
 
 const defaultSettings: DataSourceInstanceSettings = {
@@ -94,26 +51,3 @@ const defaultSettings: DataSourceInstanceSettings = {
   },
   jsonData: {},
 };
-
-const defaultQuery: DataQueryRequest<TempoQuery> = {
-  requestId: '1',
-  dashboardId: 0,
-  interval: '0',
-  intervalMs: 10,
-  panelId: 0,
-  scopedVars: {},
-  range: {
-    from: dateTime().subtract(1, 'h'),
-    to: dateTime(),
-    raw: { from: '1h', to: 'now' },
-  },
-  timezone: 'browser',
-  app: 'explore',
-  startTime: 0,
-  targets: [
-    {
-      query: '12345',
-      refId: '1',
-    },
-  ],
-};

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

@@ -1,121 +1,63 @@
 import {
-  dateMath,
-  DateTime,
-  MutableDataFrame,
-  DataSourceApi,
-  DataSourceInstanceSettings,
+  DataFrame,
+  DataQuery,
   DataQueryRequest,
   DataQueryResponse,
-  DataQuery,
+  DataSourceInstanceSettings,
   FieldType,
+  MutableDataFrame,
 } from '@grafana/data';
-import { getBackendSrv, BackendSrvRequest } from '@grafana/runtime';
-import { Observable, from, of } from 'rxjs';
+import { DataSourceWithBackend } from '@grafana/runtime';
+import { Observable } from 'rxjs';
 import { map } from 'rxjs/operators';
 
-import { getTimeSrv, TimeSrv } from 'app/features/dashboard/services/TimeSrv';
-import { serializeParams } from 'app/core/utils/fetch';
-
 export type TempoQuery = {
   query: string;
 } & DataQuery;
 
-export class TempoDatasource extends DataSourceApi<TempoQuery> {
-  constructor(private instanceSettings: DataSourceInstanceSettings, private readonly timeSrv: TimeSrv = getTimeSrv()) {
+export class TempoDatasource extends DataSourceWithBackend<TempoQuery> {
+  constructor(instanceSettings: DataSourceInstanceSettings) {
     super(instanceSettings);
   }
 
-  async metadataRequest(url: string, params?: Record<string, any>): Promise<any> {
-    const res = await this._request(url, params, { hideFromInspector: true }).toPromise();
-    return res.data.data;
-  }
-
   query(options: DataQueryRequest<TempoQuery>): Observable<DataQueryResponse> {
-    // At this moment we expect only one target. In case we somehow change the UI to be able to show multiple
-    // traces at one we need to change this.
-    const id = options.targets[0]?.query;
-    if (id) {
-      return this._request(`/api/traces/${encodeURIComponent(id)}`).pipe(
-        map((response) => {
-          return {
-            data: [
-              new MutableDataFrame({
-                fields: [
-                  {
-                    name: 'trace',
-                    type: FieldType.trace,
-                    values: response?.data?.data || [],
-                  },
-                ],
-                meta: {
-                  preferredVisualisationType: 'trace',
+    return super.query(options).pipe(
+      map((response) => {
+        if (response.error) {
+          return response;
+        }
+
+        return {
+          data: [
+            new MutableDataFrame({
+              fields: [
+                {
+                  name: 'trace',
+                  type: FieldType.trace,
+                  values: [JSON.parse((response.data as DataFrame[])[0].fields[0].values.get(0))],
                 },
-              }),
-            ],
-          };
-        })
-      );
-    } else {
-      return of({
-        data: [
-          new MutableDataFrame({
-            fields: [
-              {
-                name: 'trace',
-                type: FieldType.trace,
-                values: [],
+              ],
+              meta: {
+                preferredVisualisationType: 'trace',
               },
-            ],
-            meta: {
-              preferredVisualisationType: 'trace',
-            },
-          }),
-        ],
-      });
-    }
+            }),
+          ],
+        };
+      })
+    );
   }
 
   async testDatasource(): Promise<any> {
-    try {
-      await this._request(`/api/traces/random`).toPromise();
-    } catch (e) {
-      // If all went well this request will get back with 400 - Bad request
-      if (e?.status !== 400) {
-        throw e;
-      }
-    }
-    return { status: 'success', message: 'Data source is working' };
-  }
+    const response = await super.query({ targets: [{ query: '', refId: 'A' }] } as any).toPromise();
 
-  getTimeRange(): { start: number; end: number } {
-    const range = this.timeSrv.timeRange();
-    return {
-      start: getTime(range.from, false),
-      end: getTime(range.to, true),
-    };
+    if (!response.error?.message?.startsWith('failed to get trace')) {
+      return { status: 'error', message: 'Data source is not working' };
+    }
+
+    return { status: 'success', message: 'Data source is working' };
   }
 
   getQueryDisplayText(query: TempoQuery) {
     return query.query;
   }
-
-  private _request(apiUrl: string, data?: any, options?: Partial<BackendSrvRequest>): Observable<Record<string, any>> {
-    // Hack for proxying metadata requests
-    const baseUrl = `/api/datasources/proxy/${this.instanceSettings.id}`;
-    const params = data ? serializeParams(data) : '';
-    const url = `${baseUrl}${apiUrl}${params.length ? `?${params}` : ''}`;
-    const req = {
-      ...options,
-      url,
-    };
-
-    return from(getBackendSrv().datasourceRequest(req));
-  }
-}
-
-function getTime(date: string | DateTime, roundUp: boolean) {
-  if (typeof date === 'string') {
-    date = dateMath.parse(date, roundUp)!;
-  }
-  return date.valueOf() * 1000;
 }