AzureMonitor: Add custom header support to Azure Monitor (#60269)

* Add integration test for Azure Monitor

- Add Azure Monitor to datasource types
- Update instance creation to correctly set HTTP client options
- Combine custom azure headers and custom grafana headers on HTTP client creation
- Update HTTP client tests

* Test custom azure headers

pkg/services/datasources/models.go

@@ -27,6 +27,7 @@ const (
 	DS_ACCESS_PROXY   = "proxy"
 	DS_ES_OPEN_DISTRO = "grafana-es-open-distro-datasource"
 	DS_ES_OPENSEARCH  = "grafana-opensearch-datasource"
+	DS_AZURE_MONITOR  = "grafana-azure-monitor-datasource"
 )
 
 type DsAccess string

pkg/tests/api/azuremonitor/azuremonitor_test.go

@@ -0,0 +1,121 @@
+package azuremonitor
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/grafana/grafana/pkg/api/dtos"
+	"github.com/grafana/grafana/pkg/components/simplejson"
+	"github.com/grafana/grafana/pkg/services/datasources"
+	"github.com/grafana/grafana/pkg/services/org"
+	"github.com/grafana/grafana/pkg/services/user"
+	"github.com/grafana/grafana/pkg/tests/testinfra"
+	"github.com/stretchr/testify/require"
+)
+
+func TestIntegrationAzureMonitor(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping integration test")
+	}
+	dir, path := testinfra.CreateGrafDir(t, testinfra.GrafanaOpts{
+		DisableAnonymous: true,
+	})
+
+	grafanaListeningAddr, testEnv := testinfra.StartGrafanaEnv(t, dir, path)
+	ctx := context.Background()
+
+	testinfra.CreateUser(t, testEnv.SQLStore, user.CreateUserCommand{
+		DefaultOrgRole: string(org.RoleAdmin),
+		Password:       "admin",
+		Login:          "admin",
+	})
+
+	var outgoingRequest *http.Request
+	outgoingServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		outgoingRequest = r
+		w.WriteHeader(http.StatusUnauthorized)
+	}))
+	t.Cleanup(outgoingServer.Close)
+
+	jsonData := simplejson.NewFromAny(map[string]interface{}{
+		"httpHeaderName1": "X-CUSTOM-HEADER",
+		"clientId":        "test-client-id",
+		"tenantId":        "test-tenant-id",
+		"cloudName":       "customizedazuremonitor",
+		"customizedRoutes": map[string]interface{}{
+			"Azure Monitor": map[string]interface{}{
+				"URL": outgoingServer.URL,
+				"Headers": map[string]string{
+					"custom-azure-header": "custom-azure-value",
+				},
+			},
+		},
+	})
+	secureJSONData := map[string]string{
+		"clientSecret":     "test-client-secret",
+		"httpHeaderValue1": "custom-header-value",
+	}
+
+	uid := "azuremonitor"
+	err := testEnv.Server.HTTPServer.DataSourcesService.AddDataSource(ctx, &datasources.AddDataSourceCommand{
+		OrgId:          1,
+		Access:         datasources.DS_ACCESS_PROXY,
+		Name:           "Azure Monitor",
+		Type:           datasources.DS_AZURE_MONITOR,
+		Uid:            uid,
+		Url:            outgoingServer.URL,
+		JsonData:       jsonData,
+		SecureJsonData: secureJSONData,
+	})
+	require.NoError(t, err)
+
+	t.Run("When calling /api/ds/query should set expected headers on outgoing HTTP request", func(t *testing.T) {
+		query := simplejson.NewFromAny(map[string]interface{}{
+			"datasource": map[string]interface{}{
+				"type": "grafana-azure-monitor-datasource",
+				"uid":  uid,
+			},
+			"queryType": "Azure Monitor",
+			"azureMonitor": map[string]interface{}{
+				"resourceGroup":   "test-rg",
+				"metricNamespace": "microsoft.storage/storageaccounts",
+				"resourceName":    "testacct",
+				"timeGrain":       "auto",
+				"metricName":      "UsedCapacity",
+				"aggregation":     "Average",
+			},
+			"subscription": "test-sub",
+			"intervalMs":   30000,
+		})
+		buf1 := &bytes.Buffer{}
+		err = json.NewEncoder(buf1).Encode(dtos.MetricRequest{
+			From:    "1668078080000",
+			To:      "1668081680000",
+			Queries: []*simplejson.Json{query},
+		})
+		require.NoError(t, err)
+		u := fmt.Sprintf("http://admin:admin@%s/api/ds/query", grafanaListeningAddr)
+		// nolint:gosec
+		resp, err := http.Post(u, "application/json", buf1)
+		require.NoError(t, err)
+		require.Equal(t, http.StatusBadRequest, resp.StatusCode)
+		t.Cleanup(func() {
+			err := resp.Body.Close()
+			require.NoError(t, err)
+		})
+		_, err = io.ReadAll(resp.Body)
+		require.NoError(t, err)
+
+		require.NotNil(t, outgoingRequest)
+		require.Equal(t, "/subscriptions/test-sub/resourceGroups/test-rg/providers/microsoft.storage/storageaccounts/testacct/providers/microsoft.insights/metrics?aggregation=Average&api-version=2021-05-01&interval=PT1M&metricnames=UsedCapacity&metricnamespace=microsoft.storage%2Fstorageaccounts&timespan=2022-11-10T11%3A01%3A20Z%2F2022-11-10T12%3A01%3A20Z",
+			outgoingRequest.URL.String())
+		require.Equal(t, "custom-header-value", outgoingRequest.Header.Get("X-CUSTOM-HEADER"))
+		require.Equal(t, "custom-azure-value", outgoingRequest.Header.Get("custom-azure-header"))
+	})
+}

