Add ability to archive traces (#57)

cmd/query/app/handler.go

@@ -50,6 +50,10 @@ const (
 	defaultHTTPPrefix                 = "api"
 )
 
+var (
+	errNoArchiveSpanStorage = errors.New("archive span storage was not configured")
+)
+
 // HTTPHandler handles http requests
 type HTTPHandler interface {
 	RegisterRoutes(router *mux.Router)
@@ -71,12 +75,14 @@ type structuredError struct {
 
 // APIHandler implements the query service public API by registering routes at httpPrefix
 type APIHandler struct {
-	spanReader       spanstore.Reader
-	dependencyReader dependencystore.Reader
-	adjuster         adjuster.Adjuster
-	logger           zap.Logger
-	queryParser      queryParser
-	httpPrefix       string
+	spanReader        spanstore.Reader
+	archiveSpanReader spanstore.Reader
+	archiveSpanWriter spanstore.Writer
+	dependencyReader  dependencystore.Reader
+	adjuster          adjuster.Adjuster
+	logger            zap.Logger
+	queryParser       queryParser
+	httpPrefix        string
 }
 
 // NewAPIHandler returns an APIHandler
@@ -110,6 +116,8 @@ func NewAPIHandler(spanReader spanstore.Reader, dependencyReader dependencystore
 // RegisterRoutes registers routes for this handler on the given router
 func (aH *APIHandler) RegisterRoutes(router *mux.Router) {
 	router.HandleFunc(aH.route("/traces/{%s}", traceIDParam), aH.getTrace).Methods(http.MethodGet)
+	router.HandleFunc(aH.route("/archive/{%s}", traceIDParam), aH.getArchivedTrace).Methods(http.MethodGet)
+	router.HandleFunc(aH.route("/archive/{%s}", traceIDParam), aH.archiveTrace).Methods(http.MethodPost)
 	router.HandleFunc(aH.route(`/traces`), aH.search).Methods(http.MethodGet)
 	router.HandleFunc(aH.route(`/services`), aH.getServices).Methods(http.MethodGet)
 	// TODO change the UI to use this endpoint. Requires ?service= parameter.
@@ -279,31 +287,99 @@ func (aH *APIHandler) filterDependenciesByService(
 	return filteredDependencies
 }
 
-func (aH *APIHandler) getTrace(w http.ResponseWriter, r *http.Request) {
+// Parses trace ID from URL like /traces/{trace-id}
+func (aH *APIHandler) parseTraceID(w http.ResponseWriter, r *http.Request) (model.TraceID, bool) {
 	vars := mux.Vars(r)
 	traceIDVar := vars[traceIDParam]
 	traceID, err := model.TraceIDFromString(traceIDVar)
 	if aH.handleError(w, err, http.StatusBadRequest) {
-		return
+		return traceID, false
 	}
+	return traceID, true
+}
 
-	zTrace, err := aH.spanReader.GetTrace(traceID)
+// getTrace implements the REST API /traces/{trace-id}
+func (aH *APIHandler) getTrace(w http.ResponseWriter, r *http.Request) {
+	aH.getTraceFromReader(w, r, aH.spanReader)
+}
+
+// getTraceFromReader parses trace ID from the path, loads the trace from specified Reader,
+// formats it in the UI JSON format, and responds to the client.
+func (aH *APIHandler) getTraceFromReader(w http.ResponseWriter, r *http.Request, reader spanstore.Reader) {
+	aH.withTraceFromReader(w, r, reader, func(trace *model.Trace) {
+		var uiErrors []structuredError
+		uiTrace, uiErr := aH.convertModelToUI(trace)
+		if uiErr != nil {
+			uiErrors = append(uiErrors, *uiErr)
+		}
+
+		structuredRes := structuredResponse{
+			Data: []*ui.Trace{
+				uiTrace,
+			},
+			Errors: uiErrors,
+		}
+		aH.writeJSON(w, &structuredRes)
+	})
+}
+
+// withTraceFromReader tries to load a trace from Reader and if successful
+// execute process() function passing it that trace.
+func (aH *APIHandler) withTraceFromReader(
+	w http.ResponseWriter,
+	r *http.Request,
+	reader spanstore.Reader,
+	process func(trace *model.Trace),
+) {
+	traceID, ok := aH.parseTraceID(w, r)
+	if !ok {
+		return
+	}
+	trace, err := reader.GetTrace(traceID)
+	if err == spanstore.ErrTraceNotFound {
+		aH.handleError(w, err, http.StatusBadRequest)
+		return
+	}
 	if aH.handleError(w, err, http.StatusInternalServerError) {
 		return
 	}
-	var uiErrors []structuredError
-	uiTrace, uiErr := aH.convertModelToUI(zTrace)
-	if uiErr != nil {
-		uiErrors = append(uiErrors, *uiErr)
+	process(trace)
+}
+
+// getArchivedTrace implements the REST API GET:/archive/{trace-id}
+func (aH *APIHandler) getArchivedTrace(w http.ResponseWriter, r *http.Request) {
+	if aH.archiveSpanReader == nil {
+		aH.handleError(w, errNoArchiveSpanStorage, http.StatusInternalServerError)
+		return
 	}
+	aH.getTraceFromReader(w, r, aH.archiveSpanReader)
+}
 
-	structuredRes := structuredResponse{
-		Data: []*ui.Trace{
-			uiTrace,
-		},
-		Errors: uiErrors,
+// archiveTrace implements the REST API POST:/archive/{trace-id}.
+// It reads the trace from the main Reader and saves it to archive Writer.
+func (aH *APIHandler) archiveTrace(w http.ResponseWriter, r *http.Request) {
+	if aH.archiveSpanWriter == nil {
+		aH.handleError(w, errNoArchiveSpanStorage, http.StatusInternalServerError)
+		return
 	}
-	aH.writeJSON(w, &structuredRes)
+	aH.withTraceFromReader(w, r, aH.spanReader, func(trace *model.Trace) {
+		var writeErrors []error
+		for _, span := range trace.Spans {
+			err := aH.archiveSpanWriter.WriteSpan(span)
+			if err != nil {
+				writeErrors = append(writeErrors, err)
+			}
+		}
+		err := multierror.Wrap(writeErrors)
+		if aH.handleError(w, err, http.StatusInternalServerError) {
+			return
+		}
+		structuredRes := structuredResponse{
+			Data:   []string{}, // doens't matter, just want an empty array
+			Errors: []structuredError{},
+		}
+		aH.writeJSON(w, &structuredRes)
+	})
 }
 
 func (aH *APIHandler) handleError(w http.ResponseWriter, err error, statusCode int) bool {

cmd/query/app/handler_archive_test.go

@@ -0,0 +1,91 @@
+// Copyright (c) 2017 Uber Technologies, Inc.
+//
+// Permission is hereby granted, free of charge, to any person obtaining a copy
+// of this software and associated documentation files (the "Software"), to deal
+// in the Software without restriction, including without limitation the rights
+// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+// copies of the Software, and to permit persons to whom the Software is
+// furnished to do so, subject to the following conditions:
+//
+// The above copyright notice and this permission notice shall be included in
+// all copies or substantial portions of the Software.
+//
+// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+// THE SOFTWARE.
+
+package app
+
+import (
+	"errors"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/mock"
+
+	"github.com/uber/jaeger/model"
+	spanstoremocks "github.com/uber/jaeger/storage/spanstore/mocks"
+)
+
+func TestGetArchivedTrace_NoStorage(t *testing.T) {
+	withTestServer(t, func(ts *testServer) {
+		var response structuredResponse
+		err := getJSON(ts.server.URL+"/api/archive/"+mockTraceID.String(), &response)
+		assert.EqualError(t, err,
+			`500 error from server: {"data":null,"total":0,"limit":0,"offset":0,"errors":[{"code":500,"msg":"archive span storage was not configured"}]}`+"\n",
+		)
+	})
+}
+
+func TestGetArchivedTraceSuccess(t *testing.T) {
+	traceID := model.TraceID{Low: 123456}
+	mockReader := &spanstoremocks.Reader{}
+	mockReader.On("GetTrace", mock.AnythingOfType("model.TraceID")).
+		Return(mockTrace, nil).Once()
+	withTestServer(t, func(ts *testServer) {
+		var response structuredTraceResponse
+		err := getJSON(ts.server.URL+"/api/archive/"+mockTraceID.String(), &response)
+		assert.NoError(t, err)
+		assert.Len(t, response.Errors, 0)
+		assert.Len(t, response.Traces, 1)
+		assert.Equal(t, traceID.String(), string(response.Traces[0].TraceID))
+	}, HandlerOptions.ArchiveSpanReader(mockReader))
+}
+
+func TestArchiveTrace_NoStorage(t *testing.T) {
+	withTestServer(t, func(ts *testServer) {
+		var response structuredResponse
+		err := postJSON(ts.server.URL+"/api/archive/"+mockTraceID.String(), []string{}, &response)
+		assert.EqualError(t, err, `500 error from server: {"data":null,"total":0,"limit":0,"offset":0,"errors":[{"code":500,"msg":"archive span storage was not configured"}]}`+"\n")
+	})
+}
+
+func TestArchiveTrace_Success(t *testing.T) {
+	mockWriter := &spanstoremocks.Writer{}
+	mockWriter.On("WriteSpan", mock.AnythingOfType("*model.Span")).
+		Return(nil).Times(2)
+	withTestServer(t, func(ts *testServer) {
+		ts.spanReader.On("GetTrace", mock.AnythingOfType("model.TraceID")).
+			Return(mockTrace, nil).Once()
+		var response structuredResponse
+		err := postJSON(ts.server.URL+"/api/archive/"+mockTraceID.String(), []string{}, &response)
+		assert.NoError(t, err)
+	}, HandlerOptions.ArchiveSpanWriter(mockWriter))
+}
+
+func TestArchiveTrace_WriteErrors(t *testing.T) {
+	mockWriter := &spanstoremocks.Writer{}
+	mockWriter.On("WriteSpan", mock.AnythingOfType("*model.Span")).
+		Return(errors.New("cannot save")).Times(2)
+	withTestServer(t, func(ts *testServer) {
+		ts.spanReader.On("GetTrace", mock.AnythingOfType("model.TraceID")).
+			Return(mockTrace, nil).Once()
+		var response structuredResponse
+		err := postJSON(ts.server.URL+"/api/archive/"+mockTraceID.String(), []string{}, &response)
+		assert.EqualError(t, err, `500 error from server: {"data":null,"total":0,"limit":0,"offset":0,"errors":[{"code":500,"msg":"[cannot save, cannot save]"}]}`+"\n")
+	}, HandlerOptions.ArchiveSpanWriter(mockWriter))
+}