Share azureauth between prometheus clients (#58122)

* Move azureauth to upper package

* Refactor http transport options

pkg/tsdb/prometheus/buffered/azureauth/azure.go

@@ -1,113 +0,0 @@
-package azureauth
-
-import (
-	"fmt"
-	"net/url"
-	"path"
-
-	"github.com/grafana/grafana-azure-sdk-go/azcredentials"
-	"github.com/grafana/grafana-azure-sdk-go/azhttpclient"
-	"github.com/grafana/grafana-azure-sdk-go/azsettings"
-	"github.com/grafana/grafana-plugin-sdk-go/backend"
-	sdkhttpclient "github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
-	"github.com/grafana/grafana/pkg/tsdb/prometheus/utils"
-
-	"github.com/grafana/grafana/pkg/util/maputil"
-)
-
-var (
-	azurePrometheusScopes = map[string][]string{
-		azsettings.AzurePublic:       {"https://prometheus.monitor.azure.com/.default"},
-		azsettings.AzureChina:        {"https://prometheus.monitor.chinacloudapp.cn/.default"},
-		azsettings.AzureUSGovernment: {"https://prometheus.monitor.usgovcloudapi.net/.default"},
-	}
-)
-
-func ConfigureAzureAuthentication(settings backend.DataSourceInstanceSettings, azureSettings *azsettings.AzureSettings, opts *sdkhttpclient.Options) error {
-	jsonData, err := utils.GetJsonData(settings)
-	if err != nil {
-		return fmt.Errorf("failed to get jsonData: %w", err)
-	}
-	credentials, err := azcredentials.FromDatasourceData(jsonData, settings.DecryptedSecureJSONData)
-	if err != nil {
-		err = fmt.Errorf("invalid Azure credentials: %w", err)
-		return err
-	}
-
-	if credentials != nil {
-		var scopes []string
-
-		if scopes, err = getOverriddenScopes(jsonData); err != nil {
-			return err
-		}
-
-		if scopes == nil {
-			if scopes, err = getPrometheusScopes(azureSettings, credentials); err != nil {
-				return err
-			}
-		}
-
-		azhttpclient.AddAzureAuthentication(opts, azureSettings, credentials, scopes)
-	}
-
-	return nil
-}
-
-func getOverriddenScopes(jsonData map[string]interface{}) ([]string, error) {
-	resourceIdStr, err := maputil.GetStringOptional(jsonData, "azureEndpointResourceId")
-	if err != nil {
-		err = fmt.Errorf("overridden resource ID (audience) invalid")
-		return nil, err
-	} else if resourceIdStr == "" {
-		return nil, nil
-	}
-
-	resourceId, err := url.Parse(resourceIdStr)
-	if err != nil || resourceId.Scheme == "" || resourceId.Host == "" {
-		err = fmt.Errorf("overridden endpoint resource ID (audience) '%s' invalid", resourceIdStr)
-		return nil, err
-	}
-
-	resourceId.Path = path.Join(resourceId.Path, ".default")
-	scopes := []string{resourceId.String()}
-	return scopes, nil
-}
-
-func getPrometheusScopes(settings *azsettings.AzureSettings, credentials azcredentials.AzureCredentials) ([]string, error) {
-	// Extract cloud from credentials
-	azureCloud, err := getAzureCloudFromCredentials(settings, credentials)
-	if err != nil {
-		return nil, err
-	}
-
-	// Get scopes for the given cloud
-	if scopes, ok := azurePrometheusScopes[azureCloud]; !ok {
-		err := fmt.Errorf("the Azure cloud '%s' not supported by Prometheus datasource", azureCloud)
-		return nil, err
-	} else {
-		return scopes, nil
-	}
-}
-
-// To be part of grafana-azure-sdk-go
-func getAzureCloudFromCredentials(settings *azsettings.AzureSettings, credentials azcredentials.AzureCredentials) (string, error) {
-	switch c := credentials.(type) {
-	case *azcredentials.AzureManagedIdentityCredentials:
-		// In case of managed identity, the cloud is always same as where Grafana is hosted
-		return getDefaultAzureCloud(settings), nil
-	case *azcredentials.AzureClientSecretCredentials:
-		return c.AzureCloud, nil
-	default:
-		err := fmt.Errorf("the Azure credentials of type '%s' not supported by Prometheus datasource", c.AzureAuthType())
-		return "", err
-	}
-}
-
-// To be part of grafana-azure-sdk-go
-func getDefaultAzureCloud(settings *azsettings.AzureSettings) string {
-	cloudName := settings.Cloud
-	if cloudName == "" {
-		return azsettings.AzurePublic
-	}
-	return cloudName
-}

pkg/tsdb/prometheus/buffered/azureauth/azure_test.go

@@ -1,117 +0,0 @@
-package azureauth
-
-import (
-	"testing"
-
-	"github.com/grafana/grafana-azure-sdk-go/azsettings"
-	"github.com/grafana/grafana-plugin-sdk-go/backend"
-	sdkhttpclient "github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
-
-	"github.com/grafana/grafana/pkg/setting"
-	"github.com/stretchr/testify/assert"
-	"github.com/stretchr/testify/require"
-)
-
-func TestConfigureAzureAuthentication(t *testing.T) {
-	cfg := &setting.Cfg{
-		Azure: &azsettings.AzureSettings{},
-	}
-
-	t.Run("should set Azure middleware when JsonData contains valid credentials", func(t *testing.T) {
-		settings := backend.DataSourceInstanceSettings{
-			JSONData: []byte(`{
-					"httpMethod": "POST",
-					"azureCredentials": {
-						"authType": "msi"
-					}
-				}`),
-		}
-
-		var opts = &sdkhttpclient.Options{CustomOptions: map[string]interface{}{}}
-
-		err := ConfigureAzureAuthentication(settings, cfg.Azure, opts)
-		require.NoError(t, err)
-
-		require.NotNil(t, opts.Middlewares)
-		assert.Len(t, opts.Middlewares, 1)
-	})
-
-	t.Run("should not set Azure middleware when JsonData doesn't contain valid credentials", func(t *testing.T) {
-		settings := backend.DataSourceInstanceSettings{
-			JSONData: []byte(`{ "httpMethod": "POST" }`),
-		}
-
-		var opts = &sdkhttpclient.Options{CustomOptions: map[string]interface{}{}}
-
-		err := ConfigureAzureAuthentication(settings, cfg.Azure, opts)
-		require.NoError(t, err)
-
-		assert.NotContains(t, opts.CustomOptions, "_azureCredentials")
-	})
-
-	t.Run("should return error when JsonData contains invalid credentials", func(t *testing.T) {
-		settings := backend.DataSourceInstanceSettings{
-			JSONData: []byte(`{
-					"httpMethod":       "POST",
-					"azureCredentials": "invalid"
-				}`),
-		}
-
-		var opts = &sdkhttpclient.Options{CustomOptions: map[string]interface{}{}}
-		err := ConfigureAzureAuthentication(settings, cfg.Azure, opts)
-		assert.Error(t, err)
-	})
-
-	t.Run("should set Azure middleware when JsonData contains credentials and valid audience", func(t *testing.T) {
-		settings := backend.DataSourceInstanceSettings{
-			JSONData: []byte(`{
-					"httpMethod": "POST",
-					"azureCredentials": {
-						"authType": "msi"
-					},
-					"azureEndpointResourceId": "https://api.example.com/abd5c4ce-ca73-41e9-9cb2-bed39aa2adb5"
-				}`),
-		}
-		var opts = &sdkhttpclient.Options{CustomOptions: map[string]interface{}{}}
-
-		err := ConfigureAzureAuthentication(settings, cfg.Azure, opts)
-		require.NoError(t, err)
-
-		require.NotNil(t, opts.Middlewares)
-		assert.Len(t, opts.Middlewares, 1)
-	})
-
-	t.Run("should not set Azure middleware when JsonData doesn't contain credentials", func(t *testing.T) {
-		settings := backend.DataSourceInstanceSettings{
-			JSONData: []byte(`{
-					"httpMethod":              "POST",
-					"azureEndpointResourceId": "https://api.example.com/abd5c4ce-ca73-41e9-9cb2-bed39aa2adb5"
-				}`),
-		}
-		var opts = &sdkhttpclient.Options{CustomOptions: map[string]interface{}{}}
-
-		err := ConfigureAzureAuthentication(settings, cfg.Azure, opts)
-		require.NoError(t, err)
-
-		if opts.Middlewares != nil {
-			assert.Len(t, opts.Middlewares, 0)
-		}
-	})
-
-	t.Run("should return error when JsonData contains invalid audience", func(t *testing.T) {
-		settings := backend.DataSourceInstanceSettings{
-			JSONData: []byte(`{
-					"httpMethod": "POST",
-					"azureCredentials": {
-						"authType": "msi"
-					},
-					"azureEndpointResourceId": "invalid"
-				}`),
-		}
-
-		var opts = &sdkhttpclient.Options{CustomOptions: map[string]interface{}{}}
-
-		err := ConfigureAzureAuthentication(settings, cfg.Azure, opts)
-		assert.Error(t, err)
-	})
-}

pkg/tsdb/prometheus/buffered/client.go

@@ -1,79 +0,0 @@
-package buffered
-
-import (
-	"fmt"
-	"net/http"
-	"strings"
-
-	"github.com/grafana/grafana-plugin-sdk-go/backend"
-	sdkhttpclient "github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
-	"github.com/grafana/grafana/pkg/infra/log"
-	"github.com/grafana/grafana/pkg/setting"
-	"github.com/grafana/grafana/pkg/tsdb/prometheus/buffered/azureauth"
-	"github.com/grafana/grafana/pkg/tsdb/prometheus/middleware"
-	"github.com/grafana/grafana/pkg/tsdb/prometheus/utils"
-	"github.com/grafana/grafana/pkg/util/maputil"
-	"github.com/prometheus/client_golang/api"
-	apiv1 "github.com/prometheus/client_golang/api/prometheus/v1"
-)
-
-// CreateTransportOptions creates options for the http client. Probably should be shared and should not live in the
-// buffered package.
-func CreateTransportOptions(settings backend.DataSourceInstanceSettings, cfg *setting.Cfg, logger log.Logger) (*sdkhttpclient.Options, error) {
-	opts, err := settings.HTTPClientOptions()
-	if err != nil {
-		return nil, err
-	}
-
-	jsonData, err := utils.GetJsonData(settings)
-	if err != nil {
-		return nil, fmt.Errorf("error reading settings: %w", err)
-	}
-	httpMethod, _ := maputil.GetStringOptional(jsonData, "httpMethod")
-
-	opts.Middlewares = middlewares(logger, httpMethod)
-
-	// Set SigV4 service namespace
-	if opts.SigV4 != nil {
-		opts.SigV4.Service = "aps"
-	}
-
-	// Set Azure authentication
-	if cfg.AzureAuthEnabled {
-		err = azureauth.ConfigureAzureAuthentication(settings, cfg.Azure, &opts)
-		if err != nil {
-			return nil, fmt.Errorf("error configuring Azure auth: %v", err)
-		}
-	}
-
-	return &opts, nil
-}
-
-func CreateClient(roundTripper http.RoundTripper, url string) (apiv1.API, error) {
-	cfg := api.Config{
-		Address:      url,
-		RoundTripper: roundTripper,
-	}
-
-	client, err := api.NewClient(cfg)
-	if err != nil {
-		return nil, err
-	}
-
-	return apiv1.NewAPI(client), nil
-}
-
-func middlewares(logger log.Logger, httpMethod string) []sdkhttpclient.Middleware {
-	middlewares := []sdkhttpclient.Middleware{
-		// TODO: probably isn't needed anymore and should by done by http infra code
-		middleware.CustomQueryParameters(logger),
-		sdkhttpclient.CustomHeadersMiddleware(),
-	}
-
-	// Needed to control GET vs POST method of the requests
-	if strings.ToLower(httpMethod) == "get" {
-		middlewares = append(middlewares, middleware.ForceHttpGet(logger))
-	}
-
-	return middlewares
-}

pkg/tsdb/prometheus/buffered/client_test.go

@@ -1,44 +0,0 @@
-package buffered
-
-import (
-	"testing"
-
-	"github.com/grafana/grafana-azure-sdk-go/azsettings"
-	"github.com/grafana/grafana-plugin-sdk-go/backend"
-	"github.com/grafana/grafana/pkg/infra/log/logtest"
-	"github.com/grafana/grafana/pkg/setting"
-	"github.com/stretchr/testify/require"
-)
-
-func TestCreateTransportOptions(t *testing.T) {
-	t.Run("creates correct options object", func(t *testing.T) {
-		settings := backend.DataSourceInstanceSettings{
-			BasicAuthEnabled: false,
-			BasicAuthUser:    "",
-			JSONData:         []byte(`{"httpHeaderName1": "foo"}`),
-			DecryptedSecureJSONData: map[string]string{
-				"httpHeaderValue1": "bar",
-			},
-		}
-		opts, err := CreateTransportOptions(settings, &setting.Cfg{}, &logtest.Fake{})
-		require.NoError(t, err)
-		require.Equal(t, map[string]string{"foo": "bar"}, opts.Headers)
-		require.Equal(t, 2, len(opts.Middlewares))
-	})
-
-	t.Run("add azure credentials if configured", func(t *testing.T) {
-		settings := backend.DataSourceInstanceSettings{
-			BasicAuthEnabled: false,
-			BasicAuthUser:    "",
-			JSONData: []byte(`{
-				"azureCredentials": {
-					"authType": "msi"
-				}
-			}`),
-			DecryptedSecureJSONData: map[string]string{},
-		}
-		opts, err := CreateTransportOptions(settings, &setting.Cfg{AzureAuthEnabled: true, Azure: &azsettings.AzureSettings{}}, &logtest.Fake{})
-		require.NoError(t, err)
-		require.Equal(t, 3, len(opts.Middlewares))
-	})
-}

pkg/tsdb/prometheus/buffered/time_series_query.go

@@ -15,6 +15,7 @@ import (
 	"github.com/grafana/grafana-plugin-sdk-go/backend"
 	sdkHTTPClient "github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
 	"github.com/grafana/grafana-plugin-sdk-go/data"
+	"github.com/grafana/grafana/pkg/tsdb/prometheus/client"
 	apiv1 "github.com/prometheus/client_golang/api/prometheus/v1"
 	"github.com/prometheus/common/model"
 	"go.opentelemetry.io/otel/attribute"
@@ -68,7 +69,7 @@ type Buffered struct {
 // New creates and object capable of executing and parsing a Prometheus queries. It's "buffered" because there is
 // another implementation capable of streaming parse the response.
 func New(roundTripper http.RoundTripper, tracer tracing.Tracer, settings backend.DataSourceInstanceSettings, plog log.Logger) (*Buffered, error) {
-	promClient, err := CreateClient(roundTripper, settings.URL)
+	promClient, err := client.CreateAPIClient(roundTripper, settings.URL)
 	if err != nil {
 		return nil, fmt.Errorf("error creating prom client: %v", err)
 	}
@@ -232,7 +233,7 @@ func (b *Buffered) parseTimeSeriesQuery(req *backend.QueryDataRequest) ([]*Prome
 		if err != nil {
 			return nil, fmt.Errorf("error unmarshaling query model: %v", err)
 		}
-		//Final interval value
+		// Final interval value
 		interval, err := calculatePrometheusInterval(model, b.TimeInterval, query, b.intervalCalculator)
 		if err != nil {
 			return nil, fmt.Errorf("error calculating interval: %v", err)
@@ -301,7 +302,7 @@ func parseTimeSeriesResponse(value map[TimeSeriesQueryType]interface{}, query *P
 func calculatePrometheusInterval(model *QueryModel, timeInterval string, query backend.DataQuery, intervalCalculator intervalv2.Calculator) (time.Duration, error) {
 	queryInterval := model.Interval
 
-	//If we are using variable for interval/step, we will replace it with calculated interval
+	// If we are using variable for interval/step, we will replace it with calculated interval
 	if isVariableInterval(queryInterval) {
 		queryInterval = ""
 	}
@@ -656,7 +657,7 @@ func isVariableInterval(interval string) bool {
 	if interval == varInterval || interval == varIntervalMs || interval == varRateInterval {
 		return true
 	}
-	//Repetitive code, we should have functionality to unify these
+	// Repetitive code, we should have functionality to unify these
 	if interval == varIntervalAlt || interval == varIntervalMsAlt || interval == varRateIntervalAlt {
 		return true
 	}