From 988ce8a5e817bdc7011a17135b8827706191059b Mon Sep 17 00:00:00 2001 From: Zhang Xin Date: Tue, 17 Dec 2024 22:39:22 +0800 Subject: [PATCH] Feat: Add time window for GetTrace of anonymizer (#6368) --- cmd/anonymizer/app/flags.go | 14 ++++++++++++++ cmd/anonymizer/app/flags_test.go | 6 ++++++ cmd/anonymizer/app/query/query.go | 9 ++++++--- cmd/anonymizer/app/query/query_test.go | 16 ++++++++++++---- cmd/anonymizer/main.go | 15 ++++++++++++++- 5 files changed, 52 insertions(+), 8 deletions(-) diff --git a/cmd/anonymizer/app/flags.go b/cmd/anonymizer/app/flags.go index 51b14d88a07..f376f4af151 100644 --- a/cmd/anonymizer/app/flags.go +++ b/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) diff --git a/cmd/anonymizer/app/flags_test.go b/cmd/anonymizer/app/flags_test.go index a7f6ff28ef4..08b6bbb8538 100644 --- a/cmd/anonymizer/app/flags_test.go +++ b/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) { diff --git a/cmd/anonymizer/app/query/query.go b/cmd/anonymizer/app/query/query.go index 6cca46408ba..77e82fcd103 100644 --- a/cmd/anonymizer/app/query/query.go +++ b/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) diff --git a/cmd/anonymizer/app/query/query_test.go b/cmd/anonymizer/app/query/query_test.go index 8f350a191a0..2aad2f235c7 100644 --- a/cmd/anonymizer/app/query/query_test.go +++ b/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,16 +109,23 @@ 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") }) @@ -125,7 +133,7 @@ func TestQueryTrace(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) }) diff --git a/cmd/anonymizer/main.go b/cmd/anonymizer/main.go index 6f5e27c620b..ad57673cea2 100644 --- a/cmd/anonymizer/main.go +++ b/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 +}