Add support for retrieving unadjusted/raw spans (#615)

Signed-off-by: Yuri Shkuro <ys@uber.com>

cmd/query/app/handler.go

@@ -211,7 +211,7 @@ func (aH *APIHandler) search(w http.ResponseWriter, r *http.Request) {
 
 	uiTraces := make([]*ui.Trace, len(tracesFromStorage))
 	for i, v := range tracesFromStorage {
-		uiTrace, uiErr := aH.convertModelToUI(v)
+		uiTrace, uiErr := aH.convertModelToUI(v, true)
 		if uiErr != nil {
 			uiErrors = append(uiErrors, *uiErr)
 		}
@@ -275,11 +275,14 @@ func (aH *APIHandler) dependencies(w http.ResponseWriter, r *http.Request) {
 	aH.writeJSON(w, &structuredRes)
 }
 
-func (aH *APIHandler) convertModelToUI(traceFromStorage *model.Trace) (*ui.Trace, *structuredError) {
+func (aH *APIHandler) convertModelToUI(trace *model.Trace, adjust bool) (*ui.Trace, *structuredError) {
 	var errors []error
-	trace, err := aH.adjuster.Adjust(traceFromStorage)
-	if err != nil {
-		errors = append(errors, err)
+	if adjust {
+		var err error
+		trace, err = aH.adjuster.Adjust(trace)
+		if err != nil {
+			errors = append(errors, err)
+		}
 	}
 	uiTrace := uiconv.FromDomain(trace)
 	var uiError *structuredError
@@ -354,7 +357,7 @@ func (aH *APIHandler) getTraceFromReaders(
 ) {
 	aH.withTraceFromReader(w, r, reader, backupReader, func(trace *model.Trace) {
 		var uiErrors []structuredError
-		uiTrace, uiErr := aH.convertModelToUI(trace)
+		uiTrace, uiErr := aH.convertModelToUI(trace, shouldAdjust(r))
 		if uiErr != nil {
 			uiErrors = append(uiErrors, *uiErr)
 		}
@@ -369,6 +372,12 @@ func (aH *APIHandler) getTraceFromReaders(
 	})
 }
 
+func shouldAdjust(r *http.Request) bool {
+	raw := r.FormValue("raw")
+	isRaw, _ := strconv.ParseBool(raw)
+	return !isRaw
+}
+
 // withTraceFromReader tries to load a trace from Reader and if successful
 // execute process() function passing it that trace.
 func (aH *APIHandler) withTraceFromReader(

cmd/query/app/handler_test.go

@@ -29,6 +29,7 @@ import (
 	"github.com/gorilla/mux"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
+	"github.com/stretchr/testify/require"
 	"go.uber.org/zap"
 
 	"github.com/jaegertracing/jaeger/model"
@@ -138,22 +139,61 @@ func TestGetTraceSuccess(t *testing.T) {
 }
 
 func TestGetTrace(t *testing.T) {
-	reporter := jaeger.NewInMemoryReporter()
-	jaegerTracer, jaegerCloser := jaeger.NewTracer("test", jaeger.NewConstSampler(true), reporter)
-	defer jaegerCloser.Close()
+	testCases := []struct {
+		suffix      string
+		numSpanRefs int
+	}{
+		{suffix: "", numSpanRefs: 0},
+		{suffix: "?raw=true", numSpanRefs: 1}, // bad span reference is not filtered out
+		{suffix: "?raw=false", numSpanRefs: 0},
+	}
 
-	server, readMock, _ := initializeTestServer(HandlerOptions.Tracer(jaegerTracer))
-	defer server.Close()
-	readMock.On("GetTrace", model.TraceID{Low: 0x123456abc}).
-		Return(mockTrace, nil).Once()
+	makeMockTrace := func(t *testing.T) *model.Trace {
+		bytes, err := json.Marshal(mockTrace)
+		require.NoError(t, err)
+		var trace *model.Trace
+		require.NoError(t, json.Unmarshal(bytes, &trace))
+		trace.Spans[1].References = []model.SpanRef{
+			{TraceID: model.TraceID{High: 0, Low: 0}},
+		}
+		return trace
+	}
 
-	var response structuredResponse
-	err := getJSON(server.URL+`/api/traces/123456aBC`, &response) // trace ID in mixed lower/upper case
-	assert.NoError(t, err)
-	assert.Len(t, response.Errors, 0)
+	extractTraces := func(t *testing.T, response *structuredResponse) []ui.Trace {
+		var traces []ui.Trace
+		bytes, err := json.Marshal(response.Data)
+		require.NoError(t, err)
+		err = json.Unmarshal(bytes, &traces)
+		require.NoError(t, err)
+		return traces
+	}
+
+	for _, tc := range testCases {
+		testCase := tc // capture loop var
+		t.Run(testCase.suffix, func(t *testing.T) {
+			reporter := jaeger.NewInMemoryReporter()
+			jaegerTracer, jaegerCloser := jaeger.NewTracer("test", jaeger.NewConstSampler(true), reporter)
+			defer jaegerCloser.Close()
+
+			server, readMock, _ := initializeTestServer(HandlerOptions.Tracer(jaegerTracer))
+			defer server.Close()
 
-	assert.Len(t, reporter.GetSpans(), 1)
-	assert.Equal(t, "/api/traces/{traceID}", reporter.GetSpans()[0].(*jaeger.Span).OperationName())
+			readMock.On("GetTrace", model.TraceID{Low: 0x123456abc}).
+				Return(makeMockTrace(t), nil).Once()
+
+			var response structuredResponse
+			err := getJSON(server.URL+`/api/traces/123456aBC`+testCase.suffix, &response) // trace ID in mixed lower/upper case
+			assert.NoError(t, err)
+			assert.Len(t, response.Errors, 0)
+
+			assert.Len(t, reporter.GetSpans(), 1, "HTTP request was traced and span reported")
+			assert.Equal(t, "/api/traces/{traceID}", reporter.GetSpans()[0].(*jaeger.Span).OperationName())
+
+			traces := extractTraces(t, &response)
+			assert.Len(t, traces[0].Spans, 2)
+			assert.Len(t, traces[0].Spans[1].References, testCase.numSpanRefs)
+		})
+	}
 }
 
 func TestGetTraceDBFailure(t *testing.T) {

model/adjuster/bad_span_references.go

@@ -20,10 +20,10 @@ import (
 	"github.com/jaegertracing/jaeger/model"
 )
 
-// SpanReferences creates am adjuster that removes invalid span references, e.g. with traceID==0
+// SpanReferences creates an adjuster that removes invalid span references, e.g. with traceID==0
 func SpanReferences() Adjuster {
 	return Func(func(trace *model.Trace) (*model.Trace, error) {
-		adjuster := &spanReferenceAdjuster{}
+		adjuster := spanReferenceAdjuster{}
 		for _, span := range trace.Spans {
 			adjuster.adjust(span)
 		}
@@ -33,7 +33,7 @@ func SpanReferences() Adjuster {
 
 type spanReferenceAdjuster struct{}
 
-func (s *spanReferenceAdjuster) adjust(span *model.Span) *model.Span {
+func (s spanReferenceAdjuster) adjust(span *model.Span) *model.Span {
 	foundInvalid := false
 	for i := range span.References {
 		if !s.valid(&span.References[i]) {
@@ -56,6 +56,6 @@ func (s *spanReferenceAdjuster) adjust(span *model.Span) *model.Span {
 	return span
 }
 
-func (s *spanReferenceAdjuster) valid(ref *model.SpanRef) bool {
+func (s spanReferenceAdjuster) valid(ref *model.SpanRef) bool {
 	return ref.TraceID.High != 0 || ref.TraceID.Low != 0
 }