From 3f81f0d145c71f9afbc989a81cd7969df5c04a30 Mon Sep 17 00:00:00 2001
From: Brendan O'Handley <brendan.ohandley@grafana.com>
Date: Fri, 8 Nov 2024 15:19:14 -0600
Subject: [PATCH] Prometheus: Check both error and http status for source of
 error  (#96060)

* check error for source and check http status for source

* refactor

* add source to health check status code as well

* remove error where it will be nil
---
 pkg/promlib/querydata/request.go | 47 +++++++++++++++++++++++++-------
 1 file changed, 37 insertions(+), 10 deletions(-)

diff --git a/pkg/promlib/querydata/request.go b/pkg/promlib/querydata/request.go
index 625113a69c9..6b5c4b98266 100644
--- a/pkg/promlib/querydata/request.go
+++ b/pkg/promlib/querydata/request.go
@@ -234,9 +234,13 @@ func (s *QueryData) fetch(traceCtx context.Context, client *client.Client, q *mo
 func (s *QueryData) rangeQuery(ctx context.Context, c *client.Client, q *models.Query, enablePrometheusDataplaneFlag bool) backend.DataResponse {
 	res, err := c.QueryRange(ctx, q)
 	if err != nil {
+		return addErrorSourceToDataResponse(err)
+	}
+
+	if res.StatusCode/100 != 2 {
+		// for differentiating the source of http errors by status code
 		return backend.DataResponse{
-			Error:  err,
-			Status: backend.StatusBadGateway,
+			ErrorSource: backend.ErrorSourceFromHTTPStatus(res.StatusCode),
 		}
 	}
 
@@ -253,16 +257,22 @@ func (s *QueryData) rangeQuery(ctx context.Context, c *client.Client, q *models.
 func (s *QueryData) instantQuery(ctx context.Context, c *client.Client, q *models.Query, enablePrometheusDataplaneFlag bool) backend.DataResponse {
 	res, err := c.QueryInstant(ctx, q)
 	if err != nil {
-		return backend.DataResponse{
-			Error:  err,
-			Status: backend.StatusBadGateway,
-		}
+		// confirm that it is a downstream error
+		return addErrorSourceToDataResponse(err)
 	}
 
-	// This is only for health check fall back scenario
-	if res.StatusCode != 200 && q.RefId == "__healthcheck__" {
+	if res.StatusCode/100 != 2 {
+		// This is only for health check fall back scenario
+		// add more details to the health check error
+		if q.RefId == "__healthcheck__" {
+			return backend.DataResponse{
+				Error:       errors.New(res.Status),
+				ErrorSource: backend.ErrorSourceFromHTTPStatus(res.StatusCode),
+			}
+		}
+		// for differentiating the source of http errors by status code
 		return backend.DataResponse{
-			Error: errors.New(res.Status),
+			ErrorSource: backend.ErrorSourceFromHTTPStatus(res.StatusCode),
 		}
 	}
 
@@ -279,8 +289,13 @@ func (s *QueryData) instantQuery(ctx context.Context, c *client.Client, q *model
 func (s *QueryData) exemplarQuery(ctx context.Context, c *client.Client, q *models.Query, enablePrometheusDataplaneFlag bool) backend.DataResponse {
 	res, err := c.QueryExemplars(ctx, q)
 	if err != nil {
+		return addErrorSourceToDataResponse(err)
+	}
+
+	if res.StatusCode/100 != 2 {
+		// for differentiating the source of http errors by status code
 		return backend.DataResponse{
-			Error: err,
+			ErrorSource: backend.ErrorSourceFromHTTPStatus(res.StatusCode),
 		}
 	}
 
@@ -304,3 +319,15 @@ func addDataResponse(res *backend.DataResponse, dr *backend.DataResponse) {
 	}
 	dr.Frames = append(dr.Frames, res.Frames...)
 }
+
+func addErrorSourceToDataResponse(err error) backend.DataResponse {
+	response := backend.DataResponse{
+		Error:  err,
+		Status: backend.StatusBadGateway,
+	}
+
+	if backend.IsDownstreamHTTPError(err) {
+		response.ErrorSource = backend.ErrorSourceDownstream
+	}
+	return response
+}