Add Content-Type header on query-frontend paths

Fixes grafana#1185
wperron · Feb 23, 2022 · fb474be · fb474be
1 parent 12fff13
commit fb474be
Show file tree

Hide file tree

Showing 8 changed files with 31 additions and 5 deletions.
diff --git a/modules/frontend/frontend.go b/modules/frontend/frontend.go
@@ -155,6 +155,8 @@ func newTraceByIDMiddleware(cfg Config, logger log.Logger) Middleware {
 				span.SetTag("contentType", marshallingFormat)
 			}
 
+			resp.Header.Set(api.HeaderContentType, marshallingFormat)
+
 			return resp, err
 		})
 	})

diff --git a/modules/frontend/handler.go b/modules/frontend/handler.go
@@ -98,7 +98,8 @@ func (f *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-	// write header and body
+	// write headers, status code and body
+	copyHeader(w.Header(), resp.Header)
 	w.WriteHeader(resp.StatusCode)
 	if resp.Body != nil {
 		_, _ = io.Copy(w, resp.Body)
@@ -126,6 +127,14 @@ func (f *handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	)
 }
 
+func copyHeader(dst, src http.Header) {
+	for k, vv := range src {
+		for _, v := range vv {
+			dst.Add(k, v)
+		}
+	}
+}
+
 // writeError handles writing errors to the http.ResponseWriter. It uses weavework common
 // server.WriteError() to handle httpgrc errors. The handler handles all incoming HTTP requests
 // to the query frontend which then distributes them via httpgrpc to the queriers. As a result

diff --git a/modules/frontend/searchsharding.go b/modules/frontend/searchsharding.go
@@ -295,8 +295,10 @@ func (s searchSharder) RoundTrip(r *http.Request) (*http.Response, error) {
 	}
 
 	return &http.Response{
-		StatusCode:    http.StatusOK,
-		Header:        http.Header{},
+		StatusCode: http.StatusOK,
+		Header: http.Header{
+			api.HeaderContentType: {api.HeaderAcceptJSON},
+		},
 		Body:          io.NopCloser(strings.NewReader(bodyString)),
 		ContentLength: int64(len([]byte(bodyString))),
 	}, nil

diff --git a/modules/frontend/searchsharding_test.go b/modules/frontend/searchsharding_test.go
@@ -618,6 +618,9 @@ func TestSearchSharderRoundTrip(t *testing.T) {
 			}
 			require.NoError(t, err)
 			assert.Equal(t, tc.expectedStatus, resp.StatusCode)
+			if tc.expectedStatus == http.StatusOK {
+				assert.Equal(t, "application/json", resp.Header.Get("Content-Type"))
+			}
 			if tc.expectedResponse != nil {
 				actualResp := &tempopb.SearchResponse{}
 				bytesResp, err := io.ReadAll(resp.Body)

diff --git a/modules/frontend/tracebyidsharding.go b/modules/frontend/tracebyidsharding.go
@@ -15,6 +15,7 @@ import (
 	"github.com/go-kit/log/level"
 	"github.com/golang/protobuf/proto"
 	"github.com/grafana/tempo/modules/querier"
+	"github.com/grafana/tempo/pkg/api"
 	"github.com/grafana/tempo/pkg/model/trace"
 	"github.com/grafana/tempo/pkg/tempopb"
 	"github.com/opentracing/opentracing-go"
@@ -175,8 +176,10 @@ func (s shardQuery) RoundTrip(r *http.Request) (*http.Response, error) {
 	}
 
 	return &http.Response{
-		StatusCode:    http.StatusOK,
-		Header:        http.Header{},
+		StatusCode: http.StatusOK,
+		Header: http.Header{
+			api.HeaderContentType: {api.HeaderAcceptProtobuf},
+		},
 		Body:          io.NopCloser(bytes.NewReader(buff)),
 		ContentLength: int64(len(buff)),
 	}, nil

diff --git a/modules/frontend/tracebyidsharding_test.go b/modules/frontend/tracebyidsharding_test.go
@@ -327,6 +327,9 @@ func TestShardingWareDoRequest(t *testing.T) {
 			}
 			require.NoError(t, err)
 			assert.Equal(t, tc.expectedStatus, resp.StatusCode)
+			if tc.expectedStatus == http.StatusOK {
+				assert.Equal(t, "application/protobuf", resp.Header.Get("Content-Type"))
+			}
 			if tc.expectedTrace != nil {
 				actualResp := &tempopb.TraceByIDResponse{}
 				bytesTrace, err := io.ReadAll(resp.Body)

diff --git a/modules/querier/http.go b/modules/querier/http.go
@@ -193,6 +193,7 @@ func (q *Querier) SearchHandler(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		return
 	}
+	w.Header().Set(api.HeaderContentType, api.HeaderAcceptJSON)
 }
 
 func (q *Querier) SearchTagsHandler(w http.ResponseWriter, r *http.Request) {
@@ -217,6 +218,7 @@ func (q *Querier) SearchTagsHandler(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		return
 	}
+	w.Header().Set(api.HeaderContentType, api.HeaderAcceptJSON)
 }
 
 func (q *Querier) SearchTagValuesHandler(w http.ResponseWriter, r *http.Request) {
@@ -249,4 +251,5 @@ func (q *Querier) SearchTagValuesHandler(w http.ResponseWriter, r *http.Request)
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 		return
 	}
+	w.Header().Set(api.HeaderContentType, api.HeaderAcceptJSON)
 }
diff --git a/pkg/api/http.go b/pkg/api/http.go
@@ -38,6 +38,7 @@ const (
 	urlParamVersion       = "version"
 
 	HeaderAccept         = "Accept"
+	HeaderContentType    = "Content-Type"
 	HeaderAcceptProtobuf = "application/protobuf"
 	HeaderAcceptJSON     = "application/json"