From 07f38e2f92c276b16295175d2187773175b15d1c Mon Sep 17 00:00:00 2001
From: Yuri Shkuro <github@ysh.us>
Date: Wed, 20 Nov 2024 21:09:39 -0500
Subject: [PATCH] fix

Signed-off-by: Yuri Shkuro <github@ysh.us>
---
 cmd/query/app/apiv3/grpc_handler.go      | 37 +++++++-----------------
 cmd/query/app/apiv3/grpc_handler_test.go | 23 ++++++---------
 cmd/query/app/internal/api_v3/traces.go  | 10 +++++++
 3 files changed, 29 insertions(+), 41 deletions(-)

diff --git a/cmd/query/app/apiv3/grpc_handler.go b/cmd/query/app/apiv3/grpc_handler.go
index bff197a7ea4..f27f5560c89 100644
--- a/cmd/query/app/apiv3/grpc_handler.go
+++ b/cmd/query/app/apiv3/grpc_handler.go
@@ -8,7 +8,6 @@ import (
 	"errors"
 	"fmt"
 
-	"github.com/gogo/protobuf/types"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/status"
 
@@ -51,8 +50,8 @@ func (h *Handler) FindTraces(request *api_v3.FindTracesRequest, stream api_v3.Qu
 	if query == nil {
 		return status.Error(codes.InvalidArgument, "missing query")
 	}
-	if query.GetStartTimeMin() == nil ||
-		query.GetStartTimeMax() == nil {
+	if query.GetStartTimeMin().IsZero() ||
+		query.GetStartTimeMax().IsZero() {
 		return errors.New("start time min and max are required parameters")
 	}
 
@@ -62,33 +61,17 @@ func (h *Handler) FindTraces(request *api_v3.FindTracesRequest, stream api_v3.Qu
 		Tags:          query.GetAttributes(),
 		NumTraces:     int(query.GetSearchDepth()),
 	}
-	if query.GetStartTimeMin() != nil {
-		startTimeMin, err := types.TimestampFromProto(query.GetStartTimeMin())
-		if err != nil {
-			return err
-		}
-		queryParams.StartTimeMin = startTimeMin
+	if ts := query.GetStartTimeMin(); !ts.IsZero() {
+		queryParams.StartTimeMin = ts
 	}
-	if query.GetStartTimeMax() != nil {
-		startTimeMax, err := types.TimestampFromProto(query.GetStartTimeMax())
-		if err != nil {
-			return err
-		}
-		queryParams.StartTimeMax = startTimeMax
+	if ts := query.GetStartTimeMax(); !ts.IsZero() {
+		queryParams.StartTimeMax = ts
 	}
-	if query.GetDurationMin() != nil {
-		durationMin, err := types.DurationFromProto(query.GetDurationMin())
-		if err != nil {
-			return err
-		}
-		queryParams.DurationMin = durationMin
+	if d := query.GetDurationMin(); d != 0 {
+		queryParams.DurationMin = d
 	}
-	if query.GetDurationMax() != nil {
-		durationMax, err := types.DurationFromProto(query.GetDurationMax())
-		if err != nil {
-			return err
-		}
-		queryParams.DurationMax = durationMax
+	if d := query.GetDurationMax(); d != 0 {
+		queryParams.DurationMax = d
 	}
 
 	traces, err := h.QueryService.FindTraces(stream.Context(), queryParams)
diff --git a/cmd/query/app/apiv3/grpc_handler_test.go b/cmd/query/app/apiv3/grpc_handler_test.go
index bb3c08d9b59..7aaccb7c5d8 100644
--- a/cmd/query/app/apiv3/grpc_handler_test.go
+++ b/cmd/query/app/apiv3/grpc_handler_test.go
@@ -8,8 +8,8 @@ import (
 	"errors"
 	"net"
 	"testing"
+	"time"
 
-	"github.com/gogo/protobuf/types"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
@@ -124,7 +124,9 @@ func TestGetTraceTraceIDError(t *testing.T) {
 		}, nil).Once()
 
 	getTraceStream, err := tsc.client.GetTrace(context.Background(), &api_v3.GetTraceRequest{
-		TraceId: "Z",
+		TraceId:   "Z",
+		StartTime: time.Now().Add(-2 * time.Hour),
+		EndTime:   time.Now(),
 	})
 	require.NoError(t, err)
 	recv, err := getTraceStream.Recv()
@@ -150,10 +152,8 @@ func TestFindTraces(t *testing.T) {
 			ServiceName:   "myservice",
 			OperationName: "opname",
 			Attributes:    map[string]string{"foo": "bar"},
-			StartTimeMin:  &types.Timestamp{},
-			StartTimeMax:  &types.Timestamp{},
-			DurationMin:   &types.Duration{},
-			DurationMax:   &types.Duration{},
+			StartTimeMin:  time.Now().Add(-2 * time.Hour),
+			StartTimeMax:  time.Now(),
 		},
 	})
 	require.NoError(t, err)
@@ -172,10 +172,7 @@ func TestFindTracesQueryNil(t *testing.T) {
 	assert.Nil(t, recv)
 
 	responseStream, err = tsc.client.FindTraces(context.Background(), &api_v3.FindTracesRequest{
-		Query: &api_v3.TraceQueryParameters{
-			StartTimeMin: nil,
-			StartTimeMax: nil,
-		},
+		Query: &api_v3.TraceQueryParameters{},
 	})
 	require.NoError(t, err)
 	recv, err = responseStream.Recv()
@@ -190,10 +187,8 @@ func TestFindTracesStorageError(t *testing.T) {
 
 	responseStream, err := tsc.client.FindTraces(context.Background(), &api_v3.FindTracesRequest{
 		Query: &api_v3.TraceQueryParameters{
-			StartTimeMin: &types.Timestamp{},
-			StartTimeMax: &types.Timestamp{},
-			DurationMin:  &types.Duration{},
-			DurationMax:  &types.Duration{},
+			StartTimeMin: time.Now().Add(-2 * time.Hour),
+			StartTimeMax: time.Now(),
 		},
 	})
 	require.NoError(t, err)
diff --git a/cmd/query/app/internal/api_v3/traces.go b/cmd/query/app/internal/api_v3/traces.go
index 04b9c2a850b..1208c437c03 100644
--- a/cmd/query/app/internal/api_v3/traces.go
+++ b/cmd/query/app/internal/api_v3/traces.go
@@ -38,6 +38,16 @@ func (*TracesData) MarshalTo([]byte /* data */) (n int, err error) {
 	panic("unimplemented")
 }
 
+// MarshalToSizedBuffer is used by Gogo.
+func (td *TracesData) MarshalToSizedBuffer(buf []byte) (int, error) {
+	data, err := td.Marshal()
+	if err != nil {
+		return 0, err
+	}
+	copy(buf, data)
+	return len(data), nil
+}
+
 // MarshalJSONPB implements gogocodec.CustomType.
 func (td *TracesData) MarshalJSONPB(*jsonpb.Marshaler) ([]byte, error) {
 	return new(ptrace.JSONMarshaler).MarshalTraces(td.ToTraces())