From e569be9fe50e5734cd1f5c1c88b89da9d997a18f Mon Sep 17 00:00:00 2001
From: Karsten Jeschkies <karsten.jeschkies@grafana.com>
Date: Thu, 14 Dec 2023 13:56:47 +0100
Subject: [PATCH] Use correct error write function (#11487)

**What this PR does / why we need it**:
This fixes the last known regression for `protobuf` encoding.

**Checklist**
- [ ] Reviewed the
[`CONTRIBUTING.md`](https://github.com/grafana/loki/blob/main/CONTRIBUTING.md)
guide (**required**)
- [ ] Documentation added
- [x] Tests updated
- [ ] `CHANGELOG.md` updated
- [ ] If the change is worth mentioning in the release notes, add
`add-to-release-notes` label
- [ ] Changes that require user attention or interaction to upgrade are
documented in `docs/sources/setup/upgrade/_index.md`
- [ ] For Helm chart changes bump the Helm chart version in
`production/helm/loki/Chart.yaml` and update
`production/helm/loki/CHANGELOG.md` and
`production/helm/loki/README.md`. [Example
PR](https://github.com/grafana/loki/commit/d10549e3ece02120974929894ee333d07755d213)
- [ ] If the change is deprecating or removing a configuration option,
update the `deprecated-config.yaml` and `deleted-config.yaml` files
respectively in the `tools/deprecated-config-checker` directory.
[Example
PR](https://github.com/grafana/loki/pull/10840/commits/0d4416a4b03739583349934b96f272fb4f685d15)
---
 integration/loki_micro_services_test.go       |  7 ++++-
 .../frontend/transport/handler.go             | 17 ++---------
 .../frontend/transport/handler_test.go        | 30 -------------------
 3 files changed, 8 insertions(+), 46 deletions(-)
 delete mode 100644 pkg/lokifrontend/frontend/transport/handler_test.go

diff --git a/integration/loki_micro_services_test.go b/integration/loki_micro_services_test.go
index 5111223eaf71f..a4d03ed10a673 100644
--- a/integration/loki_micro_services_test.go
+++ b/integration/loki_micro_services_test.go
@@ -146,11 +146,16 @@ func TestMicroServicesIngestQuery(t *testing.T) {
 		assert.ElementsMatch(t, []map[string]string{{"job": "fake"}}, resp)
 	})
 
-	t.Run("stats error", func(t *testing.T) {
+	t.Run("series error", func(t *testing.T) {
 		_, err := cliQueryFrontend.Series(context.Background(), `{job="fake"}|= "search"`)
 		require.ErrorContains(t, err, "status code 400: only label matchers are supported")
 	})
 
+	t.Run("stats error", func(t *testing.T) {
+		_, err := cliQueryFrontend.Stats(context.Background(), `{job="fake"}|= "search"`)
+		require.ErrorContains(t, err, "status code 400: only label matchers are supported")
+	})
+
 	t.Run("per-request-limits", func(t *testing.T) {
 		queryLimitsPolicy := client.InjectHeadersOption(map[string][]string{querylimits.HTTPHeaderQueryLimitsKey: {`{"maxQueryLength": "1m"}`}})
 		cliQueryFrontendLimited := client.New(tenantID, "", tQueryFrontend.HTTPURL(), queryLimitsPolicy)
diff --git a/pkg/lokifrontend/frontend/transport/handler.go b/pkg/lokifrontend/frontend/transport/handler.go
index 2a4781b3bc718..03332ee046771 100644
--- a/pkg/lokifrontend/frontend/transport/handler.go
+++ b/pkg/lokifrontend/frontend/transport/handler.go
@@ -25,6 +25,7 @@ import (
 	querier_stats "github.com/grafana/loki/pkg/querier/stats"
 	"github.com/grafana/loki/pkg/util"
 	util_log "github.com/grafana/loki/pkg/util/log"
+	"github.com/grafana/loki/pkg/util/server"
 )
 
 const (
@@ -133,7 +134,7 @@ func (f *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	queryResponseTime := time.Since(startTime)
 
 	if err != nil {
-		writeError(w, err)
+		server.WriteError(err, w)
 		return
 	}
 
@@ -229,20 +230,6 @@ func formatQueryString(queryString url.Values) (fields []interface{}) {
 	return fields
 }
 
-func writeError(w http.ResponseWriter, err error) {
-	switch err {
-	case context.Canceled:
-		err = errCanceled
-	case context.DeadlineExceeded:
-		err = errDeadlineExceeded
-	default:
-		if util.IsRequestBodyTooLarge(err) {
-			err = errRequestEntityTooLarge
-		}
-	}
-	httpgrpc.WriteError(w, err)
-}
-
 func writeServiceTimingHeader(queryResponseTime time.Duration, headers http.Header, stats *querier_stats.Stats) {
 	if stats != nil {
 		parts := make([]string, 0)
diff --git a/pkg/lokifrontend/frontend/transport/handler_test.go b/pkg/lokifrontend/frontend/transport/handler_test.go
deleted file mode 100644
index 6f25963626712..0000000000000
--- a/pkg/lokifrontend/frontend/transport/handler_test.go
+++ /dev/null
@@ -1,30 +0,0 @@
-package transport
-
-import (
-	"context"
-	"net/http"
-	"net/http/httptest"
-	"testing"
-
-	"github.com/grafana/dskit/httpgrpc"
-	"github.com/pkg/errors"
-	"github.com/stretchr/testify/require"
-)
-
-func TestWriteError(t *testing.T) {
-	for _, test := range []struct {
-		status int
-		err    error
-	}{
-		{http.StatusInternalServerError, errors.New("unknown")},
-		{http.StatusGatewayTimeout, context.DeadlineExceeded},
-		{StatusClientClosedRequest, context.Canceled},
-		{http.StatusBadRequest, httpgrpc.Errorf(http.StatusBadRequest, "")},
-	} {
-		t.Run(test.err.Error(), func(t *testing.T) {
-			w := httptest.NewRecorder()
-			writeError(w, test.err)
-			require.Equal(t, test.status, w.Result().StatusCode)
-		})
-	}
-}