[grpc][v2] Implement FindTraceIDs in gRPC v2 handler

internal/storage/v2/grpc/handler.go

@@ -103,17 +103,7 @@ func (h *Handler) FindTraces(
 	req *storage.FindTracesRequest,
 	srv storage.TraceReader_FindTracesServer,
 ) error {
-	query := tracestore.TraceQueryParams{
-		ServiceName:   req.Query.ServiceName,
-		OperationName: req.Query.OperationName,
-		Attributes:    convertKeyValueListToMap(req.Query.Attributes),
-		StartTimeMin:  req.Query.StartTimeMin,
-		StartTimeMax:  req.Query.StartTimeMax,
-		DurationMin:   req.Query.DurationMin,
-		DurationMax:   req.Query.DurationMax,
-		SearchDepth:   int(req.Query.SearchDepth),
-	}
-	for traces, err := range h.traceReader.FindTraces(srv.Context(), query) {
+	for traces, err := range h.traceReader.FindTraces(srv.Context(), toTraceQueryParams(req.Query)) {
 		if err != nil {
 			return err
 		}
@@ -127,6 +117,41 @@ func (h *Handler) FindTraces(
 	return nil
 }
 
+func (h *Handler) FindTraceIDs(
+	ctx context.Context,
+	req *storage.FindTracesRequest,
+) (*storage.FindTraceIDsResponse, error) {
+	foundTraceIDs := []*storage.FoundTraceID{}
+	for traceIDs, err := range h.traceReader.FindTraceIDs(ctx, toTraceQueryParams(req.Query)) {
+		if err != nil {
+			return nil, err
+		}
+		for _, traceID := range traceIDs {
+			foundTraceIDs = append(foundTraceIDs, &storage.FoundTraceID{
+				TraceId: traceID.TraceID[:],
+				Start:   traceID.Start,
+				End:     traceID.End,
+			})
+		}
+	}
+	return &storage.FindTraceIDsResponse{
+		TraceIds: foundTraceIDs,
+	}, nil
+}
+
+func toTraceQueryParams(t *storage.TraceQueryParameters) tracestore.TraceQueryParams {
+	return tracestore.TraceQueryParams{
+		ServiceName:   t.ServiceName,
+		OperationName: t.OperationName,
+		Attributes:    convertKeyValueListToMap(t.Attributes),
+		StartTimeMin:  t.StartTimeMin,
+		StartTimeMax:  t.StartTimeMax,
+		DurationMin:   t.DurationMin,
+		DurationMax:   t.DurationMax,
+		SearchDepth:   int(t.SearchDepth),
+	}
+}
+
 func convertKeyValueListToMap(kvList []*storage.KeyValue) pcommon.Map {
 	m := pcommon.NewMap()
 	for _, kv := range kvList {

internal/storage/v2/grpc/handler_test.go

@@ -311,6 +311,84 @@ func TestHandler_FindTraces(t *testing.T) {
 	}
 }
 
+func TestHandler_FindTraceIDs(t *testing.T) {
+	reader := new(tracestoremocks.Reader)
+	query := tracestore.TraceQueryParams{
+		ServiceName:   "service",
+		OperationName: "operation",
+		Attributes:    pcommon.NewMap(),
+	}
+	now := time.Now()
+	traceIDA := [16]byte{1}
+	traceIDB := [16]byte{2}
+	tests := []struct {
+		name             string
+		traceIDs         []tracestore.FoundTraceID
+		expectedTraceIDs []*storage.FoundTraceID
+		findTraceIDsErr  error
+		expectedErr      error
+	}{
+		{
+			name: "success",
+			traceIDs: []tracestore.FoundTraceID{
+				{
+					TraceID: traceIDA,
+					Start:   now,
+					End:     now.Add(time.Minute),
+				},
+				{
+					TraceID: traceIDB,
+					Start:   now,
+					End:     now.Add(time.Hour),
+				},
+			},
+			expectedTraceIDs: []*storage.FoundTraceID{
+				{
+					TraceId: traceIDA[:],
+					Start:   now,
+					End:     now.Add(time.Minute),
+				},
+				{
+					TraceId: traceIDB[:],
+					Start:   now,
+					End:     now.Add(time.Hour),
+				},
+			},
+		},
+		{
+			name:             "empty",
+			traceIDs:         []tracestore.FoundTraceID{},
+			expectedTraceIDs: []*storage.FoundTraceID{},
+		},
+		{
+			name:            "error",
+			findTraceIDsErr: assert.AnError,
+			expectedErr:     assert.AnError,
+		},
+	}
+
+	for _, test := range tests {
+		reader.On("FindTraceIDs", mock.Anything, query).
+			Return(iter.Seq2[[]tracestore.FoundTraceID, error](func(yield func([]tracestore.FoundTraceID, error) bool) {
+				yield(test.traceIDs, test.findTraceIDsErr)
+			})).Once()
+		server := NewHandler(reader)
+
+		response, err := server.FindTraceIDs(context.Background(), &storage.FindTracesRequest{
+			Query: &storage.TraceQueryParameters{
+				ServiceName:   "service",
+				OperationName: "operation",
+			},
+		})
+		if test.expectedErr != nil {
+			require.ErrorIs(t, err, test.expectedErr)
+		} else {
+			require.NoError(t, err)
+			require.Equal(t, test.expectedTraceIDs, response.TraceIds)
+		}
+	}
+}
+
 func TestConvertKeyValueListToMap(t *testing.T) {
 	tests := []struct {
 		name     string