pkg/tsdb/azuremonitor/azuremonitor.go

@@ -68,9 +68,9 @@ type Service struct {
 	tracer          tracing.Tracer
 }
 
-func getDatasourceService(cfg *setting.Cfg, clientProvider *httpclient.Provider, dsInfo types.DatasourceInfo, routeName string) (types.DatasourceService, error) {
+func getDatasourceService(cfg *setting.Cfg, clientProvider *httpclient.Provider, dsInfo types.DatasourceInfo, routeName string, httpClientOptions httpclient.Options) (types.DatasourceService, error) {
 	route := dsInfo.Routes[routeName]
-	client, err := newHTTPClient(route, dsInfo, cfg, clientProvider)
+	client, err := newHTTPClient(route, dsInfo, cfg, clientProvider, httpClientOptions)
 	if err != nil {
 		return types.DatasourceService{}, err
 	}
@@ -86,13 +86,17 @@ func NewInstanceSettings(cfg *setting.Cfg, clientProvider *httpclient.Provider,
 		if err != nil {
 			return nil, fmt.Errorf("error reading settings: %w", err)
 		}
-
 		jsonDataObj := map[string]interface{}{}
 		err = json.Unmarshal(settings.JSONData, &jsonDataObj)
 		if err != nil {
 			return nil, fmt.Errorf("error reading settings: %w", err)
 		}
 
+		httpClientOpts, err := settings.HTTPClientOptions()
+		if err != nil {
+			return nil, fmt.Errorf("error getting http options: %w", err)
+		}
+
 		azMonitorSettings := types.AzureMonitorSettings{}
 		err = json.Unmarshal(settings.JSONData, &azMonitorSettings)
 		if err != nil {
@@ -126,7 +130,7 @@ func NewInstanceSettings(cfg *setting.Cfg, clientProvider *httpclient.Provider,
 		}
 
 		for routeName := range executors {
-			service, err := getDatasourceService(cfg, clientProvider, model, routeName)
+			service, err := getDatasourceService(cfg, clientProvider, model, routeName, httpClientOpts)
 			if err != nil {
 				return nil, err
 			}

pkg/tsdb/azuremonitor/httpclient.go

@@ -13,9 +13,9 @@ import (
 	"github.com/grafana/grafana/pkg/tsdb/azuremonitor/types"
 )
 
-func newHTTPClient(route types.AzRoute, model types.DatasourceInfo, cfg *setting.Cfg, clientProvider httpclient.Provider) (*http.Client, error) {
-	opts := sdkhttpclient.Options{
-		Headers: route.Headers,
+func newHTTPClient(route types.AzRoute, model types.DatasourceInfo, cfg *setting.Cfg, clientProvider httpclient.Provider, httpClientOptions sdkhttpclient.Options) (*http.Client, error) {
+	for header, value := range route.Headers {
+		httpClientOptions.Headers[header] = value
 	}
 
 	// Use Azure credentials if the route has OAuth scopes configured
@@ -23,8 +23,8 @@ func newHTTPClient(route types.AzRoute, model types.DatasourceInfo, cfg *setting
 		if cred, ok := model.Credentials.(*azcredentials.AzureClientSecretCredentials); ok && cred.ClientSecret == "" {
 			return nil, fmt.Errorf("unable to initialize HTTP Client: clientSecret not found")
 		}
-		azhttpclient.AddAzureAuthentication(&opts, cfg.Azure, model.Credentials, route.Scopes)
+		azhttpclient.AddAzureAuthentication(&httpClientOptions, cfg.Azure, model.Credentials, route.Scopes)
 	}
 
-	return clientProvider.New(opts)
+	return clientProvider.New(httpClientOptions)
 }

pkg/tsdb/azuremonitor/httpclient_test.go

@@ -28,7 +28,7 @@ func TestHttpClient_AzureCredentials(t *testing.T) {
 			Scopes: []string{"https://management.azure.com/.default"},
 		}
 
-		_, err := newHTTPClient(route, model, cfg, provider)
+		_, err := newHTTPClient(route, model, cfg, provider, sdkhttpclient.Options{})
 		require.NoError(t, err)
 
 		require.NotNil(t, provider.opts)
@@ -41,7 +41,7 @@ func TestHttpClient_AzureCredentials(t *testing.T) {
 			Scopes: []string{},
 		}
 
-		_, err := newHTTPClient(route, model, cfg, provider)
+		_, err := newHTTPClient(route, model, cfg, provider, sdkhttpclient.Options{})
 		require.NoError(t, err)
 
 		assert.NotNil(t, provider.opts)
@@ -50,6 +50,33 @@ func TestHttpClient_AzureCredentials(t *testing.T) {
 			assert.Len(t, provider.opts.Middlewares, 0)
 		}
 	})
+
+	t.Run("should combine custom azure and custom grafana headers", func(t *testing.T) {
+		route := types.AzRoute{
+			Headers: map[string]string{
+				"AzureHeader": "AzureValue",
+			},
+		}
+		opts := sdkhttpclient.Options{
+			Headers: map[string]string{
+				"GrafanaHeader": "GrafanaValue",
+			},
+		}
+
+		res := map[string]string{
+			"GrafanaHeader": "GrafanaValue",
+			"AzureHeader":   "AzureValue",
+		}
+		_, err := newHTTPClient(route, model, cfg, provider, opts)
+		require.NoError(t, err)
+
+		assert.NotNil(t, provider.opts)
+
+		if provider.opts.Headers != nil {
+			assert.Len(t, provider.opts.Headers, 2)
+			assert.Equal(t, res, provider.opts.Headers)
+		}
+	})
 }
 
 type fakeHttpClientProvider struct {