[storage][v2] Implement FindTraces in v2 factory adapter (jaegertra…

…cing#6328) ## Which problem is this PR solving? - Towards jaegertracing#5079 ## Description of the changes - This PR implements `FindTraces` in the v2 factory adapter ## How was this change tested? - Added unit tests ## Checklist - [x] I have read https://github.com/jaegertracing/jaeger/blob/master/CONTRIBUTING_GUIDELINES.md - [x] I have signed all commits - [x] I have added unit tests for the new functionality - [x] I have run lint and test steps successfully - for `jaeger`: `make lint test` - for `jaeger-ui`: `yarn lint` and `yarn test` --------- Signed-off-by: Mahad Zaryab <[email protected]>
Ayush-Vish · Dec 8, 2024 · 737b9be · 737b9be
1 parent 1788f03
commit 737b9be
Show file tree

Hide file tree

Showing 4 changed files with 186 additions and 18 deletions.
diff --git a/storage_v2/factoryadapter/reader.go b/storage_v2/factoryadapter/reader.go
@@ -7,6 +7,7 @@ import (
 	"context"
 	"errors"
 
+	model2otel "github.com/open-telemetry/opentelemetry-collector-contrib/pkg/translator/jaeger"
 	"go.opentelemetry.io/collector/pdata/pcommon"
 	"go.opentelemetry.io/collector/pdata/ptrace"
 
@@ -62,21 +63,25 @@ func (tr *TraceReader) GetOperations(ctx context.Context, query tracestore.Opera
 	return operations, nil
 }
 
-func (*TraceReader) FindTraces(_ context.Context, _ tracestore.TraceQueryParameters) ([]ptrace.Traces, error) {
-	panic("not implemented")
+func (tr *TraceReader) FindTraces(
+	ctx context.Context,
+	query tracestore.TraceQueryParameters,
+) ([]ptrace.Traces, error) {
+	t, err := tr.spanReader.FindTraces(ctx, query.ToSpanStoreQueryParameters())
+	if err != nil || t == nil {
+		return nil, err
+	}
+	otelTraces := []ptrace.Traces{}
+	for _, trace := range t {
+		batch := &model.Batch{Spans: trace.GetSpans()}
+		otelTrace, _ := model2otel.ProtoToTraces([]*model.Batch{batch})
+		otelTraces = append(otelTraces, otelTrace)
+	}
+	return otelTraces, nil
 }
 
 func (tr *TraceReader) FindTraceIDs(ctx context.Context, query tracestore.TraceQueryParameters) ([]pcommon.TraceID, error) {
-	t, err := tr.spanReader.FindTraceIDs(ctx, &spanstore.TraceQueryParameters{
-		ServiceName:   query.ServiceName,
-		OperationName: query.OperationName,
-		Tags:          query.Tags,
-		StartTimeMin:  query.StartTimeMin,
-		StartTimeMax:  query.StartTimeMax,
-		DurationMin:   query.DurationMin,
-		DurationMax:   query.DurationMax,
-		NumTraces:     query.NumTraces,
-	})
+	t, err := tr.spanReader.FindTraceIDs(ctx, query.ToSpanStoreQueryParameters())
 	if err != nil || t == nil {
 		return nil, err
 	}

diff --git a/storage_v2/factoryadapter/reader_test.go b/storage_v2/factoryadapter/reader_test.go
@@ -116,6 +116,11 @@ func TestTraceReader_GetOperationsDelegatesResponse(t *testing.T) {
 			operations:         nil,
 			expectedOperations: nil,
 		},
+		{
+			name:               "empty response",
+			operations:         []spanstore.Operation{},
+			expectedOperations: []tracestore.Operation{},
+		},
 		{
 			name:               "error response",
 			operations:         nil,
@@ -148,15 +153,122 @@ func TestTraceReader_GetOperationsDelegatesResponse(t *testing.T) {
 	}
 }
 
-func TestTraceReader_FindTracesPanics(t *testing.T) {
-	memstore := memory.NewStore()
+func TestTraceReader_FindTracesDelegatesSuccessResponse(t *testing.T) {
+	modelTraces := []*model.Trace{
+		{
+			Spans: []*model.Span{
+				{
+					TraceID:       model.NewTraceID(2, 3),
+					SpanID:        model.SpanID(1),
+					OperationName: "operation-a",
+				},
+				{
+					TraceID:       model.NewTraceID(4, 5),
+					SpanID:        model.SpanID(2),
+					OperationName: "operation-b",
+				},
+			},
+		},
+		{
+			Spans: []*model.Span{
+				{
+					TraceID:       model.NewTraceID(6, 7),
+					SpanID:        model.SpanID(3),
+					OperationName: "operation-c",
+				},
+			},
+		},
+	}
+	sr := new(spanStoreMocks.Reader)
+	now := time.Now()
+	sr.On(
+		"FindTraces",
+		mock.Anything,
+		&spanstore.TraceQueryParameters{
+			ServiceName:   "service",
+			OperationName: "operation",
+			Tags:          map[string]string{"tag-a": "val-a"},
+			StartTimeMin:  now,
+			StartTimeMax:  now.Add(time.Minute),
+			DurationMin:   time.Minute,
+			DurationMax:   time.Hour,
+			NumTraces:     10,
+		},
+	).Return(modelTraces, nil)
 	traceReader := &TraceReader{
-		spanReader: memstore,
+		spanReader: sr,
 	}
-	require.Panics(
-		t,
-		func() { traceReader.FindTraces(context.Background(), tracestore.TraceQueryParameters{}) },
+	traces, err := traceReader.FindTraces(
+		context.Background(),
+		tracestore.TraceQueryParameters{
+			ServiceName:   "service",
+			OperationName: "operation",
+			Tags:          map[string]string{"tag-a": "val-a"},
+			StartTimeMin:  now,
+			StartTimeMax:  now.Add(time.Minute),
+			DurationMin:   time.Minute,
+			DurationMax:   time.Hour,
+			NumTraces:     10,
+		},
 	)
+	require.NoError(t, err)
+	require.Len(t, traces, len(modelTraces))
+	traceASpans := traces[0].ResourceSpans().At(0).ScopeSpans().At(0).Spans()
+	traceBSpans := traces[1].ResourceSpans().At(0).ScopeSpans().At(0).Spans()
+	require.EqualValues(t, []byte{0, 0, 0, 0, 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 3}, traceASpans.At(0).TraceID())
+	require.EqualValues(t, []byte{0, 0, 0, 0, 0, 0, 0, 1}, traceASpans.At(0).SpanID())
+	require.Equal(t, "operation-a", traceASpans.At(0).Name())
+	require.EqualValues(t, []byte{0, 0, 0, 0, 0, 0, 0, 4, 0, 0, 0, 0, 0, 0, 0, 5}, traceASpans.At(1).TraceID())
+	require.EqualValues(t, []byte{0, 0, 0, 0, 0, 0, 0, 2}, traceASpans.At(1).SpanID())
+	require.Equal(t, "operation-b", traceASpans.At(1).Name())
+	require.EqualValues(t, []byte{0, 0, 0, 0, 0, 0, 0, 6, 0, 0, 0, 0, 0, 0, 0, 7}, traceBSpans.At(0).TraceID())
+	require.EqualValues(t, []byte{0, 0, 0, 0, 0, 0, 0, 3}, traceBSpans.At(0).SpanID())
+	require.Equal(t, "operation-c", traceBSpans.At(0).Name())
+}
+
+func TestTraceReader_FindTracesEdgeCases(t *testing.T) {
+	tests := []struct {
+		name           string
+		modelTraces    []*model.Trace
+		expectedTraces []ptrace.Traces
+		err            error
+	}{
+		{
+			name:           "nil response",
+			modelTraces:    nil,
+			expectedTraces: nil,
+		},
+		{
+			name:           "empty response",
+			modelTraces:    []*model.Trace{},
+			expectedTraces: []ptrace.Traces{},
+		},
+		{
+			name:           "error response",
+			modelTraces:    nil,
+			expectedTraces: nil,
+			err:            errors.New("test error"),
+		},
+	}
+	for _, test := range tests {
+		t.Run(test.name, func(t *testing.T) {
+			sr := new(spanStoreMocks.Reader)
+			sr.On(
+				"FindTraces",
+				mock.Anything,
+				mock.Anything,
+			).Return(test.modelTraces, test.err)
+			traceReader := &TraceReader{
+				spanReader: sr,
+			}
+			traces, err := traceReader.FindTraces(
+				context.Background(),
+				tracestore.TraceQueryParameters{},
+			)
+			require.ErrorIs(t, err, test.err)
+			require.Equal(t, test.expectedTraces, traces)
+		})
+	}
 }
 
 func TestTraceReader_FindTraceIDsDelegatesResponse(t *testing.T) {

diff --git a/storage_v2/tracestore/reader.go b/storage_v2/tracestore/reader.go
@@ -58,6 +58,19 @@ type TraceQueryParameters struct {
 	NumTraces     int
 }
 
+func (t *TraceQueryParameters) ToSpanStoreQueryParameters() *spanstore.TraceQueryParameters {
+	return &spanstore.TraceQueryParameters{
+		ServiceName:   t.ServiceName,
+		OperationName: t.OperationName,
+		Tags:          t.Tags,
+		StartTimeMin:  t.StartTimeMin,
+		StartTimeMax:  t.StartTimeMax,
+		DurationMin:   t.DurationMin,
+		DurationMax:   t.DurationMax,
+		NumTraces:     t.NumTraces,
+	}
+}
+
 // OperationQueryParameters contains parameters of query operations, empty spanKind means get operations for all kinds of span.
 type OperationQueryParameters struct {
 	ServiceName string

diff --git a/storage_v2/tracestore/reader_test.go b/storage_v2/tracestore/reader_test.go
@@ -0,0 +1,38 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package tracestore
+
+import (
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+
+	"github.com/jaegertracing/jaeger/storage/spanstore"
+)
+
+func TestToSpanStoreQueryParameters(t *testing.T) {
+	now := time.Now()
+	query := &TraceQueryParameters{
+		ServiceName:   "service",
+		OperationName: "operation",
+		Tags:          map[string]string{"tag-a": "val-a"},
+		StartTimeMin:  now,
+		StartTimeMax:  now.Add(time.Minute),
+		DurationMin:   time.Minute,
+		DurationMax:   time.Hour,
+		NumTraces:     10,
+	}
+	expected := &spanstore.TraceQueryParameters{
+		ServiceName:   "service",
+		OperationName: "operation",
+		Tags:          map[string]string{"tag-a": "val-a"},
+		StartTimeMin:  now,
+		StartTimeMax:  now.Add(time.Minute),
+		DurationMin:   time.Minute,
+		DurationMax:   time.Hour,
+		NumTraces:     10,
+	}
+	require.Equal(t, expected, query.ToSpanStoreQueryParameters())
+}