tempo-cli: add support for /api/v2/traces endpoint (#4127)

* tempo-cli: add support for traces API V2

* update CHANGELOG.md

* enable help to be more compact but still fully-specified form

* request protobuf for v2 endpoint

* default to using v2 endpoint and mark it as breaking change

* use generics in printAsJSON

* fix: support for +Inf, -Inf and NaN values in trace by id endpoints
'encoding/json' package can't handle +Inf, -Inf, NaN values and will
fail to Marshal the response from tracebyid endpoints if response has keys with
values of +Inf, -Inf or NaN.

gogo/protobuf/jsonpb package correctly handles these values so use that to
Marshal and print the response to stdout

CHANGELOG.md

@@ -1,5 +1,8 @@
 ## main / unreleased
 
+* [CHANGE] tempo-cli: add support for /api/v2/traces endpoint [#4127](https://github.com/grafana/tempo/pull/4127) (@electron0zero)
+  **BREAKING CHANGE** The `tempo-cli` now uses the `/api/v2/traces` endpoint by default, 
+  please use `--v1` flag to use `/api/traces` endpoint, which was the default in previous versions.
 * [ENHANCEMENT] Speedup collection of results from ingesters in the querier [#4100](https://github.com/grafana/tempo/pull/4100) (@electron0zero)
 * [ENHANCEMENT] Speedup DistinctValue collector and exit early for ingesters [#4104](https://github.com/grafana/tempo/pull/4104) (@electron0zero)
 * [ENHANCEMENT] Add disk caching in ingester SearchTagValuesV2 for completed blocks [#4069](https://github.com/grafana/tempo/pull/4069) (@electron0zero)

cmd/tempo-cli/cmd-query-trace-id.go

@@ -1,24 +1,54 @@
 package main
 
 import (
+	"fmt"
+	"os"
+
+	"github.com/gogo/protobuf/jsonpb"
 	"github.com/grafana/tempo/pkg/httpclient"
+	"github.com/grafana/tempo/pkg/tempopb"
 )
 
 type queryTraceIDCmd struct {
 	APIEndpoint string `arg:"" help:"tempo api endpoint"`
 	TraceID     string `arg:"" help:"trace ID to retrieve"`
 
+	V1    bool   `name:"v1" help:"Use v1 API /api/traces endpoint"`
 	OrgID string `help:"optional orgID"`
 }
 
 func (cmd *queryTraceIDCmd) Run(_ *globalOptions) error {
 	client := httpclient.New(cmd.APIEndpoint, cmd.OrgID)
-
 	// util.QueryTrace will only add orgID header if len(orgID) > 0
-	trace, err := client.QueryTrace(cmd.TraceID)
+
+	// use v1 API if specified, we default to v2
+	if cmd.V1 {
+		trace, err := client.QueryTrace(cmd.TraceID)
+		if err != nil {
+			return err
+		}
+		return printTrace(trace)
+	}
+
+	traceResp, err := client.QueryTraceV2(cmd.TraceID)
 	if err != nil {
 		return err
 	}
+	if traceResp.Message != "" {
+		// print message and status to stderr if there is one.
+		// allows users to get a clean trace on the stdout, and pipe it to a file or another commands.
+		_, _ = fmt.Fprintf(os.Stderr, "status: %s , message: %s\n", traceResp.Status, traceResp.Message)
+	}
+	return printTrace(traceResp.Trace)
+}
 
-	return printAsJSON(trace)
+func printTrace(trace *tempopb.Trace) error {
+	// tracebyid endpoints are protobuf, we are using 'gogo/protobuf/jsonpb' to marshal the
+	// trace to json because 'encoding/json' package can't handle +Inf, -Inf, NaN
+	marshaller := &jsonpb.Marshaler{}
+	err := marshaller.Marshal(os.Stdout, trace)
+	if err != nil {
+		_, _ = fmt.Fprintf(os.Stderr, "Failed to marshal trace: %v\n", err)
+	}
+	return nil
 }

cmd/tempo-cli/main.go

@@ -94,7 +94,7 @@ func main() {
 	ctx := kong.Parse(&cli,
 		kong.UsageOnError(),
 		kong.ConfigureHelp(kong.HelpOptions{
-			// Compact: true,
+			Compact: true,
 		}),
 	)
 	err := ctx.Run(&cli.globalOptions)

cmd/tempo-cli/shared.go

@@ -125,7 +125,7 @@ func loadBlock(r backend.Reader, c backend.Compactor, tenantID string, id uuid.U
 	}, nil
 }
 
-func printAsJSON(value interface{}) error {
+func printAsJSON[T any](value T) error {
 	traceJSON, err := json.Marshal(value)
 	if err != nil {
 		return err

docs/sources/tempo/operations/tempo_cli.md

@@ -70,6 +70,7 @@ Arguments:
 
 Options:
 - `--org-id <value>` Organization ID (for use in multi-tenant setup).
+- `--v1` use v1 API (use /api/traces endpoint to fetch traces, default: /api/v2/traces).
 
 **Example:**
 ```bash

pkg/httpclient/client.go

@@ -26,7 +26,8 @@ import (
 const (
 	orgIDHeader = "X-Scope-OrgID"
 
-	QueryTraceEndpoint = "/api/traces"
+	QueryTraceEndpoint   = "/api/traces"
+	QueryTraceV2Endpoint = "/api/v2/traces"
 
 	acceptHeader        = "Accept"
 	applicationProtobuf = "application/protobuf"
@@ -95,11 +96,11 @@ func (c *Client) getFor(url string, m proto.Message) (*http.Response, error) {
 	}
 
 	marshallingFormat := applicationJSON
-	if strings.Contains(url, QueryTraceEndpoint) {
+	if strings.Contains(url, QueryTraceEndpoint) || strings.Contains(url, QueryTraceV2Endpoint) {
 		marshallingFormat = applicationProtobuf
 	}
 	// Set 'Accept' header to 'application/protobuf'.
-	// This is required for the /api/traces endpoint to return a protobuf response.
+	// This is required for the /api/traces and /api/v2/traces endpoint to return a protobuf response.
 	// JSON lost backwards compatibility with the upgrade to `opentelemetry-proto` v0.18.0.
 	req.Header.Set(acceptHeader, marshallingFormat)
 
@@ -253,7 +254,18 @@ func (c *Client) QueryTrace(id string) (*tempopb.Trace, error) {
 		}
 		return nil, err
 	}
+	return m, nil
+}
 
+func (c *Client) QueryTraceV2(id string) (*tempopb.TraceByIDResponse, error) {
+	m := &tempopb.TraceByIDResponse{}
+	resp, err := c.getFor(c.BaseURL+QueryTraceV2Endpoint+"/"+id, m)
+	if err != nil {
+		if resp != nil && resp.StatusCode == http.StatusNotFound {
+			return nil, util.ErrTraceNotFound
+		}
+		return nil, err
+	}
 	return m, nil
 }