query-frontend: also extract Trace from TraceByIDResponse when return…

…ing protobuf (#1025)

* query-frontend: also extract Trace from TraceByIDResponse when returning protobuf

* Fix CHANGELOG.md

CHANGELOG.md

@@ -58,9 +58,9 @@
 * [ENHANCEMENT] Shard tenant index creation by tenant and add functionality to handle stale indexes. [#1005](https://github.com/grafana/tempo/pull/1005) (@joe-elliott)
 * [ENHANCEMENT] Added a configurable buffer for WAL writes. [#1018](https://github.com/grafana/tempo/pull/1018) (@joe-elliott)
 * [ENHANCEMENT] **BREAKING CHANGE** Support partial results from failed block queries [#1007](https://github.com/grafana/tempo/pull/1007) (@mapno)
-* [ENHANCEMENT] Add `search_default_limit` to querier config. [#1022](https://github.com/grafana/tempo/pull/1022) (@kvrhdn)
   Querier [`GET /querier/api/traces/<traceid>`](https://grafana.com/docs/tempo/latest/api_docs/#query) response's body has been modified 
   to return `tempopb.TraceByIDResponse` instead of simply `tempopb.Trace`. This will cause a disruption of the read path during rollout of the change.
+* [ENHANCEMENT] Add `search_default_limit` to querier config. [#1022](https://github.com/grafana/tempo/pull/1022) (@kvrhdn)
 * [BUGFIX] Update port spec for GCS docker-compose example [#869](https://github.com/grafana/tempo/pull/869) (@zalegrala)
 * [BUGFIX] Fix "magic number" errors and other block mishandling when an ingester forcefully shuts down [#937](https://github.com/grafana/tempo/issues/937) (@mdisibio)
 * [BUGFIX] Fix compactor memory leak [#806](https://github.com/grafana/tempo/pull/806) (@mdisibio)

modules/frontend/frontend.go

@@ -143,14 +143,13 @@ func NewTracesMiddleware(cfg Config, logger log.Logger, registerer prometheus.Re
 				marshallingFormat = util.ProtobufTypeHeaderValue
 			}
 
-			// Enforce all communication internal to Tempo to be in protobuf bytes
+			// enforce all communication internal to Tempo to be in protobuf bytes
 			r.Header.Set(util.AcceptHeaderKey, util.ProtobufTypeHeaderValue)
 
 			resp, err := rt.RoundTrip(r)
 
 			// todo : should all of this request/response content type be up a level and be used for all query types?
-			if resp != nil && resp.StatusCode == http.StatusOK && marshallingFormat == util.JSONTypeHeaderValue {
-				// if request is for application/json, unmarshal into proto object and re-marshal into json bytes
+			if resp != nil && resp.StatusCode == http.StatusOK {
 				body, err := io.ReadAll(resp.Body)
 				resp.Body.Close()
 				if err != nil {
@@ -166,13 +165,21 @@ func NewTracesMiddleware(cfg Config, logger log.Logger, registerer prometheus.Re
 					resp.StatusCode = http.StatusPartialContent
 				}
 
-				var jsonTrace bytes.Buffer
-				marshaller := &jsonpb.Marshaler{}
-				err = marshaller.Marshal(&jsonTrace, responseObject.Trace)
-				if err != nil {
-					return nil, err
+				if marshallingFormat == util.JSONTypeHeaderValue {
+					var jsonTrace bytes.Buffer
+					marshaller := &jsonpb.Marshaler{}
+					err = marshaller.Marshal(&jsonTrace, responseObject.Trace)
+					if err != nil {
+						return nil, err
+					}
+					resp.Body = io.NopCloser(bytes.NewReader(jsonTrace.Bytes()))
+				} else {
+					traceBuffer, err := proto.Marshal(responseObject.Trace)
+					if err != nil {
+						return nil, err
+					}
+					resp.Body = io.NopCloser(bytes.NewReader(traceBuffer))
 				}
-				resp.Body = io.NopCloser(bytes.NewReader(jsonTrace.Bytes()))
 			}
 			span := opentracing.SpanFromContext(r.Context())
 			if span != nil {