From bb187ce4b1d8ce49098282f9c23cc794a0f6196c Mon Sep 17 00:00:00 2001
From: Stephanie Hingtgen <stephanie.hingtgen@grafana.com>
Date: Wed, 10 Jul 2024 12:48:02 -0600
Subject: [PATCH] K8s: Match status codes from single tenant (#90153)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Co-authored-by: Jean-Philippe Quémémer <jeanphilippe.quemener@grafana.com>
---
 conf/defaults.ini                        |  2 +-
 pkg/apis/query/v0alpha1/query.go         | 10 +++---
 pkg/apis/query/v0alpha1/query_test.go    | 39 ++++++++++++++++++++++++
 pkg/registry/apis/query/client/plugin.go | 14 ++-------
 4 files changed, 48 insertions(+), 17 deletions(-)

diff --git a/conf/defaults.ini b/conf/defaults.ini
index 3f7693ad008..d365f5f98e5 100644
--- a/conf/defaults.ini
+++ b/conf/defaults.ini
@@ -1945,4 +1945,4 @@ delete_access_policy_timeout = 5s
 # The domain name used to access cms
 domain = grafana-dev.net
 # Folder used to store snapshot files. Defaults to the home dir
-snapshot_folder = ""
\ No newline at end of file
+snapshot_folder = ""
diff --git a/pkg/apis/query/v0alpha1/query.go b/pkg/apis/query/v0alpha1/query.go
index dd7ce99c390..40694086d4b 100644
--- a/pkg/apis/query/v0alpha1/query.go
+++ b/pkg/apis/query/v0alpha1/query.go
@@ -27,14 +27,14 @@ type QueryDataResponse struct {
 	backend.QueryDataResponse `json:",inline"`
 }
 
-// If errors exist, return multi-status
+// GetResponseCode return the right status code for the response by checking the responses.
 func GetResponseCode(rsp *backend.QueryDataResponse) int {
 	if rsp == nil {
-		return http.StatusInternalServerError
+		return http.StatusBadRequest
 	}
-	for _, v := range rsp.Responses {
-		if v.Error != nil {
-			return http.StatusMultiStatus
+	for _, res := range rsp.Responses {
+		if res.Error != nil {
+			return http.StatusBadRequest
 		}
 	}
 	return http.StatusOK
diff --git a/pkg/apis/query/v0alpha1/query_test.go b/pkg/apis/query/v0alpha1/query_test.go
index 03656d5b419..a2aadb8c1f1 100644
--- a/pkg/apis/query/v0alpha1/query_test.go
+++ b/pkg/apis/query/v0alpha1/query_test.go
@@ -2,9 +2,12 @@ package v0alpha1_test
 
 import (
 	"encoding/json"
+	"fmt"
 	"testing"
 
+	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	data "github.com/grafana/grafana-plugin-sdk-go/experimental/apis/data/v0alpha1"
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
 	query "github.com/grafana/grafana/pkg/apis/query/v0alpha1"
@@ -76,3 +79,39 @@ func TestParseQueriesIntoQueryDataRequest(t *testing.T) {
 		}
 	  }`, string(out))
 }
+
+func TestGetResponseCode(t *testing.T) {
+	t.Run("return 200 if no errors in responses", func(t *testing.T) {
+		assert.Equal(t, 200, query.GetResponseCode(&backend.QueryDataResponse{
+			Responses: map[string]backend.DataResponse{
+				"A": {
+					Error: nil,
+				},
+				"B": {
+					Error: nil,
+				},
+			},
+		}))
+	})
+	t.Run("return 400 if there is an error in the responses", func(t *testing.T) {
+		assert.Equal(t, 400, query.GetResponseCode(&backend.QueryDataResponse{
+			Responses: map[string]backend.DataResponse{
+				"A": {
+					Error: fmt.Errorf("some wild error"),
+				},
+			},
+		}))
+	})
+	t.Run("return 400 if there is a partial error", func(t *testing.T) {
+		assert.Equal(t, 400, query.GetResponseCode(&backend.QueryDataResponse{
+			Responses: map[string]backend.DataResponse{
+				"A": {
+					Error: nil,
+				},
+				"B": {
+					Error: fmt.Errorf("some partial error"),
+				},
+			},
+		}))
+	})
+}
diff --git a/pkg/registry/apis/query/client/plugin.go b/pkg/registry/apis/query/client/plugin.go
index bf97b9b486f..dd0e3d2960f 100644
--- a/pkg/registry/apis/query/client/plugin.go
+++ b/pkg/registry/apis/query/client/plugin.go
@@ -80,19 +80,11 @@ func (d *pluginClient) QueryData(ctx context.Context, req data.QueryDataRequest)
 		return http.StatusBadRequest, nil, err
 	}
 
-	code := http.StatusOK
 	rsp, err := d.pluginClient.QueryData(ctx, qdr)
-	if err == nil {
-		for _, v := range rsp.Responses {
-			if v.Error != nil {
-				code = http.StatusMultiStatus
-				break
-			}
-		}
-	} else {
-		code = http.StatusInternalServerError
+	if err != nil {
+		return http.StatusInternalServerError, rsp, err
 	}
-	return code, rsp, err
+	return query.GetResponseCode(rsp), rsp, err
 }
 
 // GetDatasourceAPI implements DataSourceRegistry.