Add optional time window for GetTrace & SearchTrace

Signed-off-by: Zhang Xin <[email protected]>

Author: rim99

cmd/query/app/http_handler.go

@@ -231,7 +231,12 @@ func (aH *APIHandler) search(w http.ResponseWriter, r *http.Request) {
 	var uiErrors []structuredError
 	var tracesFromStorage []*model.Trace
 	if len(tQuery.traceIDs) > 0 {
-		tracesFromStorage, uiErrors, err = aH.tracesByIDs(r.Context(), tQuery.traceIDs)
+		tracesFromStorage, uiErrors, err = aH.tracesByIDs(
+			r.Context(),
+			tQuery.traceIDs,
+			tQuery.StartTimeMin,
+			tQuery.StartTimeMax,
+		)
 		if aH.handleError(w, err, http.StatusInternalServerError) {
 			return
 		}
@@ -262,13 +267,14 @@ func (aH *APIHandler) tracesToResponse(traces []*model.Trace, adjust bool, uiErr
 	}
 }
 
-func (aH *APIHandler) tracesByIDs(ctx context.Context, traceIDs []model.TraceID) ([]*model.Trace, []structuredError, error) {
+func (aH *APIHandler) tracesByIDs(ctx context.Context, traceIDs []model.TraceID, startTime time.Time, endTime time.Time) ([]*model.Trace, []structuredError, error) {
 	var traceErrors []structuredError
 	retMe := make([]*model.Trace, 0, len(traceIDs))
 	for _, traceID := range traceIDs {
-		// TODO: add start time & end time
 		query := spanstore.GetTraceParameters{
-			TraceID: traceID,
+			TraceID:   traceID,
+			StartTime: startTime,
+			EndTime:   endTime,
 		}
 		if trc, err := aH.queryService.GetTrace(ctx, query); err != nil {
 			if !errors.Is(err, spanstore.ErrTraceNotFound) {
@@ -433,10 +439,21 @@ func (aH *APIHandler) getTrace(w http.ResponseWriter, r *http.Request) {
 	if !ok {
 		return
 	}
-	// TODO: add start time & end time
 	query := spanstore.GetTraceParameters{
 		TraceID: traceID,
 	}
+	startTime := r.FormValue(startTimeParam)
+	if startTime != "" {
+		// No need to handle err here, since time window is optional
+		t, _ := aH.queryParser.parseTime(r, startTimeParam, time.Microsecond)
+		query.StartTime = t
+	}
+	endTime := r.FormValue(endTimeParam)
+	if endTime != "" {
+		// No need to handle err here, since time window is optional
+		t, _ := aH.queryParser.parseTime(r, endTimeParam, time.Microsecond)
+		query.EndTime = t
+	}
 	trc, err := aH.queryService.GetTrace(r.Context(), query)
 	if errors.Is(err, spanstore.ErrTraceNotFound) {
 		aH.handleError(w, err, http.StatusNotFound)
@@ -466,9 +483,13 @@ func (aH *APIHandler) archiveTrace(w http.ResponseWriter, r *http.Request) {
 	}
 
 	// QueryService.ArchiveTrace can now archive this traceID.
-	// TODO: add start time & end time
+	// No need to handle err here, since time window is optional
+	startTime, _ := aH.queryParser.parseTime(r, startTimeParam, time.Microsecond)
+	endTime, _ := aH.queryParser.parseTime(r, endTimeParam, time.Microsecond)
 	query := spanstore.GetTraceParameters{
-		TraceID: traceID,
+		TraceID:   traceID,
+		StartTime: startTime,
+		EndTime:   endTime,
 	}
 	err := aH.queryService.ArchiveTrace(r.Context(), query)
 	if errors.Is(err, spanstore.ErrTraceNotFound) {

cmd/query/app/http_handler_test.go

@@ -198,6 +198,22 @@ func TestGetTraceDedupeSuccess(t *testing.T) {
 	}
 }
 
+func TestGetTraceWithTimeWindowSuccess(t *testing.T) {
+	ts := initializeTestServer(t)
+	start_time := time.Date(1970, time.January, 1, 0, 0, 0, 1000, time.UTC).Local()
+	end_time := time.Date(1970, time.January, 1, 0, 0, 0, 2000, time.UTC).Local()
+	ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), spanstore.GetTraceParameters{
+		TraceID:   model.TraceID{High: 0, Low: 0x123456},
+		StartTime: start_time,
+		EndTime:   end_time,
+	}).Return(mockTrace, nil).Once()
+
+	var response structuredResponse
+	err := getJSON(ts.server.URL+`/api/traces/123456?start=1&end=2`, &response)
+	require.NoError(t, err)
+	assert.Empty(t, response.Errors)
+}
+
 func TestLogOnServerError(t *testing.T) {
 	zapCore, logs := observer.New(zap.InfoLevel)
 	logger := zap.New(zapCore)
@@ -411,6 +427,24 @@ func TestSearchByTraceIDSuccess(t *testing.T) {
 	assert.Len(t, response.Data, 2)
 }
 
+func TestSearchByTraceIDWithTimeWindowSuccess(t *testing.T) {
+	ts := initializeTestServer(t)
+	expectedTraceId, _ := model.TraceIDFromString("1")
+	expectedQuery := spanstore.GetTraceParameters{
+		TraceID:   expectedTraceId,
+		StartTime: time.UnixMicro(1),
+		EndTime:   time.UnixMicro(2),
+	}
+	ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), expectedQuery).
+		Return(mockTrace, nil)
+
+	var response structuredResponse
+	err := getJSON(ts.server.URL+`/api/traces?traceID=1&start=1&end=2`, &response)
+	require.NoError(t, err)
+	assert.Empty(t, response.Errors)
+	assert.Len(t, response.Data, 1)
+}
+
 func TestSearchByTraceIDSuccessWithArchive(t *testing.T) {
 	archiveReadMock := &spanstoremocks.Reader{}
 	ts := initializeTestServerWithOptions(t, &tenancy.Manager{}, querysvc.QueryServiceOptions{
@@ -428,6 +462,29 @@ func TestSearchByTraceIDSuccessWithArchive(t *testing.T) {
 	assert.Len(t, response.Data, 2)
 }
 
+func TestSearchByTraceIDSuccessWithArchiveAndTimeWindow(t *testing.T) {
+	archiveReadMock := &spanstoremocks.Reader{}
+	ts := initializeTestServerWithOptions(t, &tenancy.Manager{}, querysvc.QueryServiceOptions{
+		ArchiveSpanReader: archiveReadMock,
+	})
+	expectedTraceId, _ := model.TraceIDFromString("1")
+	expectedQuery := spanstore.GetTraceParameters{
+		TraceID:   expectedTraceId,
+		StartTime: time.UnixMicro(1),
+		EndTime:   time.UnixMicro(2),
+	}
+	ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), expectedQuery).
+		Return(nil, spanstore.ErrTraceNotFound)
+	archiveReadMock.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), expectedQuery).
+		Return(mockTrace, nil)
+
+	var response structuredResponse
+	err := getJSON(ts.server.URL+`/api/traces?traceID=1&start=1&end=2`, &response)
+	require.NoError(t, err)
+	assert.Empty(t, response.Errors)
+	assert.Len(t, response.Data, 1)
+}
+
 func TestSearchByTraceIDNotFound(t *testing.T) {
 	ts := initializeTestServer(t)
 	ts.spanReader.On("GetTrace", mock.AnythingOfType("*context.valueCtx"), mock.AnythingOfType("spanstore.GetTraceParameters")).