Feat: Add time window for GetTrace of anonymizer

cmd/anonymizer/app/flags.go

@@ -17,6 +17,8 @@ type Options struct {
 	HashCustomTags    bool
 	HashLogs          bool
 	HashProcess       bool
+	StartTime         int64
+	EndTime           int64
 }
 
 const (
@@ -28,6 +30,8 @@ const (
 	hashLogsFlag          = "hash-logs"
 	hashProcessFlag       = "hash-process"
 	maxSpansCount         = "max-spans-count"
+	startTime             = "start-time"
+	endTime               = "end-time"
 )
 
 // AddFlags adds flags for anonymizer main program
@@ -72,6 +76,16 @@ func (o *Options) AddFlags(command *cobra.Command) {
 		maxSpansCount,
 		-1,
 		"The maximum number of spans to anonymize")
+	command.Flags().Int64Var(
+		&o.StartTime,
+		startTime,
+		0,
+		"The start time of time window for searching trace, timestampe in unix nanoseconds")
+	command.Flags().Int64Var(
+		&o.EndTime,
+		endTime,
+		0,
+		"The end time of time window for searching trace, timestampe in unix nanoseconds")
 
 	// mark traceid flag as mandatory
 	command.MarkFlagRequired(traceIDFlag)

cmd/anonymizer/app/flags_test.go

@@ -24,6 +24,8 @@ func TestOptionsWithDefaultFlags(t *testing.T) {
 	assert.False(t, o.HashLogs)
 	assert.False(t, o.HashProcess)
 	assert.Equal(t, -1, o.MaxSpansCount)
+	assert.Equal(t, int64(0), o.StartTime)
+	assert.Equal(t, int64(0), o.EndTime)
 }
 
 func TestOptionsWithFlags(t *testing.T) {
@@ -40,6 +42,8 @@ func TestOptionsWithFlags(t *testing.T) {
 		"--hash-logs",
 		"--hash-process",
 		"--max-spans-count=100",
+		"--start-time=1",
+		"--end-time=2",
 	})
 
 	assert.Equal(t, "192.168.1.10:16686", o.QueryGRPCHostPort)
@@ -50,6 +54,8 @@ func TestOptionsWithFlags(t *testing.T) {
 	assert.True(t, o.HashLogs)
 	assert.True(t, o.HashProcess)
 	assert.Equal(t, 100, o.MaxSpansCount)
+	assert.Equal(t, int64(1), o.StartTime)
+	assert.Equal(t, int64(2), o.EndTime)
 }
 
 func TestMain(m *testing.M) {

cmd/anonymizer/app/query/query.go

@@ -9,6 +9,7 @@ import (
 	"fmt"
 	"io"
 	"strings"
+	"time"
 
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials/insecure"
@@ -49,14 +50,16 @@ func unwrapNotFoundErr(err error) error {
 }
 
 // QueryTrace queries for a trace and returns all spans inside it
-func (q *Query) QueryTrace(traceID string) ([]model.Span, error) {
+func (q *Query) QueryTrace(traceID string, startTime time.Time, endTime time.Time) ([]model.Span, error) {
 	mTraceID, err := model.TraceIDFromString(traceID)
 	if err != nil {
 		return nil, fmt.Errorf("failed to convert the provided trace id: %w", err)
 	}
-	// TODO: add start time & end time
+
 	request := api_v2.GetTraceRequest{
-		TraceID: mTraceID,
+		TraceID:   mTraceID,
+		StartTime: startTime,
+		EndTime:   endTime,
 	}
 
 	stream, err := q.client.GetTrace(context.Background(), &request)

cmd/anonymizer/app/query/query_test.go

@@ -7,6 +7,7 @@ import (
 	"net"
 	"sync"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
@@ -108,24 +109,31 @@ func TestQueryTrace(t *testing.T) {
 	defer q.Close()
 
 	t.Run("No error", func(t *testing.T) {
-		s.spanReader.On("GetTrace", matchContext, matchGetTraceParameters).Return(
+		startTime := time.Date(1970, time.January, 1, 0, 0, 0, 1000, time.UTC)
+		endTime := time.Date(1970, time.January, 1, 0, 0, 0, 2000, time.UTC)
+		expectedGetTraceParameters := spanstore.GetTraceParameters{
+			TraceID:   mockTraceID,
+			StartTime: startTime,
+			EndTime:   endTime,
+		}
+		s.spanReader.On("GetTrace", matchContext, expectedGetTraceParameters).Return(
 			mockTraceGRPC, nil).Once()
 
-		spans, err := q.QueryTrace(mockTraceID.String())
+		spans, err := q.QueryTrace(mockTraceID.String(), startTime, endTime)
 		require.NoError(t, err)
 		assert.Equal(t, len(spans), len(mockTraceGRPC.Spans))
 	})
 
 	t.Run("Invalid TraceID", func(t *testing.T) {
-		_, err := q.QueryTrace(mockInvalidTraceID)
+		_, err := q.QueryTrace(mockInvalidTraceID, time.Time{}, time.Time{})
 		assert.ErrorContains(t, err, "failed to convert the provided trace id")
 	})
 
 	t.Run("Trace not found", func(t *testing.T) {
 		s.spanReader.On("GetTrace", matchContext, matchGetTraceParameters).Return(
 			nil, spanstore.ErrTraceNotFound).Once()
 
-		spans, err := q.QueryTrace(mockTraceID.String())
+		spans, err := q.QueryTrace(mockTraceID.String(), time.Time{}, time.Time{})
 		assert.Nil(t, spans)
 		assert.ErrorIs(t, err, spanstore.ErrTraceNotFound)
 	})

cmd/anonymizer/main.go

@@ -7,6 +7,7 @@ import (
 	"errors"
 	"fmt"
 	"os"
+	"time"
 
 	"github.com/spf13/cobra"
 	"go.uber.org/zap"
@@ -53,7 +54,11 @@ func main() {
 				logger.Fatal("error while creating query object", zap.Error(err))
 			}
 
-			spans, err := query.QueryTrace(options.TraceID)
+			spans, err := query.QueryTrace(
+				options.TraceID,
+				initTime(options.StartTime),
+				initTime(options.EndTime),
+			)
 			if err != nil {
 				logger.Fatal("error while querying for trace", zap.Error(err))
 			}
@@ -93,3 +98,11 @@ func main() {
 		os.Exit(1)
 	}
 }
+
+func initTime(ts int64) time.Time {
+	var t time.Time
+	if ts != 0 {
+		t = time.Unix(0, ts)
+	}
+	return t
+}