From c0401cb3106deccf6e42273a5536248fb7fdc568 Mon Sep 17 00:00:00 2001
From: Grzegorz Piwowarek <gpiwowarek@gmail.com>
Date: Tue, 21 May 2024 18:43:57 +0200
Subject: [PATCH] Move header processing to responseFromQuesma() (#172)

---
 quesma/quesma/mux/mux.go |  7 +++++++
 quesma/quesma/quesma.go  | 21 +++++++++++----------
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/quesma/quesma/mux/mux.go b/quesma/quesma/mux/mux.go
index e9a6a77e7..598ddfe4a 100644
--- a/quesma/quesma/mux/mux.go
+++ b/quesma/quesma/mux/mux.go
@@ -46,6 +46,13 @@ type (
 
 type RequestMatcherFunc func(req *Request) bool
 
+func ServerErrorResult() *Result {
+	return &Result{
+		StatusCode: http.StatusInternalServerError,
+		Meta:       map[string]string{"Content-Type": "text/plain"},
+	}
+}
+
 func (f RequestMatcherFunc) Matches(req *Request) bool {
 	return f(req)
 }
diff --git a/quesma/quesma/quesma.go b/quesma/quesma/quesma.go
index d7a2df2f0..631e89aec 100644
--- a/quesma/quesma/quesma.go
+++ b/quesma/quesma/quesma.go
@@ -55,12 +55,18 @@ func responseFromElastic(ctx context.Context, elkResponse *http.Response, w http
 	elkResponse.Body.Close()
 }
 
-func responseFromQuesma(ctx context.Context, unzipped []byte, w http.ResponseWriter, elkResponse *http.Response, statusCode int, zip bool) {
+func responseFromQuesma(ctx context.Context, unzipped []byte, w http.ResponseWriter, elkResponse *http.Response, quesmaResponse *mux.Result, zip bool) {
 	id := ctx.Value(tracing.RequestIdCtxKey).(string)
 	logger.Debug().Str(logger.RID, id).Msg("responding from Quesma")
 
+	for key, value := range quesmaResponse.Meta {
+		w.Header().Set(key, value)
+	}
+	if zip {
+		w.Header().Set("Content-Encoding", "gzip")
+	}
 	w.Header().Set(quesmaSourceHeader, quesmaSourceClickhouse)
-	w.WriteHeader(statusCode)
+	w.WriteHeader(quesmaResponse.StatusCode)
 	if zip {
 		zipped, err := gzip.Zip(unzipped)
 		if err != nil {
@@ -190,13 +196,8 @@ func (r *router) reroute(ctx context.Context, w http.ResponseWriter, req *http.R
 				logger.WarnWithCtx(ctx).Msg("empty response from Clickhouse")
 			}
 			addProductAndContentHeaders(req.Header, w.Header())
-			for key, value := range quesmaResponse.Meta {
-				w.Header().Set(key, value)
-			}
-			if zip {
-				w.Header().Set("Content-Encoding", "gzip")
-			}
-			responseFromQuesma(ctx, unzipped, w, elkResponse, quesmaResponse.StatusCode, zip)
+
+			responseFromQuesma(ctx, unzipped, w, elkResponse, quesmaResponse, zip)
 
 		} else {
 			if elkResponse != nil && r.config.Mode == config.DualWriteQueryClickhouseFallback {
@@ -213,7 +214,7 @@ func (r *router) reroute(ctx context.Context, w http.ResponseWriter, req *http.R
 
 				// We should not send our error message to the client. There can be sensitive information in it.
 				// We will send ID of failed request instead
-				responseFromQuesma(ctx, []byte(fmt.Sprintf("Internal server error. Request ID: %s\n", requestId)), w, elkResponse, 500, zip)
+				responseFromQuesma(ctx, []byte(fmt.Sprintf("Internal server error. Request ID: %s\n", requestId)), w, elkResponse, mux.ServerErrorResult(), zip)
 			}
 		}
 	} else {