From f29339defdf0b6cc2027487c961e349466624394 Mon Sep 17 00:00:00 2001 From: Sven Grossmann Date: Mon, 21 Nov 2022 16:44:35 +0100 Subject: [PATCH 1/3] Loki: Add compression to `callResource` --- pkg/tsdb/loki/api.go | 25 ++++++++++++++++++++----- pkg/tsdb/loki/api_mock.go | 31 +++++++++++++++++++++++++++++++ pkg/tsdb/loki/api_test.go | 25 +++++++++++++------------ pkg/tsdb/loki/loki.go | 24 ++++++++++++++++-------- 4 files changed, 80 insertions(+), 25 deletions(-) diff --git a/pkg/tsdb/loki/api.go b/pkg/tsdb/loki/api.go index 553cc9bdf4dd..15743e5efebe 100644 --- a/pkg/tsdb/loki/api.go +++ b/pkg/tsdb/loki/api.go @@ -24,6 +24,11 @@ type LokiAPI struct { headers map[string]string } +type RawLokiResponse struct { + Body []byte + Encoding string +} + func newLokiAPI(client *http.Client, url string, log log.Logger, headers map[string]string) *LokiAPI { return &LokiAPI{client: client, url: url, log: log, headers: headers} } @@ -204,15 +209,15 @@ func makeRawRequest(ctx context.Context, lokiDsUrl string, resourcePath string, return req, nil } -func (api *LokiAPI) RawQuery(ctx context.Context, resourcePath string) ([]byte, error) { +func (api *LokiAPI) RawQuery(ctx context.Context, resourcePath string) (RawLokiResponse, error) { req, err := makeRawRequest(ctx, api.url, resourcePath, api.headers) if err != nil { - return nil, err + return RawLokiResponse{}, err } resp, err := api.client.Do(req) if err != nil { - return nil, err + return RawLokiResponse{}, err } defer func() { @@ -222,8 +227,18 @@ func (api *LokiAPI) RawQuery(ctx context.Context, resourcePath string) ([]byte, }() if resp.StatusCode/100 != 2 { - return nil, makeLokiError(resp.Body) + return RawLokiResponse{}, makeLokiError(resp.Body) + } + + body, err := io.ReadAll(resp.Body) + if err != nil { + return RawLokiResponse{}, err + } + + encodedBytes := RawLokiResponse{ + Body: body, + Encoding: resp.Header.Get("Content-Encoding"), } - return io.ReadAll(resp.Body) + return encodedBytes, nil } diff --git a/pkg/tsdb/loki/api_mock.go b/pkg/tsdb/loki/api_mock.go index ae54fb852932..425c0a885a03 100644 --- a/pkg/tsdb/loki/api_mock.go +++ b/pkg/tsdb/loki/api_mock.go @@ -32,6 +32,29 @@ func (mockedRT *mockedRoundTripper) RoundTrip(req *http.Request) (*http.Response }, nil } +type mockedCompressedRoundTripper struct { + statusCode int + responseBytes []byte + contentType string + requestCallback mockRequestCallback +} + +func (mockedRT *mockedCompressedRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + requestCallback := mockedRT.requestCallback + if requestCallback != nil { + requestCallback(req) + } + + header := http.Header{} + header.Add("Content-Type", mockedRT.contentType) + header.Add("Content-Encoding", "gzip") + return &http.Response{ + StatusCode: mockedRT.statusCode, + Header: header, + Body: io.NopCloser(bytes.NewReader(mockedRT.responseBytes)), + }, nil +} + func makeMockedAPI(statusCode int, contentType string, responseBytes []byte, requestCallback mockRequestCallback) *LokiAPI { return makeMockedAPIWithUrl("http://localhost:9999", statusCode, contentType, responseBytes, requestCallback) } @@ -43,3 +66,11 @@ func makeMockedAPIWithUrl(url string, statusCode int, contentType string, respon return newLokiAPI(&client, url, log.New("test"), nil) } + +func makeCompressedMockedAPIWithUrl(url string, statusCode int, contentType string, responseBytes []byte, requestCallback mockRequestCallback) *LokiAPI { + client := http.Client{ + Transport: &mockedCompressedRoundTripper{statusCode: statusCode, contentType: contentType, responseBytes: responseBytes, requestCallback: requestCallback}, + } + + return newLokiAPI(&client, url, log.New("test"), nil) +} diff --git a/pkg/tsdb/loki/api_test.go b/pkg/tsdb/loki/api_test.go index 4db38dd14f38..017d4a395f42 100644 --- a/pkg/tsdb/loki/api_test.go +++ b/pkg/tsdb/loki/api_test.go @@ -135,18 +135,19 @@ func TestApiUrlHandling(t *testing.T) { require.True(t, called) }) } +} - for _, test := range queryTestData { - t.Run("Loki should build the metadata query URL correctly when "+test.name, func(t *testing.T) { - called := false - api := makeMockedAPIWithUrl(test.dsUrl, 200, "application/json", response, func(req *http.Request) { - called = true - require.Equal(t, test.metaUrl, req.URL.String()) - }) - - _, err := api.RawQuery(context.Background(), "/loki/api/v1/labels?start=1&end=2") - require.NoError(t, err) - require.True(t, called) +func TestApiReturnValues(t *testing.T) { + t.Run("Loki should return the right encoding", func(t *testing.T) { + called := false + api := makeCompressedMockedAPIWithUrl("http://localhost:3100", 200, "application/json", []byte("foo"), func(req *http.Request) { + called = true }) - } + + encodedBytes, err := api.RawQuery(context.Background(), "/loki/api/v1/labels?start=1&end=2") + require.NoError(t, err) + require.True(t, called) + require.Equal(t, "gzip", encodedBytes.Encoding) + require.Equal(t, []byte("foo"), encodedBytes.Body) + }) } diff --git a/pkg/tsdb/loki/loki.go b/pkg/tsdb/loki/loki.go index 49c1f46b9094..12a2d8561a7a 100644 --- a/pkg/tsdb/loki/loki.go +++ b/pkg/tsdb/loki/loki.go @@ -119,7 +119,7 @@ func (s *Service) CallResource(ctx context.Context, req *backend.CallResourceReq return callResource(ctx, req, sender, dsInfo, logger.FromContext(ctx)) } -func getAuthHeadersForCallResource(headers map[string][]string) map[string]string { +func getHeadersForCallResource(headers map[string][]string) map[string]string { data := make(map[string]string) if auth := arrayHeaderFirstValue(headers["Authorization"]); auth != "" { @@ -134,6 +134,10 @@ func getAuthHeadersForCallResource(headers map[string][]string) map[string]strin data["X-ID-Token"] = idToken } + if encType := arrayHeaderFirstValue(headers["Accept-Encoding"]); encType != "" { + data["Accept-Encoding"] = encType + } + return data } @@ -151,19 +155,23 @@ func callResource(ctx context.Context, req *backend.CallResourceRequest, sender } lokiURL := fmt.Sprintf("/loki/api/v1/%s", url) - api := newLokiAPI(dsInfo.HTTPClient, dsInfo.URL, plog, getAuthHeadersForCallResource(req.Headers)) - bytes, err := api.RawQuery(ctx, lokiURL) + api := newLokiAPI(dsInfo.HTTPClient, dsInfo.URL, plog, getHeadersForCallResource(req.Headers)) + encodedBytes, err := api.RawQuery(ctx, lokiURL) if err != nil { return err } + respHeaders := map[string][]string{ + "content-type": {"application/json"}, + } + if encodedBytes.Encoding != "" { + respHeaders["content-encoding"] = []string{encodedBytes.Encoding} + } return sender.Send(&backend.CallResourceResponse{ - Status: http.StatusOK, - Headers: map[string][]string{ - "content-type": {"application/json"}, - }, - Body: bytes, + Status: http.StatusOK, + Headers: respHeaders, + Body: encodedBytes.Body, }) } From 00ef35246784c914ffb72c0e6645b87b6f7c369a Mon Sep 17 00:00:00 2001 From: Sven Grossmann Date: Mon, 21 Nov 2022 17:09:08 +0100 Subject: [PATCH 2/3] add missing tests --- pkg/tsdb/loki/api_test.go | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/pkg/tsdb/loki/api_test.go b/pkg/tsdb/loki/api_test.go index 017d4a395f42..7ac6015eb808 100644 --- a/pkg/tsdb/loki/api_test.go +++ b/pkg/tsdb/loki/api_test.go @@ -135,6 +135,20 @@ func TestApiUrlHandling(t *testing.T) { require.True(t, called) }) } + + for _, test := range queryTestData { + t.Run("Loki should build the metadata query URL correctly when "+test.name, func(t *testing.T) { + called := false + api := makeMockedAPIWithUrl(test.dsUrl, 200, "application/json", response, func(req *http.Request) { + called = true + require.Equal(t, test.metaUrl, req.URL.String()) + }) + + _, err := api.RawQuery(context.Background(), "/loki/api/v1/labels?start=1&end=2") + require.NoError(t, err) + require.True(t, called) + }) + } } func TestApiReturnValues(t *testing.T) { From 6cfc4e80958f0bfd296bd6d27ed2e73ffb93630a Mon Sep 17 00:00:00 2001 From: Sven Grossmann Date: Tue, 22 Nov 2022 13:16:08 +0100 Subject: [PATCH 3/3] fix formatting --- pkg/tsdb/loki/api_test.go | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pkg/tsdb/loki/api_test.go b/pkg/tsdb/loki/api_test.go index 7ac6015eb808..2b64e383bf4c 100644 --- a/pkg/tsdb/loki/api_test.go +++ b/pkg/tsdb/loki/api_test.go @@ -135,7 +135,7 @@ func TestApiUrlHandling(t *testing.T) { require.True(t, called) }) } - + for _, test := range queryTestData { t.Run("Loki should build the metadata query URL correctly when "+test.name, func(t *testing.T) { called := false