From 21fe78302b8fbeaed588bc820dee851a222b7c86 Mon Sep 17 00:00:00 2001
From: Ivana Huckova <30407135+ivanahuckova@users.noreply.github.com>
Date: Fri, 3 Dec 2021 16:28:05 +0100
Subject: [PATCH] Prometheus: Fix passing of custom params when running queries
 (#42728)

* Prometheus: Fix passing of custom params

* Remove prints
---
 .../prometheus/custom_query_params_middleware.go    | 13 ++++++++++++-
 .../custom_query_params_middleware_test.go          |  8 ++++++--
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/pkg/tsdb/prometheus/custom_query_params_middleware.go b/pkg/tsdb/prometheus/custom_query_params_middleware.go
index b5c75c1cb83..3bdab4cfe9b 100644
--- a/pkg/tsdb/prometheus/custom_query_params_middleware.go
+++ b/pkg/tsdb/prometheus/custom_query_params_middleware.go
@@ -11,14 +11,25 @@ import (
 const (
 	customQueryParametersMiddlewareName = "prom-custom-query-parameters"
 	customQueryParametersKey            = "customQueryParameters"
+	grafanaDataKey                      = "grafanaData"
 )
 
 func customQueryParametersMiddleware(logger log.Logger) sdkhttpclient.Middleware {
 	return sdkhttpclient.NamedMiddlewareFunc(customQueryParametersMiddlewareName, func(opts sdkhttpclient.Options, next http.RoundTripper) http.RoundTripper {
-		customQueryParamsVal, exists := opts.CustomOptions[customQueryParametersKey]
+		grafanaData, exists := opts.CustomOptions[grafanaDataKey]
 		if !exists {
 			return next
 		}
+
+		data, ok := grafanaData.(map[string]interface{})
+		if !ok {
+			return next
+		}
+		customQueryParamsVal, exists := data[customQueryParametersKey]
+		if !exists {
+			return next
+		}
+
 		customQueryParams, ok := customQueryParamsVal.(string)
 		if !ok || customQueryParams == "" {
 			return next
diff --git a/pkg/tsdb/prometheus/custom_query_params_middleware_test.go b/pkg/tsdb/prometheus/custom_query_params_middleware_test.go
index 643b8c3753c..c4e41ee49e2 100644
--- a/pkg/tsdb/prometheus/custom_query_params_middleware_test.go
+++ b/pkg/tsdb/prometheus/custom_query_params_middleware_test.go
@@ -114,7 +114,9 @@ func TestCustomQueryParametersMiddleware(t *testing.T) {
 		mw := customQueryParametersMiddleware(log.New("test"))
 		rt := mw.CreateMiddleware(httpclient.Options{
 			CustomOptions: map[string]interface{}{
-				customQueryParametersKey: "custom=par/am&second=f oo",
+				grafanaDataKey: map[string]interface{}{
+					customQueryParametersKey: "custom=par/am&second=f oo",
+				},
 			},
 		}, finalRoundTripper)
 		require.NotNil(t, rt)
@@ -144,7 +146,9 @@ func TestCustomQueryParametersMiddleware(t *testing.T) {
 		mw := customQueryParametersMiddleware(log.New("test"))
 		rt := mw.CreateMiddleware(httpclient.Options{
 			CustomOptions: map[string]interface{}{
-				customQueryParametersKey: "custom=par/am&second=f oo",
+				grafanaDataKey: map[string]interface{}{
+					customQueryParametersKey: "custom=par/am&second=f oo",
+				},
 			},
 		}, finalRoundTripper)
 		require.NotNil(t, rt)