use helper functions to extract team headers instead of core functions

pkg/tsdb/tempo/utils/stream_utils.go

@@ -2,11 +2,10 @@ package stream_utils
 
 import (
 	"context"
+	"encoding/json"
 	"fmt"
 
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
-	"github.com/grafana/grafana/pkg/components/simplejson"
-	"github.com/grafana/grafana/pkg/services/datasources"
 	"google.golang.org/grpc/metadata"
 )
 
@@ -51,28 +50,61 @@ func getTeamHTTPHeaders(plugin backend.PluginContext) (map[string]string, error)
 	headers := map[string]string{}
 	// Grab the JSON data from the datasource instance settings
 	jsonData := plugin.DataSourceInstanceSettings.JSONData
-	js, err := simplejson.NewJson(jsonData)
+	var data map[string]interface{}
+	err := json.Unmarshal(jsonData, &data)
 	if err != nil {
 		return nil, err
 	}
 
 	// fetch team http headers
-	teamHttpHeaders, err := datasources.GetTeamHTTPHeaders(js)
-	if err != nil {
-		return nil, err
-	}
+	if teamHttpHeaders, ok := data["teamHttpHeaders"]; ok {
+		// team headers have the following structure
+		// headers: [<team_id>: [{header: <header_name>, value: <header_value>}]]
+		// header_value is whatever the user has set under LBAC permissions for their given rule.
+		if lbacHeaders, ok := teamHttpHeaders.(map[string]interface{})["headers"]; ok {
+			headerMap := lbacHeaders.(map[string]interface{})
+			labelPolicyKey, labelPolicyValue := getLabelPolicyKeyValue(headerMap)
 
-	// if present, set the Team HTTP Headers
-	if teamHttpHeaders != nil {
-		for _, ruleValue := range teamHttpHeaders.Headers {
-			for _, accessRule := range ruleValue {
-				headers[accessRule.Header] = accessRule.LBACRule
+			if labelPolicyKey != "" && labelPolicyValue != "" {
+				headers[labelPolicyKey] = labelPolicyValue
 			}
 		}
 	}
 	return headers, nil
 }
 
+func getLabelPolicyKeyValue(headerWithRules map[string]interface{}) (string, string) {
+	labelPolicyKey := ""
+	labelPolicyValue := ""
+	// we go through each teams' rule and ignoring the team, go through their set rules and prepare them to be all appended for the X-Prom-Label-Policy header value
+	// the result will be a comma separated list of the rules:
+	// "<rule_num>:<rule_value>, <rule_num>:<rule_value>"
+	for _, accessRuleValue := range headerWithRules {
+		rules := accessRuleValue.([]interface{})
+		for _, accessRule := range rules {
+			header := accessRule.(map[string]interface{})
+			for key, value := range header {
+				// for now, team headers only contain a single header key value, but in case in the future more are introduced, we make sure we only set the one we care about.
+				if key == "header" && value == "X-Prom-Label-Policy" {
+					labelPolicyKey = value.(string)
+					continue
+				}
+				if key == "value" {
+					if valueStr, ok := value.(string); ok {
+						if labelPolicyValue == "" {
+							labelPolicyValue = valueStr
+						} else {
+							labelPolicyValue += "," + valueStr
+						}
+					}
+				}
+			}
+		}
+	}
+
+	return labelPolicyKey, labelPolicyValue
+}
+
 func getClientOptionsHeaders(ctx context.Context, plugin backend.PluginContext) (map[string]string, error) {
 	headers := map[string]string{}
 	opts, err := plugin.DataSourceInstanceSettings.HTTPClientOptions(ctx)

pkg/tsdb/tempo/utils/stream_utils_test.go

@@ -37,8 +37,8 @@ func TestSetHeadersFromIncomingContext_MergesTeamAndClientHeaders(t *testing.T)
 		"teamHttpHeaders": {
 			"headers": {
 				"101": [
-					{"header": "X-Team", "value": "team-value"},
-					{"header": "X-Shared", "value": "team-wins"}
+					{"header": "X-Prom-Label-Policy", "value": "1:team-value"},
+					{"header": "X-Prom-Label-Policy", "value": "2:team-wins"}
 				]
 			}
 		},
@@ -61,9 +61,9 @@ func TestSetHeadersFromIncomingContext_MergesTeamAndClientHeaders(t *testing.T)
 	require.NoError(t, err)
 
 	expected := map[string]string{
-		"X-Client": "client-value",
-		"X-Team":   "team-value",
-		"X-Shared": "team-wins",
+		"X-Client":            "client-value",
+		"X-Prom-Label-Policy": "1:team-value,2:team-wins",
+		"X-Shared":            "client-overridden",
 	}
 	assert.Equal(t, expected, headers)
 }
@@ -80,6 +80,48 @@ func TestGetTeamHTTPHeaders_NoTeamHeaders(t *testing.T) {
 	assert.Empty(t, headers)
 }
 
+func TestGetTeamHTTPHeaders_LabelPolicyValue(t *testing.T) {
+	pluginCtx := backend.PluginContext{
+		DataSourceInstanceSettings: &backend.DataSourceInstanceSettings{
+			JSONData: []byte(`{
+				"teamHttpHeaders": {
+					"headers": {
+						"101": [
+							{"header": "X-Prom-Label-Policy", "value": "1:team-value"},
+							{"header": "X-Prom-Label-Policy", "value": "2:team-wins"}
+						]
+					}
+				}
+			}`),
+		},
+	}
+
+	headers, err := getTeamHTTPHeaders(pluginCtx)
+	require.NoError(t, err)
+	assert.Equal(t, map[string]string{
+		"X-Prom-Label-Policy": "1:team-value,2:team-wins",
+	}, headers)
+}
+
+func TestGetLabelPolicyKeyValue_AppendsValues(t *testing.T) {
+	headerWithRules := map[string]interface{}{
+		"101": []interface{}{
+			map[string]interface{}{
+				"header": "X-Prom-Label-Policy",
+				"value":  "1:alpha",
+			},
+			map[string]interface{}{
+				"header": "X-Prom-Label-Policy",
+				"value":  "2:beta",
+			},
+		},
+	}
+
+	key, value := getLabelPolicyKeyValue(headerWithRules)
+	assert.Equal(t, "X-Prom-Label-Policy", key)
+	assert.Equal(t, "1:alpha,2:beta", value)
+}
+
 func TestGetClientOptionsHeaders_ParsesHeaders(t *testing.T) {
 	pluginCtx := backend.PluginContext{
 		DataSourceInstanceSettings: &backend.DataSourceInstanceSettings{